OmniRoute/open-sse/executors/gemini-cli.ts

import { BaseExecutor, mergeUpstreamExtraHeaders, mergeAbortSignals } from "./base.ts";
import { randomUUID } from "crypto";
import { PROVIDERS, OAUTH_ENDPOINTS, FETCH_TIMEOUT_MS } from "../config/constants.ts";
import { getGeminiCliHeaders } from "../services/geminiCliHeaders.ts";
import { scrubProxyAndFingerprintHeaders } from "../services/antigravityHeaderScrub.ts";
import { obfuscateSensitiveWords } from "../services/antigravityObfuscation.ts";
import {
  shouldStripCloudCodeThinking,
  stripCloudCodeThinkingConfig,
} from "../services/cloudCodeThinking.ts";

const LOAD_CODE_ASSIST_URL = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist";
const ONBOARD_USER_URL = "https://cloudcode-pa.googleapis.com/v1internal:onboardUser";
const PROJECT_TTL_MS = 30_000; // 30 seconds — matches native Gemini CLI
const MAX_CACHE_SIZE = 100;
const LOAD_CODE_ASSIST_TIMEOUT_MS = 10_000; // 10 seconds timeout
const ONBOARD_TIMEOUT_MS = 30_000;
const ONBOARD_MAX_ATTEMPTS = 10;
const ONBOARD_DELAY_MS = 5_000;
const DEFAULT_PROJECT_ID = "default-project";
const DEFAULT_ONBOARD_TIER = "free-tier";
const LOAD_CODE_ASSIST_METADATA = Object.freeze({
  ideType: "IDE_UNSPECIFIED",
  platform: "PLATFORM_UNSPECIFIED",
  pluginType: "GEMINI",
});
const ONBOARD_METADATA = Object.freeze({
  ideType: "IDE_UNSPECIFIED",
  pluginType: "GEMINI",
});

// Per-account cache: accessToken -> { projectId, expiresAt }
const projectCache = new Map<string, { projectId: string; expiresAt: number }>();
// In-flight deduplication: prevents thundering herd on cache miss
const inflightRefresh = new Map<string, Promise<string | null>>();

type LoadCodeAssistResponse = {
  cloudaicompanionProject?: string | { id?: string | null } | null;
  allowedTiers?: Array<{ id?: string | null; isDefault?: boolean | null }> | null;
};

type OnboardOptions = {
  attempts?: number;
  delayMs?: number;
};

function normalizeGeminiModel(model: string): string {
  return typeof model === "string" && model.trim().length > 0
    ? model.replace(/^models\//, "").trim()
    : "unknown";
}

function generateGeminiCliRequestId(): string {
  return `agent-${randomUUID()}`;
}

function generateGeminiCliSessionId(): string {
  return `-${Date.now()}`;
}

function cloneGeminiCliRecord(value: Record<string, any>): Record<string, any> {
  if (typeof structuredClone === "function") {
    return structuredClone(value);
  }

  return JSON.parse(JSON.stringify(value));
}

function extractProjectId(payload: unknown): string {
  if (!payload || typeof payload !== "object") return "";
  const data = payload as LoadCodeAssistResponse;
  if (typeof data.cloudaicompanionProject === "string") {
    return data.cloudaicompanionProject.trim();
  }
  if (typeof data.cloudaicompanionProject?.id === "string") {
    return data.cloudaicompanionProject.id.trim();
  }
  return "";
}

function extractDefaultTierId(payload: unknown): string {
  if (!payload || typeof payload !== "object") return DEFAULT_ONBOARD_TIER;
  const tiers = Array.isArray((payload as LoadCodeAssistResponse).allowedTiers)
    ? (payload as LoadCodeAssistResponse).allowedTiers
    : [];
  for (const tier of tiers) {
    if (tier?.isDefault && typeof tier.id === "string" && tier.id.trim()) {
      return tier.id.trim();
    }
  }
  return DEFAULT_ONBOARD_TIER;
}

function cacheProject(accessToken: string, projectId: string): void {
  if (projectCache.size >= MAX_CACHE_SIZE) {
    const now = Date.now();
    for (const [key, val] of projectCache) {
      if (val.expiresAt <= now) projectCache.delete(key);
    }
    if (projectCache.size >= MAX_CACHE_SIZE) {
      const firstKey = projectCache.keys().next().value;
      if (firstKey !== undefined) projectCache.delete(firstKey);
    }
  }

  projectCache.set(accessToken, {
    projectId,
    expiresAt: Date.now() + PROJECT_TTL_MS,
  });
}

function sleep(ms: number): Promise<void> {
  if (!ms || ms <= 0) return Promise.resolve();
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export class GeminiCLIExecutor extends BaseExecutor {
  constructor() {
    super("gemini-cli", PROVIDERS["gemini-cli"]);
  }

  buildUrl(model, stream, urlIndex = 0) {
    void model;
    void urlIndex;
    const action = stream ? "streamGenerateContent?alt=sse" : "generateContent";
    return `${this.config.baseUrl}:${action}`;
  }

  buildHeaders(
    credentials,
    stream = true,
    clientHeaders?: Record<string, string> | null,
    model?: string
  ) {
    void clientHeaders;

    // Fallback to internal tracker if model not explicitly passed (matches older interface calls)
    const activeModel = model || this._currentModel || "unknown";

    const raw = getGeminiCliHeaders(
      normalizeGeminiModel(activeModel),
      credentials.accessToken,
      stream ? "*/*" : "application/json"
    );

    if (credentials.apiKey) {
      raw["x-goog-api-key"] = credentials.apiKey;
      // getGeminiCliHeaders adds Authorization: Bearer undefined if accessToken is empty, so we clean it up
      if (!credentials.accessToken) {
        delete raw["Authorization"];
      }
    }

    return scrubProxyAndFingerprintHeaders(raw);
  }

  async onboardManagedProject(
    accessToken: string,
    tierId = DEFAULT_ONBOARD_TIER,
    options: OnboardOptions = {},
    model = "unknown"
  ): Promise<string | null> {
    const currentModel = normalizeGeminiModel(model);
    const attempts =
      Number.isInteger(options.attempts) && options.attempts! > 0
        ? Number(options.attempts)
        : ONBOARD_MAX_ATTEMPTS;
    const delayMs =
      typeof options.delayMs === "number" &&
      Number.isFinite(options.delayMs) &&
      options.delayMs >= 0
        ? options.delayMs
        : ONBOARD_DELAY_MS;

    const requestBody = {
      tierId: tierId || DEFAULT_ONBOARD_TIER,
      metadata: { ...ONBOARD_METADATA },
    };

    for (let attempt = 0; attempt < attempts; attempt += 1) {
      try {
        const controller = new AbortController();
        const timeoutId = setTimeout(() => controller.abort(), ONBOARD_TIMEOUT_MS);

        let response;
        try {
          response = await fetch(ONBOARD_USER_URL, {
            method: "POST",
            headers: getGeminiCliHeaders(currentModel, accessToken, "application/json"),
            body: JSON.stringify(requestBody),
            signal: controller.signal,
          });
        } finally {
          clearTimeout(timeoutId);
        }

        if (response.ok) {
          const payload = await response.json();
          const managedProjectId = extractProjectId(payload?.response);

          if (payload?.done === true && managedProjectId) {
            return managedProjectId;
          }

          if (payload?.done === true) {
            return DEFAULT_PROJECT_ID;
          }
        } else {
          console.warn(
            `[OmniRoute] onboardUser returned ${response.status} on attempt ${attempt + 1}`
          );
        }
      } catch (error) {
        const msg = error instanceof Error ? error.message : String(error);
        console.warn(`[OmniRoute] onboardUser attempt ${attempt + 1} failed (${msg})`);
      }

      if (attempt < attempts - 1) {
        await sleep(delayMs);
      }
    }

    return null;
  }

  /**
   * Fetch the current cloudaicompanionProject via loadCodeAssist API.
   * Native Gemini CLI refreshes this every 30 seconds — OmniRoute stores it once
   * at OAuth connection time, so it goes stale. This method keeps it fresh.
   */
  async refreshProject(accessToken: string, model = "unknown"): Promise<string | null> {
    // Check cache
    const cached = projectCache.get(accessToken);
    if (cached && cached.expiresAt > Date.now()) {
      return cached.projectId;
    }

    // Deduplicate in-flight requests (thundering herd prevention)
    const inflight = inflightRefresh.get(accessToken);
    if (inflight) return inflight;

    const promise = this._doRefresh(accessToken, model);
    inflightRefresh.set(accessToken, promise);
    try {
      return await promise;
    } finally {
      inflightRefresh.delete(accessToken);
    }
  }

  async _doRefresh(accessToken: string, model = "unknown"): Promise<string | null> {
    const currentModel = normalizeGeminiModel(model);
    try {
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), LOAD_CODE_ASSIST_TIMEOUT_MS);

      let response;
      try {
        response = await fetch(LOAD_CODE_ASSIST_URL, {
          method: "POST",
          headers: getGeminiCliHeaders(currentModel, accessToken, "application/json"),
          body: JSON.stringify({
            metadata: { ...LOAD_CODE_ASSIST_METADATA },
          }),
          signal: controller.signal,
        });
      } finally {
        clearTimeout(timeoutId);
      }

      if (!response.ok) {
        console.warn(
          `[OmniRoute] loadCodeAssist returned ${response.status} — falling back to stored projectId`
        );
        return null;
      }

      const data = (await response.json()) as LoadCodeAssistResponse;
      let projectId = extractProjectId(data);

      if (!projectId) {
        console.warn(
          "[OmniRoute] loadCodeAssist returned no project — attempting managed project onboarding"
        );
        projectId = await this.onboardManagedProject(
          accessToken,
          extractDefaultTierId(data),
          {},
          currentModel
        );
      }

      if (!projectId) {
        console.warn(
          "[OmniRoute] managed project onboarding failed — falling back to stored projectId"
        );
        return null;
      }

      cacheProject(accessToken, projectId);

      return projectId;
    } catch (error) {
      const msg = error instanceof Error ? error.message : String(error);
      console.warn(`[OmniRoute] loadCodeAssist failed (${msg}) — falling back to stored projectId`);
      return null;
    }
  }

  async transformRequest(model, body, stream, credentials) {
    const currentModel = normalizeGeminiModel(model);
    const normalizedBody =
      shouldStripCloudCodeThinking(this.provider, currentModel) && body && typeof body === "object"
        ? stripCloudCodeThinkingConfig(body)
        : body;

    const bodyRecord =
      normalizedBody && typeof normalizedBody === "object"
        ? (normalizedBody as Record<string, any>)
        : { request: {} };
    const requestRecord =
      bodyRecord.request && typeof bodyRecord.request === "object"
        ? cloneGeminiCliRecord(bodyRecord.request as Record<string, any>)
        : {};

    const storedProject =
      bodyRecord.project ||
      credentials.projectId ||
      (credentials.providerSpecificData as Record<string, unknown>)?.projectId ||
      "";

    const envelope: Record<string, any> = {
      model: currentModel,
      project: storedProject,
      user_prompt_id: bodyRecord.user_prompt_id || generateGeminiCliRequestId(),
      request: {
        ...requestRecord,
        session_id: requestRecord.session_id || generateGeminiCliSessionId(),
      },
    };

    for (const [key, value] of Object.entries(bodyRecord)) {
      if (!(key in envelope) && key !== "request") {
        envelope[key] = value;
      }
    }

    // Native Gemini CLI refreshes the Cloud Code project periodically because
    // stored project IDs can go stale. Keep the stored value as a fallback.
    if (credentials.accessToken) {
      const freshProject = await this.refreshProject(credentials.accessToken, currentModel);
      if (freshProject) {
        envelope.project = freshProject;
      }
    }

    if (envelope.request) {
      // Obfuscate sensitive client names in user content
      const contents = envelope.request?.contents;
      if (Array.isArray(contents)) {
        for (const msg of contents) {
          if (Array.isArray(msg.parts)) {
            for (const part of msg.parts) {
              if (typeof part.text === "string") {
                part.text = obfuscateSensitiveWords(part.text);
              }
            }
          }
        }
      }
    }
    return envelope;
  }

  async execute({
    model,
    body,
    stream,
    credentials,
    signal,
    log,
    upstreamExtraHeaders,
  }: ExecuteInput) {
    const fallbackCount = this.getFallbackCount();
    let lastError = null;
    let lastStatus = 0;
    const MAX_AUTO_RETRIES = 3;
    const retryAttemptsByUrl: Record<number, number> = {};

    for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) {
      const url = this.buildUrl(model, stream, urlIndex);
      const headers = this.buildHeaders(credentials, stream, null, model);
      mergeUpstreamExtraHeaders(headers, upstreamExtraHeaders);

      const transformed = await this.transformRequest(model, body, stream, credentials);
      if (transformed instanceof Response) {
        return { response: transformed, url, headers, transformedBody: body };
      }
      const transformedBody = transformed;

      if (!retryAttemptsByUrl[urlIndex]) {
        retryAttemptsByUrl[urlIndex] = 0;
      }

      try {
        log?.debug?.(
          "TELEMETRY",
          `[Gemini CLI] Execute - URL: ${url}, Model: ${model}, Retry: ${retryAttemptsByUrl[urlIndex]}`
        );

        const timeoutSignal = AbortSignal.timeout(FETCH_TIMEOUT_MS);
        const mergedSignal = signal ? mergeAbortSignals(signal, timeoutSignal) : timeoutSignal;
        const response = await fetch(url, {
          method: "POST",
          headers,
          body: JSON.stringify(transformedBody),
          signal: mergedSignal,
        });

        if (!response.ok) {
          log?.warn?.(
            "TELEMETRY",
            `[Gemini CLI] Error Response - URL: ${url}, Status: ${response.status}`
          );

          let retryMs: number | null = null;
          if (response.status === 429 || response.status === 503) {
            try {
              const errorBody = await response.clone().text();
              retryMs = this.parseRetryFromErrorMessage(errorBody);
            } catch {
              /* ignore parse error */
            }

            if ((!retryMs || retryMs <= 60000) && retryAttemptsByUrl[urlIndex] < MAX_AUTO_RETRIES) {
              retryAttemptsByUrl[urlIndex]++;
              const backoffMs =
                retryMs || Math.min(1000 * 2 ** retryAttemptsByUrl[urlIndex], 30000);
              log?.debug?.(
                "RETRY",
                `Gemini CLI 429 retry ${retryAttemptsByUrl[urlIndex]} after ${backoffMs}ms`
              );
              await sleep(backoffMs);
              urlIndex--;
              continue;
            }
          }
        }

        if (this.shouldRetry(response.status, urlIndex)) {
          lastStatus = response.status;
          continue;
        }

        return { response, url, headers, transformedBody };
      } catch (error) {
        lastError = error;
        if (urlIndex + 1 < fallbackCount) continue;
        throw error;
      }
    }

    throw lastError || new Error(`All ${fallbackCount} URLs failed with status ${lastStatus}`);
  }

  async refreshCredentials(credentials, log) {
    if (!credentials.refreshToken) return null;

    try {
      const response = await fetch(OAUTH_ENDPOINTS.google.token, {
        method: "POST",
        headers: {
          "Content-Type": "application/x-www-form-urlencoded",
          Accept: "application/json",
        },
        body: new URLSearchParams({
          grant_type: "refresh_token",
          refresh_token: credentials.refreshToken,
          client_id: this.config.clientId,
          client_secret: this.config.clientSecret,
        }),
      });

      if (!response.ok) {
        const errorText = await response.text().catch(() => "");
        log?.error?.("TOKEN", "Gemini CLI refresh failed", {
          status: response.status,
          error: errorText.slice(0, 200),
        });
        // Match refreshGoogleToken's pattern: invalid_grant means the refresh
        // token was revoked / replaced — surface as unrecoverable so the caller
        // marks the account expired instead of retrying forever with a dead token.
        try {
          const errorBody = JSON.parse(errorText);
          if (errorBody?.error === "invalid_grant") {
            return { error: "unrecoverable_refresh_error", code: "invalid_grant" } as never;
          }
        } catch {
          // not JSON — fall through
        }
        return null;
      }

      const tokens = await response.json();
      log?.info?.("TOKEN", "Gemini CLI refreshed");

      const refreshed: Record<string, unknown> = {
        accessToken: tokens.access_token,
        refreshToken: tokens.refresh_token || credentials.refreshToken,
        expiresIn: tokens.expires_in,
        projectId: credentials.projectId,
      };
      if (credentials.providerSpecificData !== undefined) {
        refreshed.providerSpecificData = credentials.providerSpecificData;
      }
      return refreshed as never;
    } catch (error) {
      log?.error?.("TOKEN", `Gemini CLI refresh error: ${error.message}`);
      return null;
    }
  }

  // Parse retry time from Gemini error message body
  // Format: "Your quota will reset after 2h7m23s"
  parseRetryFromErrorMessage(errorMessage: unknown): number | null {
    if (!errorMessage || typeof errorMessage !== "string") return null;

    const match = errorMessage.match(/reset (?:after|in) (\d+h)?(\d+m)?(\d+s)?/i);
    if (!match) return null;

    let totalMs = 0;
    if (match[1]) totalMs += parseInt(match[1]) * 3600 * 1000;
    if (match[2]) totalMs += parseInt(match[2]) * 60 * 1000;
    if (match[3]) totalMs += parseInt(match[3]) * 1000;

    return totalMs || 2_000;
  }
}

export default GeminiCLIExecutor;