zclaw_openfang/desktop/src/lib/saas-relay.ts

/**
 * SaaS Relay Methods — Mixin
 *
 * Installs relay-related methods (tasks, chat completion, usage) onto
 * SaaSClient.prototype. Uses the same mixin pattern as gateway-api.ts.
 */

import type {
  RelayTaskInfo,
} from './saas-types';
import { createLogger } from './logger';
const logger = createLogger('SaaSRelay');

export function installRelayMethods(ClientClass: { prototype: any }): void {
  const proto = ClientClass.prototype;

  // --- Relay Task Management ---

  /** List relay tasks for the current user */
  proto.listRelayTasks = async function (this: { request<T>(method: string, path: string, body?: unknown): Promise<T> }, query?: { status?: string; page?: number; page_size?: number }): Promise<RelayTaskInfo[]> {
    const params = new URLSearchParams();
    if (query?.status) params.set('status', query.status);
    if (query?.page) params.set('page', String(query.page));
    if (query?.page_size) params.set('page_size', String(query.page_size));
    const qs = params.toString();
    return this.request<RelayTaskInfo[]>('GET', `/api/v1/relay/tasks${qs ? '?' + qs : ''}`);
  };

  /** Retry a failed relay task (admin only) */
  proto.retryRelayTask = async function (this: { request<T>(method: string, path: string, body?: unknown): Promise<T> }, taskId: string): Promise<{ ok: boolean; task_id: string }> {
    return this.request<{ ok: boolean; task_id: string }>('POST', `/api/v1/relay/tasks/${taskId}/retry`);
  };

  // --- Chat Relay ---

  /**
   * Send a chat completion request via the SaaS relay.
   * Returns the raw Response object to support both streaming and non-streaming.
   *
   * Includes one retry on 401 (auto token refresh) and on network errors.
   * The caller is responsible for:
   * - Reading the response body (JSON or SSE stream)
   * - Handling errors from the response
   */
  proto.chatCompletion = async function (
    this: {
      baseUrl: string;
      token: string | null;
      _serverReachable: boolean;
      _isAuthEndpoint(path: string): boolean;
      refreshToken(): Promise<string>;
      refreshMutex(): Promise<string>;
    },
    body: unknown,
    signal?: AbortSignal,
  ): Promise<Response> {
    const maxAttempts = 2; // 1 initial + 1 retry

    for (let attempt = 0; attempt < maxAttempts; attempt++) {
      const headers: Record<string, string> = {
        'Content-Type': 'application/json',
      };
      if (this.token) {
        headers['Authorization'] = `Bearer ${this.token}`;
      }

      // Use caller's AbortSignal if provided, otherwise default 5min timeout
      const effectiveSignal = signal ?? AbortSignal.timeout(300_000);

      try {
        const response = await fetch(
          `${this.baseUrl}/api/v1/relay/chat/completions`,
          {
            method: 'POST',
            headers,
            credentials: 'include', // Send HttpOnly cookies
            body: JSON.stringify(body),
            signal: effectiveSignal,
          },
        );

        // On 401, attempt token refresh once
        if (response.status === 401 && attempt === 0 && !this._isAuthEndpoint('/api/v1/relay/chat/completions')) {
          try {
            const newToken = await this.refreshMutex();
            if (newToken) continue; // Retry with refreshed token
          } catch (e) {
            logger.debug('Token refresh failed', { error: e });
            // Refresh failed, return the 401 response
          }
        }

        this._serverReachable = true;
        return response;
      } catch (err: unknown) {
        this._serverReachable = false;
        const isNetworkError = err instanceof TypeError
          && (err.message.includes('Failed to fetch') || err.message.includes('NetworkError'));

        if (isNetworkError && attempt < maxAttempts - 1) {
          // Brief backoff before retry
          await new Promise((r) => setTimeout(r, 1000 * (attempt + 1)));
          continue;
        }

        throw err;
      }
    }

    // Unreachable but TypeScript needs it
    throw new Error('chatCompletion: all attempts exhausted');
  };

}