diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index adf303aa4..03674ddf4 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -14,6 +14,15 @@ - `resolveAnthropicMetadataUserId()` auto-generates a cloaking user ID for OAuth requests when `metadata.user_id` is absent or invalid - `AnthropicOAuthFlow` is now exported for direct use - OAuth callback server timeout extended from 2 min to 5 min +- `parseGeminiCliCredentials()` parses Google Cloud credential JSON with support for legacy (`{token,projectId}`), alias (`project_id`/`refresh`/`expires`), and enriched formats +- `shouldRefreshGeminiCliCredentials()` and proactive token refresh before requests for both Gemini CLI and Antigravity providers (60s pre-expiry buffer) +- `normalizeAntigravityTools()` converts `parametersJsonSchema` → `parameters` in function declarations for Antigravity compatibility +- `ANTIGRAVITY_SYSTEM_INSTRUCTION` is now exported for use by search and other consumers +- `ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA` constant exported from OAuth module with `ANTIGRAVITY` ideType +- Antigravity project onboarding: `onboardProjectWithRetries()` provisions a new project via `onboardUser` LRO when `loadCodeAssist` returns no existing project (up to 5 attempts, 2s interval) +- `getOAuthApiKey` now includes `refreshToken`, `expiresAt`, `email`, and `accountId` in the Gemini/Antigravity JSON credential payload to enable proactive refresh +- Antigravity model discovery now tries the production daily endpoint first, with sandbox as fallback +- `ANTIGRAVITY_DISCOVERY_DENYLIST` filters low-quality/internal models from discovery results ### Changed @@ -29,6 +38,17 @@ - Anthropic OAuth scopes reduced to `org:create_api_key user:profile user:inference` - OAuth code exchange now strips URL fragment from callback code, using the fragment as state override when present - Claude usage headers aligned: user-agent updated to `claude-cli/2.1.63 (external, cli)`, anthropic-beta extended with full beta set +- Antigravity session ID format changed to signed decimal (negative int63 derived from SHA-256 of first user message, or random bounded int63) +- Antigravity `requestId` now uses `agent-{uuid}` format; non-Antigravity requests no longer include requestId/userAgent/requestType in the payload +- `ANTIGRAVITY_DAILY_ENDPOINT` corrected to `daily-cloudcode-pa.googleapis.com`; sandbox endpoint kept as fallback only +- Antigravity discovery: removed `recommended`/`agentModelSorts` filter; now includes all non-internal, non-denylisted models +- Antigravity discovery no longer sends `project` in the request body +- Gemini/Antigravity OAuth flows no longer use PKCE (code_challenge removed) +- Antigravity `loadCodeAssist` metadata ideType changed from `IDE_UNSPECIFIED` to `ANTIGRAVITY` +- Antigravity `discoverProject` now uses a single canonical production endpoint; falls back to project onboarding instead of a hardcoded default project ID +- `VALIDATED` tool calling config applied to Antigravity requests with Claude models +- `maxOutputTokens` removed from Antigravity generation config for non-Claude models +- System instruction injection for Antigravity scoped to Claude and `gemini-3-pro-high` models only ## [13.3.14] - 2026-02-28 diff --git a/packages/ai/src/providers/google-gemini-cli.ts b/packages/ai/src/providers/google-gemini-cli.ts index 18c064314..9d98125a9 100644 --- a/packages/ai/src/providers/google-gemini-cli.ts +++ b/packages/ai/src/providers/google-gemini-cli.ts @@ -3,7 +3,7 @@ * Shared implementation for both google-gemini-cli and google-antigravity providers. * Uses the Cloud Code Assist API endpoint to access Gemini and Claude models. */ -import { createHash } from "node:crypto"; +import { createHash, randomBytes, randomUUID } from "node:crypto"; import type { Content, FunctionCallingConfigMode, ThinkingConfig } from "@google/genai"; import { abortableSleep, readSseJson } from "@oh-my-pi/pi-utils"; import { calculateCost } from "../models"; @@ -20,6 +20,8 @@ import type { } from "../types"; import { AssistantMessageEventStream } from "../utils/event-stream"; import { appendRawHttpRequestDumpFor400, type RawHttpRequestDump, withHttpStatus } from "../utils/http-inspector"; +import { refreshAntigravityToken } from "../utils/oauth/google-antigravity"; +import { refreshGoogleCloudToken } from "../utils/oauth/google-gemini-cli"; import { sanitizeSurrogates } from "../utils/sanitize-unicode"; import { convertMessages, @@ -56,8 +58,9 @@ export interface GoogleGeminiCliOptions extends StreamOptions { } const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com"; -const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"; -const ANTIGRAVITY_ENDPOINT_FALLBACKS = [ANTIGRAVITY_DAILY_ENDPOINT, DEFAULT_ENDPOINT] as const; +const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com"; +const ANTIGRAVITY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"; +const ANTIGRAVITY_ENDPOINT_FALLBACKS = [ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_SANDBOX_ENDPOINT] as const; const GEMINI_CLI_USER_AGENT = process.env.PI_AI_GEMINI_CLI_USER_AGENT || "google-api-nodejs-client/9.15.1"; @@ -110,7 +113,7 @@ export function getAntigravityUserAgent() { } // Antigravity system instruction (compact version from CLIProxyAPI). -const ANTIGRAVITY_SYSTEM_INSTRUCTION = +export const ANTIGRAVITY_SYSTEM_INSTRUCTION = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding." + "You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question." + "**Absolute paths only**" + @@ -126,6 +129,8 @@ const MAX_EMPTY_STREAM_RETRIES = 2; const EMPTY_STREAM_BASE_DELAY_MS = 500; const RATE_LIMIT_BUDGET_MS = 5 * 60 * 1000; const CLAUDE_THINKING_BETA_HEADER = "interleaved-thinking-2025-05-14"; +const GOOGLE_GEMINI_REFRESH_SKEW_MS = 60_000; +const ANTIGRAVITY_REFRESH_SKEW_MS = 60_000; /** * Extract retry delay from Gemini error response (in milliseconds). @@ -227,11 +232,20 @@ export function extractRetryDelay(errorText: string, response?: Response | Heade return undefined; } +function isClaudeModel(modelId: string): boolean { + return modelId.toLowerCase().includes("claude"); +} + function isClaudeThinkingModel(modelId: string): boolean { const normalized = modelId.toLowerCase(); return normalized.includes("claude") && normalized.includes("thinking"); } +function shouldInjectAntigravitySystemInstruction(modelId: string): boolean { + const normalized = modelId.toLowerCase(); + return normalized.includes("claude") || normalized.includes("gemini-3-pro-high"); +} + /** * Check if an error is retryable (rate limit, server error, network error, etc.) */ @@ -258,6 +272,107 @@ function extractErrorMessage(errorText: string): string { return errorText; } +interface GeminiCliApiKeyPayload { + token?: unknown; + projectId?: unknown; + project_id?: unknown; + refreshToken?: unknown; + expiresAt?: unknown; + refresh?: unknown; + expires?: unknown; +} +interface ParsedGeminiCliCredentials { + accessToken: string; + projectId: string; + refreshToken?: string; + expiresAt?: number; +} + +function normalizeExpiryMs(value: unknown): number | undefined { + if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) { + return undefined; + } + return value < 10_000_000_000 ? value * 1000 : value; +} + +export function parseGeminiCliCredentials(apiKeyRaw: string): ParsedGeminiCliCredentials { + const invalidCredentialsMessage = "Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate."; + const missingCredentialsMessage = + "Missing token or projectId in Google Cloud credentials. Use /login to re-authenticate."; + + let parsed: GeminiCliApiKeyPayload; + try { + parsed = JSON.parse(apiKeyRaw) as GeminiCliApiKeyPayload; + } catch { + throw new Error(invalidCredentialsMessage); + } + + const projectId = + typeof parsed.projectId === "string" + ? parsed.projectId + : typeof parsed.project_id === "string" + ? parsed.project_id + : undefined; + + if (typeof parsed.token !== "string" || typeof projectId !== "string") { + throw new Error(missingCredentialsMessage); + } + + const refreshToken = + typeof parsed.refreshToken === "string" + ? parsed.refreshToken + : typeof parsed.refresh === "string" + ? parsed.refresh + : undefined; + const expiresAt = normalizeExpiryMs(parsed.expiresAt ?? parsed.expires); + + return { + accessToken: parsed.token, + projectId, + refreshToken, + expiresAt, + }; +} + +export function shouldRefreshGeminiCliCredentials( + expiresAt: number | undefined, + isAntigravity: boolean, + nowMs = Date.now(), +): boolean { + if (expiresAt === undefined) { + return false; + } + + const skewMs = isAntigravity ? ANTIGRAVITY_REFRESH_SKEW_MS : GOOGLE_GEMINI_REFRESH_SKEW_MS; + return nowMs + skewMs >= expiresAt; +} + +async function refreshGeminiCliCredentialsIfNeeded( + credentials: ParsedGeminiCliCredentials, + isAntigravity: boolean, +): Promise { + if (!credentials.refreshToken || !shouldRefreshGeminiCliCredentials(credentials.expiresAt, isAntigravity)) { + return credentials; + } + + try { + const refreshed = isAntigravity + ? await refreshAntigravityToken(credentials.refreshToken, credentials.projectId) + : await refreshGoogleCloudToken(credentials.refreshToken, credentials.projectId); + return { + accessToken: refreshed.access, + projectId: credentials.projectId, + refreshToken: refreshed.refresh, + expiresAt: refreshed.expires, + }; + } catch (error) { + if (credentials.expiresAt !== undefined && Date.now() >= credentials.expiresAt) { + const reason = error instanceof Error ? error.message : String(error); + throw new Error(`OAuth token refresh failed before request: ${reason}`); + } + return credentials; + } +} interface CloudCodeAssistRequest { project: string; model: string; @@ -349,28 +464,16 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = ( let rawRequestDump: RawHttpRequestDump | undefined; try { - // apiKey is JSON-encoded: { token, projectId } const apiKeyRaw = options?.apiKey; if (!apiKeyRaw) { throw new Error("Google Cloud Code Assist requires OAuth authentication. Use /login to authenticate."); } - let accessToken: string; - let projectId: string; - - try { - const parsed = JSON.parse(apiKeyRaw) as { token: string; projectId: string }; - accessToken = parsed.token; - projectId = parsed.projectId; - } catch { - throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate."); - } - - if (!accessToken || !projectId) { - throw new Error("Missing token or projectId in Google Cloud credentials. Use /login to re-authenticate."); - } - const isAntigravity = model.provider === "google-antigravity"; + const parsedCredentials = parseGeminiCliCredentials(apiKeyRaw); + const activeCredentials = await refreshGeminiCliCredentialsIfNeeded(parsedCredentials, isAntigravity); + const { accessToken, projectId } = activeCredentials; + const baseUrl = model.baseUrl?.trim(); const endpoints = baseUrl ? [baseUrl] : isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT]; @@ -383,7 +486,9 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = ( "Content-Type": "application/json", Accept: "text/event-stream", ...headers, - ...(isClaudeThinkingModel(model.id) ? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER } : {}), + ...(!isAntigravity && isClaudeThinkingModel(model.id) + ? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER } + : {}), ...(options?.headers ?? {}), }; const requestBodyJson = JSON.stringify(requestBody); @@ -791,33 +896,89 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = ( return stream; }; -function deriveSessionId(context: Context): string | undefined { +const INT63_MASK = (1n << 63n) - 1n; +const ANTIGRAVITY_RANDOM_BOUND = 9_000_000_000_000_000_000n; + +function formatSignedDecimalSessionId(value: bigint): string { + return `-${value.toString()}`; +} + +function deriveSignedDecimalFromHash(text: string): string { + const digest = createHash("sha256").update(text).digest(); + let value = 0n; + for (let index = 0; index < 8; index += 1) { + value = (value << 8n) | BigInt(digest[index] ?? 0); + } + return formatSignedDecimalSessionId(value & INT63_MASK); +} + +function randomBoundedInt63(maxExclusive: bigint): bigint { + while (true) { + const bytes = randomBytes(8); + let value = 0n; + for (const byte of bytes) { + value = (value << 8n) | BigInt(byte); + } + value &= INT63_MASK; + if (value < maxExclusive) { + return value; + } + } +} + +function randomSignedDecimalSessionId(): string { + return formatSignedDecimalSessionId(randomBoundedInt63(ANTIGRAVITY_RANDOM_BOUND)); +} + +function getFirstUserTextForAntigravitySession(context: Context): string | undefined { for (const message of context.messages) { - if (message.role !== "user" && message.role !== "developer") { + if (message.role !== "user") { continue; } - let text = ""; if (typeof message.content === "string") { - text = message.content; - } else if (Array.isArray(message.content)) { - text = message.content - .filter((item): item is TextContent => item.type === "text") - .map(item => item.text) - .join("\n"); + return message.content; } - if (!text || text.trim().length === 0) { - return undefined; + if (Array.isArray(message.content)) { + const firstTextPart = message.content.find((item): item is TextContent => item.type === "text"); + return firstTextPart?.text; } - const hash = createHash("sha256").update(text).digest("hex"); - return hash.slice(0, 32); + return undefined; } return undefined; } +function deriveAntigravitySessionId(context: Context): string { + const text = getFirstUserTextForAntigravitySession(context); + if (text && text.trim().length > 0) { + return deriveSignedDecimalFromHash(text); + } + + return randomSignedDecimalSessionId(); +} + +function normalizeAntigravityTools( + tools: CloudCodeAssistRequest["request"]["tools"], +): CloudCodeAssistRequest["request"]["tools"] { + return tools?.map(tool => ({ + ...tool, + functionDeclarations: tool.functionDeclarations.map(declaration => { + if (!("parametersJsonSchema" in declaration)) { + return declaration; + } + + const { parametersJsonSchema, ...rest } = declaration; + return { + ...rest, + parameters: parametersJsonSchema, + }; + }), + })); +} + export function buildRequest( model: Model<"google-gemini-cli">, context: Context, @@ -868,9 +1029,8 @@ export function buildRequest( contents, }; - const sessionId = deriveSessionId(context); - if (sessionId) { - request.sessionId = sessionId; + if (isAntigravity) { + request.sessionId = deriveAntigravitySessionId(context); } // System instruction must be object with parts, not plain string @@ -885,7 +1045,8 @@ export function buildRequest( } if (context.tools && context.tools.length > 0) { - request.tools = convertTools(context.tools, model); + const convertedTools = convertTools(context.tools, model); + request.tools = isAntigravity ? normalizeAntigravityTools(convertedTools) : convertedTools; if (options.toolChoice) { request.toolConfig = { functionCallingConfig: { @@ -895,7 +1056,22 @@ export function buildRequest( } } - if (isAntigravity) { + if (isAntigravity && !isClaudeModel(model.id) && request.generationConfig?.maxOutputTokens !== undefined) { + delete request.generationConfig.maxOutputTokens; + if (Object.keys(request.generationConfig).length === 0) { + delete request.generationConfig; + } + } + + if (isAntigravity && isClaudeModel(model.id)) { + request.toolConfig = { + functionCallingConfig: { + mode: "VALIDATED" as FunctionCallingConfigMode, + }, + }; + } + + if (isAntigravity && shouldInjectAntigravitySystemInstruction(model.id)) { const existingParts = request.systemInstruction?.parts ?? []; request.systemInstruction = { role: "user", @@ -911,8 +1087,12 @@ export function buildRequest( project: projectId, model: model.id, request, - ...(isAntigravity ? { requestType: "agent" } : {}), - userAgent: isAntigravity ? "antigravity" : "pi-coding-agent", - requestId: `${isAntigravity ? "agent" : "pi"}-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`, + ...(isAntigravity + ? { + requestType: "agent", + userAgent: "antigravity", + requestId: `agent-${randomUUID()}`, + } + : {}), }; } diff --git a/packages/ai/src/utils/discovery/antigravity.ts b/packages/ai/src/utils/discovery/antigravity.ts index 7d5d0dfee..8022edfea 100644 --- a/packages/ai/src/utils/discovery/antigravity.ts +++ b/packages/ai/src/utils/discovery/antigravity.ts @@ -3,11 +3,21 @@ import { getAntigravityUserAgent } from "../../providers/google-gemini-cli"; import type { Model } from "../../types"; import { toPositiveNumber } from "../../utils"; -const DEFAULT_ANTIGRAVITY_DISCOVERY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"; +const DEFAULT_ANTIGRAVITY_DISCOVERY_ENDPOINTS = [ + "https://daily-cloudcode-pa.googleapis.com", + "https://daily-cloudcode-pa.sandbox.googleapis.com", +] as const; const FETCH_AVAILABLE_MODELS_PATH = "/v1internal:fetchAvailableModels"; const DEFAULT_CONTEXT_WINDOW = 200_000; const DEFAULT_MAX_TOKENS = 64_000; +const ANTIGRAVITY_DISCOVERY_DENYLIST = new Set([ + "chat_20706", + "chat_23310", + "gemini-2.5-flash-thinking", + "gemini-3-pro-low", + "gemini-2.5-pro", +]); /** * Raw model metadata returned by Antigravity's `fetchAvailableModels` endpoint. @@ -142,9 +152,9 @@ const AntigravityDiscoveryApiResponseSchema: z.ZodType`. */ token: string; - /** Optional endpoint override. Defaults to Antigravity daily endpoint. */ + /** Optional endpoint override. Defaults to Antigravity fallback endpoints. */ endpoint?: string; - /** Optional project id. Defaults to an empty string for discovery. */ + /** Deprecated and ignored for antigravity discovery parity. */ project?: string; /** Optional user agent override. */ userAgent?: string; @@ -164,87 +174,78 @@ export async function fetchAntigravityDiscoveryModels( options: FetchAntigravityDiscoveryModelsOptions, ): Promise[] | null> { const fetcher = options.fetcher ?? fetch; - const endpoint = trimTrailingSlashes(options.endpoint ?? DEFAULT_ANTIGRAVITY_DISCOVERY_ENDPOINT); + const endpoints = options.endpoint + ? [trimTrailingSlashes(options.endpoint)] + : DEFAULT_ANTIGRAVITY_DISCOVERY_ENDPOINTS.map(trimTrailingSlashes); - let response: Response; - try { - response = await fetcher(`${endpoint}${FETCH_AVAILABLE_MODELS_PATH}`, { - method: "POST", - headers: { - Authorization: `Bearer ${options.token}`, - "Content-Type": "application/json", - "User-Agent": options.userAgent ?? getAntigravityUserAgent(), - }, - body: JSON.stringify({ project: options.project ?? "" }), - signal: options.signal, - }); - } catch { - return null; - } - - if (!response.ok) { - return null; - } - - let payload: unknown; - try { - payload = await response.json(); - } catch { - return null; - } - - const parsed = parseAntigravityDiscoveryResponse(payload); - if (!parsed) { - return null; - } - - const recommendedIds = collectRecommendedModelIds(parsed.agentModelSorts ?? []); - const models: Model<"google-gemini-cli">[] = []; - - for (const [modelId, model] of Object.entries(parsed.models ?? {})) { - if (model.isInternal === true) { - continue; - } - if (model.recommended !== true && !recommendedIds.has(modelId)) { + for (const endpoint of endpoints) { + let response: Response; + try { + response = await fetcher(`${endpoint}${FETCH_AVAILABLE_MODELS_PATH}`, { + method: "POST", + headers: { + Authorization: `Bearer ${options.token}`, + "Content-Type": "application/json", + "User-Agent": options.userAgent ?? getAntigravityUserAgent(), + }, + body: JSON.stringify({}), + signal: options.signal, + }); + } catch { continue; } - const supportsImages = model.supportsImages === true; - models.push({ - id: modelId, - name: model.displayName ? `${model.displayName} (Antigravity)` : modelId, - api: "google-gemini-cli", - provider: "google-antigravity", - baseUrl: endpoint, - reasoning: model.supportsThinking === true, - input: supportsImages ? ["text", "image"] : ["text"], - cost: { - input: 0, - output: 0, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: toPositiveNumber(model.maxTokens, DEFAULT_CONTEXT_WINDOW), - maxTokens: toPositiveNumber(model.maxOutputTokens, DEFAULT_MAX_TOKENS), - }); - } + if (!response.ok) { + continue; + } - models.sort((a, b) => a.name.localeCompare(b.name) || a.id.localeCompare(b.id)); - return models; -} + let payload: unknown; + try { + payload = await response.json(); + } catch { + continue; + } -function collectRecommendedModelIds(sorts: AntigravityDiscoveryAgentModelSort[]): Set { - const ids = new Set(); - for (const sort of sorts) { - for (const group of sort.groups ?? []) { - for (const modelId of group.modelIds ?? []) { - if (typeof modelId === "string" && modelId.length > 0) { - ids.add(modelId); - } + const parsed = parseAntigravityDiscoveryResponse(payload); + if (!parsed) { + continue; + } + + const models: Model<"google-gemini-cli">[] = []; + + for (const [modelId, model] of Object.entries(parsed.models ?? {})) { + if (ANTIGRAVITY_DISCOVERY_DENYLIST.has(modelId)) { + continue; } + if (model.isInternal === true) { + continue; + } + + const supportsImages = model.supportsImages === true; + models.push({ + id: modelId, + name: model.displayName ? `${model.displayName} (Antigravity)` : modelId, + api: "google-gemini-cli", + provider: "google-antigravity", + baseUrl: endpoint, + reasoning: model.supportsThinking === true, + input: supportsImages ? ["text", "image"] : ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: toPositiveNumber(model.maxTokens, DEFAULT_CONTEXT_WINDOW), + maxTokens: toPositiveNumber(model.maxOutputTokens, DEFAULT_MAX_TOKENS), + }); } + + models.sort((a, b) => a.name.localeCompare(b.name) || a.id.localeCompare(b.id)); + return models; } - return ids; + + return null; } function parseAntigravityDiscoveryResponse(value: unknown): AntigravityDiscoveryApiResponse | null { diff --git a/packages/ai/src/utils/oauth/google-antigravity.ts b/packages/ai/src/utils/oauth/google-antigravity.ts index ab181b069..67caf6c53 100644 --- a/packages/ai/src/utils/oauth/google-antigravity.ts +++ b/packages/ai/src/utils/oauth/google-antigravity.ts @@ -4,7 +4,6 @@ */ import { getAntigravityAuthHeaders } from "../../providers/google-gemini-cli"; import { OAuthCallbackFlow } from "./callback-server"; -import { generatePKCE } from "./pkce"; import type { OAuthController, OAuthCredentials } from "./types"; const decode = (s: string) => atob(s); @@ -25,7 +24,10 @@ const SCOPES = [ const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"; const TOKEN_URL = "https://oauth2.googleapis.com/token"; -const DEFAULT_PROJECT_ID = "rising-fact-p41fc"; +const CLOUD_CODE_ENDPOINT = "https://cloudcode-pa.googleapis.com"; +const TIER_LEGACY = "legacy-tier"; +const PROJECT_ONBOARD_MAX_ATTEMPTS = 5; +const PROJECT_ONBOARD_INTERVAL_MS = 2000; interface LoadCodeAssistPayload { cloudaicompanionProject?: string | { id?: string }; @@ -33,6 +35,79 @@ interface LoadCodeAssistPayload { allowedTiers?: Array<{ id?: string; isDefault?: boolean }>; } +interface LongRunningOperationResponse { + done?: boolean; + response?: { + cloudaicompanionProject?: string | { id?: string }; + }; +} + +export const ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA = Object.freeze({ + ideType: "ANTIGRAVITY", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", +}); + +function readProjectId(value: string | { id?: string } | undefined): string | undefined { + if (typeof value === "string" && value.length > 0) { + return value; + } + if (value && typeof value === "object" && typeof value.id === "string" && value.id.length > 0) { + return value.id; + } + return undefined; +} + +function getDefaultTierId(allowedTiers?: Array<{ id?: string; isDefault?: boolean }>): string { + if (!allowedTiers || allowedTiers.length === 0) { + return TIER_LEGACY; + } + const defaultTier = allowedTiers.find(tier => tier.isDefault && typeof tier.id === "string" && tier.id.length > 0); + if (defaultTier?.id) { + return defaultTier.id; + } + return TIER_LEGACY; +} + +async function onboardProjectWithRetries( + endpoint: string, + headers: Record, + onboardBody: { tierId: string; metadata: typeof ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA }, + onProgress?: (message: string) => void, +): Promise { + for (let attempt = 1; attempt <= PROJECT_ONBOARD_MAX_ATTEMPTS; attempt += 1) { + if (attempt > 1) { + onProgress?.(`Waiting for project provisioning (attempt ${attempt}/${PROJECT_ONBOARD_MAX_ATTEMPTS})...`); + await Bun.sleep(PROJECT_ONBOARD_INTERVAL_MS); + } + + const onboardResponse = await fetch(`${endpoint}/v1internal:onboardUser`, { + method: "POST", + headers, + body: JSON.stringify(onboardBody), + }); + + if (!onboardResponse.ok) { + const errorText = await onboardResponse.text(); + throw new Error(`onboardUser failed: ${onboardResponse.status} ${onboardResponse.statusText}: ${errorText}`); + } + + const operation = (await onboardResponse.json()) as LongRunningOperationResponse; + if (!operation.done) { + continue; + } + + const projectId = readProjectId(operation.response?.cloudaicompanionProject); + if (projectId) { + return projectId; + } + } + + throw new Error( + `onboardUser did not return a provisioned project id after ${PROJECT_ONBOARD_MAX_ATTEMPTS} attempts`, + ); +} + async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise { const headers = { Authorization: `Bearer ${accessToken}`, @@ -40,45 +115,41 @@ async function discoverProject(accessToken: string, onProgress?: (message: strin ...getAntigravityAuthHeaders(), }; - const endpoints = ["https://cloudcode-pa.googleapis.com", "https://daily-cloudcode-pa.sandbox.googleapis.com"]; - onProgress?.("Checking for existing project..."); + const endpoint = CLOUD_CODE_ENDPOINT; + try { + const loadResponse = await fetch(`${endpoint}/v1internal:loadCodeAssist`, { + method: "POST", + headers, + body: JSON.stringify({ + metadata: ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA, + }), + }); - for (const endpoint of endpoints) { - try { - const loadResponse = await fetch(`${endpoint}/v1internal:loadCodeAssist`, { - method: "POST", - headers, - body: JSON.stringify({ - metadata: { - ideType: "IDE_UNSPECIFIED", - platform: "PLATFORM_UNSPECIFIED", - pluginType: "GEMINI", - }, - }), - }); - - if (loadResponse.ok) { - const data = (await loadResponse.json()) as LoadCodeAssistPayload; - - if (typeof data.cloudaicompanionProject === "string" && data.cloudaicompanionProject) { - return data.cloudaicompanionProject; - } - if ( - data.cloudaicompanionProject && - typeof data.cloudaicompanionProject === "object" && - data.cloudaicompanionProject.id - ) { - return data.cloudaicompanionProject.id; - } - } - } catch { - // Try next endpoint + if (!loadResponse.ok) { + const errorText = await loadResponse.text(); + throw new Error(`loadCodeAssist failed: ${loadResponse.status} ${loadResponse.statusText}: ${errorText}`); } - } - onProgress?.("Using default project..."); - return DEFAULT_PROJECT_ID; + const loadPayload = (await loadResponse.json()) as LoadCodeAssistPayload; + const existingProject = readProjectId(loadPayload.cloudaicompanionProject); + if (existingProject) { + return existingProject; + } + + const tierId = getDefaultTierId(loadPayload.allowedTiers); + onProgress?.("Provisioning project..."); + const onboardBody = { + tierId, + metadata: ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA, + }; + const provisionedProject = await onboardProjectWithRetries(endpoint, headers, onboardBody, onProgress); + return provisionedProject; + } catch (error) { + throw new Error( + `Could not discover or provision an Antigravity project. ${error instanceof Error ? error.message : String(error)}`, + ); + } } async function getUserEmail(accessToken: string): Promise { @@ -98,25 +169,16 @@ async function getUserEmail(accessToken: string): Promise { } class AntigravityOAuthFlow extends OAuthCallbackFlow { - #verifier: string = ""; - #challenge: string = ""; - constructor(ctrl: OAuthController) { super(ctrl, CALLBACK_PORT, CALLBACK_PATH); } async generateAuthUrl(state: string, redirectUri: string): Promise<{ url: string; instructions?: string }> { - const pkce = await generatePKCE(); - this.#verifier = pkce.verifier; - this.#challenge = pkce.challenge; - const authParams = new URLSearchParams({ client_id: CLIENT_ID, response_type: "code", redirect_uri: redirectUri, scope: SCOPES.join(" "), - code_challenge: this.#challenge, - code_challenge_method: "S256", state, access_type: "offline", prompt: "consent", @@ -138,7 +200,6 @@ class AntigravityOAuthFlow extends OAuthCallbackFlow { code, grant_type: "authorization_code", redirect_uri: redirectUri, - code_verifier: this.#verifier, }), }); diff --git a/packages/ai/src/utils/oauth/google-gemini-cli.ts b/packages/ai/src/utils/oauth/google-gemini-cli.ts index 99b9e0f6a..a2a047ab0 100644 --- a/packages/ai/src/utils/oauth/google-gemini-cli.ts +++ b/packages/ai/src/utils/oauth/google-gemini-cli.ts @@ -6,7 +6,6 @@ import { $env } from "@oh-my-pi/pi-utils"; import { getGeminiCliHeaders } from "../../providers/google-gemini-cli"; import { OAuthCallbackFlow } from "./callback-server"; -import { generatePKCE } from "./pkce"; import type { OAuthController, OAuthCredentials } from "./types"; const decode = (s: string) => atob(s); @@ -228,25 +227,16 @@ async function getUserEmail(accessToken: string): Promise { } class GeminiCliOAuthFlow extends OAuthCallbackFlow { - #verifier: string = ""; - #challenge: string = ""; - constructor(ctrl: OAuthController) { super(ctrl, CALLBACK_PORT, CALLBACK_PATH); } async generateAuthUrl(state: string, redirectUri: string): Promise<{ url: string; instructions?: string }> { - const pkce = await generatePKCE(); - this.#verifier = pkce.verifier; - this.#challenge = pkce.challenge; - const authParams = new URLSearchParams({ client_id: CLIENT_ID, response_type: "code", redirect_uri: redirectUri, scope: SCOPES.join(" "), - code_challenge: this.#challenge, - code_challenge_method: "S256", state, access_type: "offline", prompt: "consent", @@ -268,7 +258,6 @@ class GeminiCliOAuthFlow extends OAuthCallbackFlow { code, grant_type: "authorization_code", redirect_uri: redirectUri, - code_verifier: this.#verifier, }), }); diff --git a/packages/ai/src/utils/oauth/index.ts b/packages/ai/src/utils/oauth/index.ts index 0318e45bc..04af7be8f 100644 --- a/packages/ai/src/utils/oauth/index.ts +++ b/packages/ai/src/utils/oauth/index.ts @@ -375,8 +375,8 @@ function getPerplexityJwtExpiryMs(token: string): number | undefined { * Get API key for a provider from OAuth credentials. * Automatically refreshes expired tokens. * - * For google-gemini-cli and antigravity, returns JSON-encoded { token, projectId } - * + * For google-gemini-cli and antigravity, returns JSON-encoded credentials including token/projectId + * plus refresh/expiry metadata for proactive refresh support. * @returns API key string, or null if no credentials * @throws Error if refresh fails */ @@ -416,7 +416,16 @@ export async function getOAuthApiKey( } // For providers that need projectId, return JSON const needsProjectId = provider === "google-gemini-cli" || provider === "google-antigravity"; - const apiKey = needsProjectId ? JSON.stringify({ token: creds.access, projectId: creds.projectId }) : creds.access; + const apiKey = needsProjectId + ? JSON.stringify({ + token: creds.access, + projectId: creds.projectId, + refreshToken: creds.refresh, + expiresAt: creds.expires, + email: creds.email, + accountId: creds.accountId, + }) + : creds.access; return { newCredentials: creds, apiKey }; } diff --git a/packages/ai/test/anthropic-alignment.test.ts b/packages/ai/test/anthropic-alignment.test.ts index c190e610e..c4555056f 100644 --- a/packages/ai/test/anthropic-alignment.test.ts +++ b/packages/ai/test/anthropic-alignment.test.ts @@ -284,8 +284,8 @@ describe("Anthropic request fingerprint alignment", () => { }); it("treats tool prefix helpers as no-ops when prefix is empty", () => { - expect(applyClaudeToolPrefix("Read")).toBe("Read"); - expect(stripClaudeToolPrefix("proxy_Read")).toBe("proxy_Read"); + expect(applyClaudeToolPrefix("Read", "")).toBe("Read"); + expect(stripClaudeToolPrefix("proxy_Read", "")).toBe("proxy_Read"); }); it("does not prefix built-in Anthropic tool names when prefix is configured", () => { diff --git a/packages/ai/test/google-antigravity-auth.test.ts b/packages/ai/test/google-antigravity-auth.test.ts new file mode 100644 index 000000000..abea29c76 --- /dev/null +++ b/packages/ai/test/google-antigravity-auth.test.ts @@ -0,0 +1,29 @@ +import { describe, expect, it } from "bun:test"; +import { getAntigravityAuthHeaders } from "../src/providers/google-gemini-cli"; +import { ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA } from "../src/utils/oauth/google-antigravity"; + +describe("Google Antigravity auth alignment", () => { + it("uses ANTIGRAVITY ideType in loadCodeAssist metadata payload", () => { + expect(ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA).toEqual({ + ideType: "ANTIGRAVITY", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", + }); + }); + + it("keeps auth header client-metadata aligned with CLI defaults", () => { + const headers = getAntigravityAuthHeaders(); + const rawMetadata = headers["Client-Metadata"]; + expect(rawMetadata).toBeDefined(); + const metadata = JSON.parse(rawMetadata) as { + ideType?: string; + platform?: string; + pluginType?: string; + }; + expect(metadata).toEqual({ + ideType: "IDE_UNSPECIFIED", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", + }); + }); +}); diff --git a/packages/ai/test/google-gemini-cli-alignment.test.ts b/packages/ai/test/google-gemini-cli-alignment.test.ts new file mode 100644 index 000000000..5496d1705 --- /dev/null +++ b/packages/ai/test/google-gemini-cli-alignment.test.ts @@ -0,0 +1,144 @@ +import { describe, expect, it } from "bun:test"; +import { + buildRequest, + parseGeminiCliCredentials, + shouldRefreshGeminiCliCredentials, +} from "../src/providers/google-gemini-cli"; +import type { Context, Model } from "../src/types"; +import { getOAuthApiKey } from "../src/utils/oauth"; + +function createModel(provider: "google-gemini-cli" | "google-antigravity"): Model<"google-gemini-cli"> { + return { + id: provider === "google-antigravity" ? "gemini-3-flash" : "gemini-2.5-flash", + name: provider, + api: "google-gemini-cli", + provider, + baseUrl: "https://example.com", + reasoning: false, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 200000, + maxTokens: 8192, + }; +} + +function createContext(): Context { + return { + messages: [{ role: "user", content: "implement token refresh", timestamp: Date.now() }], + }; +} + +describe("Google Gemini CLI alignment", () => { + it("encodes enriched OAuth JSON while preserving token + projectId", async () => { + const expiresAt = Date.now() + 60 * 60 * 1000; + const result = await getOAuthApiKey("google-gemini-cli", { + "google-gemini-cli": { + access: "access-token", + refresh: "refresh-token", + expires: expiresAt, + projectId: "proj-123", + email: "dev@example.com", + accountId: "acct-1", + }, + }); + + expect(result).not.toBeNull(); + const payload = JSON.parse(result!.apiKey) as { + token?: string; + projectId?: string; + refreshToken?: string; + expiresAt?: number; + email?: string; + accountId?: string; + }; + expect(payload.token).toBe("access-token"); + expect(payload.projectId).toBe("proj-123"); + expect(payload.refreshToken).toBe("refresh-token"); + expect(payload.expiresAt).toBe(expiresAt); + expect(payload.email).toBe("dev@example.com"); + expect(payload.accountId).toBe("acct-1"); + }); + + it("accepts legacy, alias, and enriched OAuth JSON payloads", () => { + const legacy = parseGeminiCliCredentials(JSON.stringify({ token: "legacy-token", projectId: "proj-legacy" })); + expect(legacy).toEqual({ + accessToken: "legacy-token", + projectId: "proj-legacy", + refreshToken: undefined, + expiresAt: undefined, + }); + + const aliasPayload = parseGeminiCliCredentials( + JSON.stringify({ + token: "alias-token", + project_id: "proj-alias", + refresh: "refresh-alias", + expires: 1_737_000_000, + }), + ); + expect(aliasPayload).toEqual({ + accessToken: "alias-token", + projectId: "proj-alias", + refreshToken: "refresh-alias", + expiresAt: 1_737_000_000_000, + }); + + const enriched = parseGeminiCliCredentials( + JSON.stringify({ + token: "enriched-token", + projectId: "proj-enriched", + refreshToken: "refresh-token", + expiresAt: 1_737_000_000_000, + }), + ); + expect(enriched).toEqual({ + accessToken: "enriched-token", + projectId: "proj-enriched", + refreshToken: "refresh-token", + expiresAt: 1_737_000_000_000, + }); + }); + + it("avoids excessive antigravity refresh churn with pre-buffered OAuth expiry", () => { + const issuedAt = 1_700_000_000_000; + const preBufferedExpiry = issuedAt + 55 * 60 * 1000; + + expect(shouldRefreshGeminiCliCredentials(preBufferedExpiry, true, issuedAt + 10 * 60 * 1000)).toBe(false); + expect(shouldRefreshGeminiCliCredentials(preBufferedExpiry, true, issuedAt + 54 * 60 * 1000)).toBe(true); + expect(shouldRefreshGeminiCliCredentials(preBufferedExpiry, false, issuedAt + 54 * 60 * 1000)).toBe(true); + }); + it("omits antigravity-only metadata in non-antigravity request payloads", () => { + const model = createModel("google-gemini-cli"); + const payload = buildRequest(model, createContext(), "proj-123", {}, false) as { + request: { sessionId?: string }; + requestType?: string; + userAgent?: string; + requestId?: string; + }; + + expect(payload.request.sessionId).toBeUndefined(); + expect(payload.requestType).toBeUndefined(); + expect(payload.userAgent).toBeUndefined(); + expect(payload.requestId).toBeUndefined(); + }); + + it("keeps antigravity metadata in antigravity request payloads", () => { + const model = createModel("google-antigravity"); + const payload = buildRequest(model, createContext(), "proj-123", {}, true) as { + request: { sessionId?: string }; + requestType?: string; + userAgent?: string; + requestId?: string; + }; + + expect(payload.request.sessionId).toMatch(/^-[0-9]+$/); + expect(payload.requestType).toBe("agent"); + expect(payload.userAgent).toBe("antigravity"); + expect(payload.requestId).toMatch(/^agent-/); + }); +}); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 4ad014742..74649307c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -12,6 +12,15 @@ - Both AST tools now report `scopePath`, `files`, and per-file match/replacement counts in tool details - Task item `id` max length raised from 32 to 48 characters - Anthropic web search provider now uses `buildAnthropicSearchHeaders` (dedicated search header builder separate from inference headers) +- Gemini web search provider: endpoint fallback (daily → sandbox) with retry on 429/5xx +- Gemini web search now injects Antigravity system instruction and aligned request metadata (`requestType`, `userAgent`, `requestId`) for Antigravity credentials +- `buildGeminiRequestTools()` helper for composable Gemini tool configuration (googleSearch, codeExecution, urlContext) +- Web search schema exposes `max_tokens`, `temperature`, and `num_search_results` as tool parameters +- Web search provider fallback: when an explicit provider is unavailable, resolves the auto chain instead of returning empty results + +### Fixed + +- `gemini_image` tool: corrected `responseModalities` values from `'Image'`/`'Text'` to uppercase `'IMAGE'`/`'TEXT'` matching the API enum ### Breaking Changes diff --git a/packages/coding-agent/src/tools/gemini-image.ts b/packages/coding-agent/src/tools/gemini-image.ts index 4362e948f..3a9e31fac 100644 --- a/packages/coding-agent/src/tools/gemini-image.ts +++ b/packages/coding-agent/src/tools/gemini-image.ts @@ -27,7 +27,7 @@ interface ImageApiKey { projectId?: string; } -const responseModalitySchema = StringEnum(["Image", "Text"]); +const responseModalitySchema = StringEnum(["IMAGE", "TEXT"]); const aspectRatioSchema = StringEnum(["1:1", "3:4", "4:3", "9:16", "16:9"], { description: "Aspect ratio (1:1, 3:4, 4:3, 9:16, 16:9).", }); @@ -536,7 +536,7 @@ function buildAntigravityRequest( contents: [{ role: "user", parts }], systemInstruction: { parts: [{ text: IMAGE_SYSTEM_INSTRUCTION }] }, generationConfig: { - responseModalities: ["Image"], + responseModalities: ["IMAGE"], imageConfig, candidateCount: 1, }, @@ -788,7 +788,7 @@ export const geminiImageTool: CustomTool; details: SearchRenderDetails }> { const providers = - params.provider && params.provider !== "auto" && params.no_fallback + params.provider && params.provider !== "auto" ? (await getSearchProvider(params.provider).isAvailable()) ? [getSearchProvider(params.provider)] - : [] + : await resolveProviderChain("auto") : await resolveProviderChain(params.provider); - if (providers.length === 0) { const message = "No web search provider configured."; return { diff --git a/packages/coding-agent/src/web/search/providers/base.ts b/packages/coding-agent/src/web/search/providers/base.ts index 7023b989f..e09f13505 100644 --- a/packages/coding-agent/src/web/search/providers/base.ts +++ b/packages/coding-agent/src/web/search/providers/base.ts @@ -10,6 +10,9 @@ export interface SearchParams { maxOutputTokens?: number; numSearchResults?: number; temperature?: number; + googleSearch?: Record; + codeExecution?: Record; + urlContext?: Record; } /** Base class for web search providers. */ diff --git a/packages/coding-agent/src/web/search/providers/gemini.ts b/packages/coding-agent/src/web/search/providers/gemini.ts index 1dac642a7..93d1222d3 100644 --- a/packages/coding-agent/src/web/search/providers/gemini.ts +++ b/packages/coding-agent/src/web/search/providers/gemini.ts @@ -5,7 +5,12 @@ * Requires OAuth credentials stored in agent.db for provider "google-gemini-cli" or "google-antigravity". * Returns synthesized answers with citations and source metadata from grounding chunks. */ -import { getAntigravityHeaders, getGeminiCliHeaders, refreshGoogleCloudToken } from "@oh-my-pi/pi-ai"; +import { + ANTIGRAVITY_SYSTEM_INSTRUCTION, + getAntigravityHeaders, + getGeminiCliHeaders, + refreshGoogleCloudToken, +} from "@oh-my-pi/pi-ai"; import { getAgentDbPath } from "@oh-my-pi/pi-utils"; import { AgentStorage } from "../../../session/agent-storage"; import type { SearchCitation, SearchResponse, SearchSource } from "../../../web/search/types"; @@ -14,10 +19,18 @@ import type { SearchParams } from "./base"; import { SearchProvider } from "./base"; const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com"; -const ANTIGRAVITY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"; +const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com"; +const ANTIGRAVITY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"; +const ANTIGRAVITY_ENDPOINT_FALLBACKS = [ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_SANDBOX_ENDPOINT] as const; const DEFAULT_MODEL = "gemini-2.5-flash"; -export interface GeminiSearchParams { +interface GeminiToolParams { + google_search?: Record; + code_execution?: Record; + url_context?: Record; +} + +export interface GeminiSearchParams extends GeminiToolParams { query: string; system_prompt?: string; num_results?: number; @@ -27,6 +40,17 @@ export interface GeminiSearchParams { temperature?: number; } +export function buildGeminiRequestTools(params: GeminiToolParams): Array>> { + const tools: Array>> = [{ googleSearch: params.google_search ?? {} }]; + if (params.code_execution !== undefined) { + tools.push({ codeExecution: params.code_execution }); + } + if (params.url_context !== undefined) { + tools.push({ urlContext: params.url_context }); + } + return tools; +} + /** OAuth credential stored in agent.db */ interface GeminiOAuthCredential { type: "oauth"; @@ -48,15 +72,15 @@ interface GeminiAuth { /** * Finds valid Gemini OAuth credentials from agent.db. - * Checks google-antigravity first (daily sandbox, more quota), then google-gemini-cli (prod). + * Checks google-gemini-cli first (stable prod), then google-antigravity (daily sandbox). * @returns OAuth credential with access token and project ID, or null if none found */ export async function findGeminiAuth(): Promise { const expiryBuffer = 5 * 60 * 1000; // 5 minutes const now = Date.now(); - // Try providers in order: antigravity first (more quota), then gemini-cli - const providers = ["google-antigravity", "google-gemini-cli"] as const; + // Try providers in deterministic order: gemini-cli first, then antigravity + const providers = ["google-gemini-cli", "google-antigravity"] as const; try { const storage = await AgentStorage.open(getAgentDbPath()); @@ -180,6 +204,7 @@ async function callGeminiSearch( systemPrompt?: string, maxOutputTokens?: number, temperature?: number, + toolParams: GeminiToolParams = {}, ): Promise<{ answer: string; sources: SearchSource[]; @@ -188,10 +213,31 @@ async function callGeminiSearch( model: string; usage?: { inputTokens: number; outputTokens: number; totalTokens: number }; }> { - const endpoint = auth.isAntigravity ? ANTIGRAVITY_ENDPOINT : DEFAULT_ENDPOINT; - const url = `${endpoint}/v1internal:streamGenerateContent?alt=sse`; + const endpoints = auth.isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT]; const headers = auth.isAntigravity ? getAntigravityHeaders() : getGeminiCliHeaders(); + const requestMetadata = auth.isAntigravity + ? { + requestType: "agent", + userAgent: "antigravity", + requestId: `agent-${crypto.randomUUID()}`, + } + : { + userAgent: "pi-coding-agent", + requestId: `pi-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`, + }; + + const normalizedSystemPrompt = systemPrompt?.toWellFormed(); + const systemInstructionParts: Array<{ text: string }> = [ + ...(auth.isAntigravity + ? [ + { text: ANTIGRAVITY_SYSTEM_INSTRUCTION }, + { text: `Please ignore following [ignore]${ANTIGRAVITY_SYSTEM_INSTRUCTION}[/ignore]` }, + ] + : []), + ...(normalizedSystemPrompt ? [{ text: normalizedSystemPrompt }] : []), + ]; + const requestBody: Record = { project: auth.projectId, model: DEFAULT_MODEL, @@ -202,16 +248,15 @@ async function callGeminiSearch( parts: [{ text: query }], }, ], - // Add googleSearch tool for grounding - tools: [{ googleSearch: {} }], - ...(systemPrompt && { + tools: buildGeminiRequestTools(toolParams), + ...(systemInstructionParts.length > 0 && { systemInstruction: { - parts: [{ text: systemPrompt }], + ...(auth.isAntigravity ? { role: "user" } : {}), + parts: systemInstructionParts, }, }), }, - userAgent: "pi-web-search", - requestId: `search-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`, + ...requestMetadata, }; if (maxOutputTokens !== undefined || temperature !== undefined) { @@ -224,17 +269,52 @@ async function callGeminiSearch( } (requestBody.request as Record).generationConfig = generationConfig; } + let response: Response | undefined; + for (let endpointIndex = 0; endpointIndex < endpoints.length; endpointIndex++) { + const url = `${endpoints[endpointIndex]}/v1internal:streamGenerateContent?alt=sse`; + try { + response = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${auth.accessToken}`, + "Content-Type": "application/json", + Accept: "text/event-stream", + ...headers, + }, + body: JSON.stringify(requestBody), + }); + } catch (error) { + if (auth.isAntigravity && endpointIndex < endpoints.length - 1) { + continue; + } + throw error; + } - const response = await fetch(url, { - method: "POST", - headers: { - Authorization: `Bearer ${auth.accessToken}`, - "Content-Type": "application/json", - Accept: "text/event-stream", - ...headers, - }, - body: JSON.stringify(requestBody), - }); + if (response.ok) { + break; + } + + const errorText = await response.text(); + const isRetryableStatus = + response.status === 429 || + response.status === 500 || + response.status === 502 || + response.status === 503 || + response.status === 504; + if (auth.isAntigravity && isRetryableStatus && endpointIndex < endpoints.length - 1) { + continue; + } + + throw new SearchProviderError( + "gemini", + `Gemini Cloud Code API error (${response.status}): ${errorText}`, + response.status, + ); + } + + if (!response) { + throw new SearchProviderError("gemini", "Gemini API request failed", 500); + } if (!response.ok) { const errorText = await response.text(); @@ -396,6 +476,11 @@ export async function searchGemini(params: GeminiSearchParams): Promise | null; +}; + +const SSE_RESPONSE = + 'data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"Gemini answer"}]}}],"modelVersion":"gemini-2.5-flash"}}\n\n'; + +describe("searchGemini tools serialization", () => { + const originalFetch = globalThis.fetch; + let capturedRequest: CapturedRequest | null = null; + + beforeEach(() => { + capturedRequest = null; + vi.spyOn(AgentStorage, "open").mockResolvedValue({ + listAuthCredentials: () => [ + { + id: 1, + credential: { + type: "oauth", + access: "test-access-token", + expires: Date.now() + 600_000, + projectId: "test-project", + }, + }, + ], + updateAuthCredential: () => undefined, + } as unknown as AgentStorage); + + globalThis.fetch = mock(async (_url: string | URL | Request, init?: RequestInit) => { + capturedRequest = { + body: init?.body ? (JSON.parse(init.body as string) as Record) : null, + }; + return new Response(SSE_RESPONSE, { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }); + }) as unknown as typeof fetch; + }); + + afterEach(() => { + vi.restoreAllMocks(); + globalThis.fetch = originalFetch; + capturedRequest = null; + }); + + it("sends default googleSearch tool when no passthrough payloads are provided", async () => { + await searchGemini({ query: "default tools" }); + + expect(capturedRequest).not.toBeNull(); + expect(capturedRequest?.body?.request).toMatchObject({ + tools: [{ googleSearch: {} }], + }); + }); + + it("passes through google_search payload into googleSearch tool", async () => { + await searchGemini({ + query: "google payload", + google_search: { dynamicRetrievalConfig: { mode: "MODE_DYNAMIC" } }, + }); + + expect(capturedRequest).not.toBeNull(); + expect(capturedRequest?.body?.request).toMatchObject({ + tools: [{ googleSearch: { dynamicRetrievalConfig: { mode: "MODE_DYNAMIC" } } }], + }); + }); + + it("includes codeExecution and urlContext tools when provided", async () => { + await searchGemini({ + query: "extended tools", + code_execution: {}, + url_context: { allowedDomains: ["example.com"] }, + }); + + expect(capturedRequest).not.toBeNull(); + expect(capturedRequest?.body?.request).toMatchObject({ + tools: [{ googleSearch: {} }, { codeExecution: {} }, { urlContext: { allowedDomains: ["example.com"] } }], + }); + }); +});