diff --git a/docs-site/src/content/docs/guides/providers.md b/docs-site/src/content/docs/guides/providers.md index e1451dcfbd..ac0cf8c99b 100644 --- a/docs-site/src/content/docs/guides/providers.md +++ b/docs-site/src/content/docs/guides/providers.md @@ -115,7 +115,7 @@ ocx logout | `kimi` | `openai-chat` | `https://api.kimi.com/coding/v1` | Kimi K2.7/K2.6/K2.5 coding models. | | `nous` | `openai-chat` | `https://inference-api.nousresearch.com/v1` | Nous Research subscription gateway (same backend Hermes Agent uses). Device-grant login against `portal.nousresearch.com`; the access token is the per-request inference JWT. Mixed paid + `:free` model catalog (`tencent/hy3:free`, `stepfun/step-3.7-flash:free`, ...) discovered live from the signed-in account. Refresh tokens are single-use and rotated on every refresh. | | `kiro` | `kiro` | `https://runtime.us-east-1.kiro.dev` | Initial login imports the installed, signed-in `kiro-cli` session (on Unix, install with `curl -fsSL https://cli.kiro.dev/install` | `bash`; on Windows PowerShell, use `irm 'https://cli.kiro.dev/install.ps1'` | `iex`; then run `kiro-cli login`). **Add account** logs `kiro-cli` out, starts a fresh browser login that switches the account used by `kiro-cli`, and stores account-scoped profile metadata. Existing OpenCodex accounts are preserved, and cancellation or failure restores the previous `kiro-cli` session. | -| `google-antigravity` | `google` | `https://daily-cloudcode-pa.googleapis.com` | Google OAuth over the Cloud Code Assist wire. Live discovery uses CCA's authenticated `v1internal:fetchAvailableModels` endpoint and publishes the agent models available to the signed-in account; the maintained catalog remains the fallback. | +| `google-antigravity` | `google` | `https://daily-cloudcode-pa.googleapis.com` | Google OAuth over the Cloud Code Assist wire. Live discovery uses CCA's authenticated `v1internal:fetchAvailableModels` endpoint and publishes the agent models available to the signed-in account; the maintained catalog remains the fallback. See [Claude on Antigravity](#claude-on-antigravity-cloud-code-assist) below. | | `cursor` | `cursor` | `https://api2.cursor.sh` | Experimental PKCE login, live HTTP/2 transport with an opt-in HTTP/1.1 compatibility path, and account-filtered model discovery. | | `github-copilot` | `openai-chat` | `https://api.githubcopilot.com` | Experimental. GitHub device flow + `copilot_internal` exchange (VS Code OAuth client). Requires an active Copilot subscription; not an official third-party API. | @@ -153,6 +153,23 @@ cat accounts.json | ocx account import google-antigravity --format cockpit-tools Inline JSON and extra positional arguments are rejected. Keep exported files private and delete or store them securely after import. +### Claude on Antigravity (Cloud Code Assist) + +The `google-antigravity` provider routes Claude models through Google's Cloud Code Assist (Antigravity) +wire rather than Anthropic's native API. opencodex translates requests and responses at the Gemini +format envelope: tool use/result pairing follows Anthropic semantics (including stable `functionCall.id` +/ `functionResponse.id` fields), and Claude thinking blocks keep their `thoughtSignature` values across +turns. + +CCA Claude models reject histories that end with an assistant (model) turn — upstream treats that as +prefill. opencodex strips trailing model turns when safe and appends a `(continue)` user nudge when the +history would otherwise end on model output (for example after context compaction or interrupted-turn +replay). Histories that already end on a user message or tool result are left unchanged. + +Antigravity exposes only SSE transport. Unary (non-streaming) callers still go through the same +`parseStream` path; plain JSON bodies without `data:` framing are rejected as truncated SSE rather +than parsed as a separate JSON response format. + ### OAuth reliability opencodex coordinates token refresh and Codex pool routing so concurrent requests do not race the diff --git a/src/adapters/google-antigravity-hosts.ts b/src/adapters/google-antigravity-hosts.ts new file mode 100644 index 0000000000..a2c1fe096d --- /dev/null +++ b/src/adapters/google-antigravity-hosts.ts @@ -0,0 +1,26 @@ +const DAILY_ANTIGRAVITY_HOST = "https://daily-cloudcode-pa.googleapis.com"; +const PROD_ANTIGRAVITY_HOST = "https://cloudcode-pa.googleapis.com"; + +/** + * Return the configured Antigravity endpoint and, for Google's known daily/prod hosts + * only, its daily/production peer. Custom baseUrl values stay single-host. + */ +export function antigravityHostCandidates(configuredBase: string): string[] { + const configured = configuredBase.replace(/\/+$/, ""); + if (configured === DAILY_ANTIGRAVITY_HOST) { + return [DAILY_ANTIGRAVITY_HOST, PROD_ANTIGRAVITY_HOST]; + } + if (configured === PROD_ANTIGRAVITY_HOST) { + return [PROD_ANTIGRAVITY_HOST, DAILY_ANTIGRAVITY_HOST]; + } + return [configured]; +} + +/** OAuth bearer requests must not use a cleartext host, even if generic baseUrl config allows http. */ +export function isAntigravityHttpsHost(host: string): boolean { + try { + return new URL(host).protocol === "https:"; + } catch { + return false; + } +} diff --git a/src/adapters/google-antigravity-tools.ts b/src/adapters/google-antigravity-tools.ts new file mode 100644 index 0000000000..15a930c3b6 --- /dev/null +++ b/src/adapters/google-antigravity-tools.ts @@ -0,0 +1,94 @@ +import type { + OcxAssistantMessage, + OcxMessage, + OcxToolCall, + OcxToolResultMessage, +} from "../types"; + +function isAssistantToolCall(message: OcxMessage): message is OcxAssistantMessage { + return message.role === "assistant"; +} + +function isToolResult(message: OcxMessage): message is OcxToolResultMessage { + return message.role === "toolResult"; +} + +/** + * Repair incomplete tool exchanges before assigning provider-visible ids. + * + * CCA translates Gemini function calls and responses into Anthropic tool blocks, + * which requires both sides of every exchange. A result is valid only when its + * call appeared earlier in the history, and a call is valid only when a result + * appears later. Filtering the history first also prevents orphan results from + * reserving ids in the request-scoped allocator. + * + * The allocator maps one raw id to one wire id, so a second complete exchange + * that reuses the same raw id would serialize as a colliding pair. Keep only + * the first matched occurrence per raw id. + */ +export function repairGoogleToolPairs(messages: readonly OcxMessage[]): OcxMessage[] { + const pendingCalls = new Map>(); + const seenRawCallIds = new Set(); + const matchedCallParts = new Set(); + const matchedResultIndexes = new Set(); + + const enqueueCall = (id: string, messageIndex: number, partIndex: number) => { + if (seenRawCallIds.has(id)) return; + seenRawCallIds.add(id); + const queue = pendingCalls.get(id) ?? []; + queue.push({ messageIndex, partIndex }); + pendingCalls.set(id, queue); + }; + + for (let messageIndex = 0; messageIndex < messages.length; messageIndex++) { + const message = messages[messageIndex]!; + if (isAssistantToolCall(message)) { + message.content.forEach((part, partIndex) => { + if (part.type !== "toolCall") return; + enqueueCall((part as OcxToolCall).id, messageIndex, partIndex); + }); + continue; + } + if (!isToolResult(message)) continue; + const queue = pendingCalls.get(message.toolCallId); + const slot = queue?.shift(); + if (!slot) continue; + matchedCallParts.add(`${slot.messageIndex}:${slot.partIndex}`); + matchedResultIndexes.add(messageIndex); + } + + const repaired: OcxMessage[] = []; + for (const [messageIndex, message] of messages.entries()) { + if (isToolResult(message)) { + if (matchedResultIndexes.has(messageIndex)) repaired.push(message); + continue; + } + if (!isAssistantToolCall(message)) { + repaired.push(message); + continue; + } + + const content = message.content.filter((part, partIndex) => + part.type !== "toolCall" || matchedCallParts.has(`${messageIndex}:${partIndex}`)); + if (content.length > 0) { + repaired.push(content.length === message.content.length ? message : { ...message, content }); + } + } + return repaired; +} + +/** + * Claude interprets a final model turn as a prefilled assistant response. + * CCA expects the next turn to be generated instead, except when that model + * turn is the entire conversation and must remain as the initial context. + */ +export function stripTrailingClaudePrefill(contents: unknown[]): boolean { + let strippedModelTail = false; + while (contents.length >= 2) { + const last = contents[contents.length - 1]; + if (typeof last !== "object" || last === null || (last as { role?: unknown }).role !== "model") break; + contents.pop(); + strippedModelTail = true; + } + return strippedModelTail; +} diff --git a/src/adapters/google-errors.ts b/src/adapters/google-errors.ts index 69e6d0cef5..371d2b70b9 100644 --- a/src/adapters/google-errors.ts +++ b/src/adapters/google-errors.ts @@ -15,6 +15,12 @@ function googleErrorDetail(payloadText: string): { message?: string; status?: st }; } +const ANTIGRAVITY_GEO_BLOCKED_MARKER = "user location is not supported for the api use"; + +export function isAntigravityGeoBlockedBody(payloadText: string): boolean { + return payloadText.toLowerCase().includes(ANTIGRAVITY_GEO_BLOCKED_MARKER); +} + function classifyGoogle(label: string, status: number | undefined, enumStatus: string | undefined, text: string): string { const lower = `${enumStatus ?? ""} ${text}`.toLowerCase(); const quotaExhausted = @@ -29,6 +35,7 @@ function classifyGoogle(label: string, status: number | undefined, enumStatus: s if (status === 401 || enumStatus === "UNAUTHENTICATED" || lower.includes("unauthenticated") || lower.includes("invalid authentication") || lower.includes("expired")) { return `${label} authentication failed`; } + if (isAntigravityGeoBlockedBody(lower)) return `${label} location not supported`; if (status === 403 || enumStatus === "PERMISSION_DENIED" || lower.includes("permission denied") || lower.includes("access denied")) { return `${label} access denied`; } diff --git a/src/adapters/google.ts b/src/adapters/google.ts index f152d91741..28cf10db81 100644 --- a/src/adapters/google.ts +++ b/src/adapters/google.ts @@ -21,6 +21,8 @@ import { fetchAntigravityWithRetry, fetchVertexWithRetry } from "./google-http"; import { safeAntigravityHttpErrorMessage, safeVertexHttpErrorMessage } from "./google-errors"; import { isVertexTruncatedTurn, vertexTruncationErrorMessage } from "./google-truncation"; import { ANTIGRAVITY_REQUEST_UA, antigravitySessionId, isLikelyRealThoughtSignature, sanitizeAntigravityClaudeSignatures } from "./google-antigravity-wire"; +import { repairGoogleToolPairs, stripTrailingClaudePrefill } from "./google-antigravity-tools"; +import { isAntigravityHttpsHost } from "./google-antigravity-hosts"; import { compileGoogleWireBody } from "./google-wire-compiler"; import { identifyRoutedModel } from "./identity"; import { antigravityUsesReplayCache, applyAntigravityReplay, clearAntigravityReplay, observeAntigravityReplay } from "./google-antigravity-replay"; @@ -159,9 +161,10 @@ function messagesToGeminiFormat( const systemInstruction = { parts: [{ text: systemText }] }; const contents: unknown[] = []; + const messages = repairGoogleToolPairs(parsed.context.messages); const callIds = createToolCallIdAllocator(); - for (const msg of parsed.context.messages) { + for (const msg of messages) { if (msg.role === "assistant") { for (const part of (msg as OcxAssistantMessage).content) { if (part.type === "toolCall") callIds.reserve((part as OcxToolCall).id); @@ -170,7 +173,7 @@ function messagesToGeminiFormat( callIds.reserve((msg as OcxToolResultMessage).toolCallId); } } - for (const msg of parsed.context.messages) { + for (const msg of messages) { switch (msg.role) { case "user": case "developer": { @@ -315,6 +318,29 @@ function usageFromGemini(usage: Record | undefined): OcxUsage | */ const MAX_RESPONSE_BYTES = 100 * 1024 * 1024; const MAX_SSE_FRAME_BYTES = MAX_RESPONSE_BYTES; +let sseFrameMaxBytes = MAX_SSE_FRAME_BYTES; + +/** Test-only: lower the SSE frame byte cap without allocating a 100 MiB fixture. */ +export function setGoogleSseFrameMaxBytesForTests(bytes?: number): void { + sseFrameMaxBytes = bytes ?? MAX_SSE_FRAME_BYTES; +} + +function scanSseLineBytes(incompleteLineBytes: number, incoming: Uint8Array): { + maximum: number; + residual: number; +} { + let lineBytes = incompleteLineBytes; + let maximum = lineBytes; + for (const byte of incoming) { + if (byte === 0x0a) { + lineBytes = 0; + continue; + } + lineBytes += 1; + maximum = Math.max(maximum, lineBytes); + } + return { maximum, residual: lineBytes }; +} // Note: imagen-* models use a different API surface (prediction/image-generation // schema) and must NOT be treated as responseModalities-capable Gemini models. @@ -448,8 +474,9 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte } if (Object.keys(generationConfig).length > 0) body.generationConfig = generationConfig; - const method = parsed.stream ? "streamGenerateContent" : "generateContent"; - const streamParam = parsed.stream ? "?alt=sse" : ""; + const ccaAlwaysSse = provider.googleMode === "cloud-code-assist"; + const method = ccaAlwaysSse || parsed.stream ? "streamGenerateContent" : "generateContent"; + const streamParam = ccaAlwaysSse || parsed.stream ? "?alt=sse" : ""; const headers: Record = { "Content-Type": "application/json" }; if (provider.headers) Object.assign(headers, provider.headers); @@ -459,7 +486,10 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte if (!token) throw new Error("google-antigravity oauth token missing — run ocx login google-antigravity"); const base = provider.baseUrl?.trim(); if (!base) throw new Error("google-antigravity requires a non-empty baseUrl"); - const url = `${base}/v1internal:${method}${streamParam}`; + if (!isAntigravityHttpsHost(base)) { + throw new Error("google-antigravity requires an HTTPS baseUrl"); + } + const url = `${base.replace(/\/+$/, "")}/v1internal:${method}${streamParam}`; const project = provider.project; if (!project) throw new Error("Antigravity requires a discovered Cloud Code Assist project id (re-run `ocx login google-antigravity`)."); const sessionId = antigravitySessionId(parsed); @@ -484,6 +514,9 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte // nested) — matching CLIProxyAPI `generateStableSessionID`. An extra top-level/snake_case // spelling is a non-first-party key, so we send the single canonical location. const draftRequest: Record = { ...body, sessionId }; + if (systemInstruction) { + draftRequest.preambleConfig = { mode: "SYSTEM_INSTRUCTION_MODE_REPLACE" }; + } // Claude-on-Antigravity forces VALIDATED function calling (the real client always sets it). if (/claude/i.test(wireModelId)) { // VALIDATED would defeat a client's tool_choice "none": honor it by dropping the @@ -498,10 +531,14 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte } const compiled = compileGoogleWireBody(draftRequest); const request = compiled.body; + if (systemInstruction) { + request.preambleConfig = { mode: "SYSTEM_INSTRUCTION_MODE_REPLACE" }; + } restoreGoogleToolName = compiled.restoreToolName; // Compile names before replay: signatures are keyed by the exact provider-visible name. if (Array.isArray((request as { contents?: unknown[] }).contents)) { const contents = (request as { contents: unknown[] }).contents; + const strippedModelTail = /claude/i.test(wireModelId) ? stripTrailingClaudePrefill(contents) : false; if (antigravityUsesReplayCache(wireModelId)) { applyAntigravityReplay(wireModelId, sessionId, contents); } else { @@ -512,9 +549,10 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte // must end with a user message." Context compaction, previous_response_id expansion, // and interrupted-turn replay can all produce a model-tail history. Append a user // "(continue)" nudge, mirroring the anthropic adapter's tail guard (src/adapters/anthropic.ts). + // When a trailing model turn was stripped, append even if the history now ends with user. if (/claude/i.test(wireModelId)) { const last = contents.length > 0 ? contents[contents.length - 1] as { role?: string } : undefined; - if (!last || last.role === "model") { + if (strippedModelTail || !last || last.role === "model") { contents.push({ role: "user", parts: [{ text: "(continue)" }] }); } } @@ -532,6 +570,9 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte }; headers["User-Agent"] = ANTIGRAVITY_REQUEST_UA; headers["Authorization"] = `Bearer ${token}`; + if (/claude/i.test(wireModelId)) { + headers["anthropic-beta"] = "interleaved-thinking-2025-05-14"; + } return { url, method: "POST", headers, body: JSON.stringify(envelope) }; } @@ -598,6 +639,10 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte const budgetEncoder = new TextEncoder(); let buffer = ""; let bufferBytes = 0; + // Raw unterminated-line bytes, independent of TextDecoder's pending UTF-8 + // state. `bufferBytes` is the decoded residual and can undercount by 1–3 + // bytes when a chunk ends mid-character. + let incompleteLineBytes = 0; let pendingUsage: OcxUsage | undefined; let toolCallsStarted = 0; let lastFinishReason: string | undefined; @@ -608,8 +653,8 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte const handleDataLine = async function* (line: string): AsyncGenerator { const payload = line.slice(5).trim(); if (!payload) return "continue"; - if (payload.length > MAX_SSE_FRAME_BYTES) { - yield { type: "error", message: `upstream SSE data frame exceeds ${MAX_SSE_FRAME_BYTES} bytes` }; + if (budgetEncoder.encode(payload).byteLength > sseFrameMaxBytes) { + yield { type: "error", message: `upstream SSE data frame exceeds ${sseFrameMaxBytes} bytes` }; return "terminate"; } let emittedContentEvent = false; @@ -756,6 +801,20 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte while (true) { const { done, value } = await reader.read(); if (done) break; + const incoming = value ?? new Uint8Array(); + // Cap each incomplete line on raw bytes before decode and before waiting + // for a newline — otherwise a single unterminated data: payload can grow + // without bound, buffer.length is UTF-16 units, and a mid-character + // decode residual undercounts the true line. Reset at each newline so + // several sub-cap frames in one network chunk are not rejected as one + // oversized frame. + const lineScan = scanSseLineBytes(incompleteLineBytes, incoming); + if (lineScan.maximum > sseFrameMaxBytes) { + yield { type: "error", message: `upstream SSE data frame exceeds ${sseFrameMaxBytes} bytes` }; + try { await reader.cancel(); } catch { /* ignore */ } + return; + } + incompleteLineBytes = lineScan.residual; const nextBuffer = buffer + decoder.decode(value, { stream: true }); const nextBufferBytes = budgetEncoder.encode(nextBuffer).byteLength; const appendReservation = budget.reserveTransient(nextBufferBytes, { kind: "live_transient" }); @@ -763,13 +822,6 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte appendReservation.commitRetained(); budget.releaseRetained(bufferBytes, { kind: "live_transient" }); bufferBytes = nextBufferBytes; - // Cap incomplete frames before waiting for a newline — otherwise a single - // unterminated data: payload can grow without bound. - if (buffer.length > MAX_SSE_FRAME_BYTES) { - yield { type: "error", message: `upstream SSE data frame exceeds ${MAX_SSE_FRAME_BYTES} bytes` }; - try { await reader.cancel(); } catch { /* ignore */ } - return; - } const lines = buffer.split("\n"); buffer = lines.pop() ?? ""; @@ -842,6 +894,15 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte }, async parseResponse(response: Response, budget: TranslatorBudget): Promise { + // Cloud Code Assist exposes only the SSE transport. Unary callers still use this + // buffered adapter entry point, so collect the exact same events parseStream emits + // instead of maintaining a second CCA JSON parser. + if (provider.googleMode === "cloud-code-assist") { + const events: AdapterEvent[] = []; + for await (const event of this.parseStream(response, budget)) events.push(event); + retainTranslatedEventBatch(events, budget); + return events; + } // Reject oversized responses before JSON parse. Prefer Content-Length when // present and truthful; always stream-read with a hard byte cap so a missing // or lying Content-Length cannot force a full in-memory buffer + parse. @@ -912,15 +973,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte const err = raw.error as { message?: string }; return finish([{ type: "error", message: err.message ?? "upstream error" }]); } - // Antigravity (CCA) nests the standard Gemini payload under `response`; unwrap it. - let json = raw; - if (provider.googleMode === "cloud-code-assist") { - const wrapped = raw.response; - if (!wrapped || typeof wrapped !== "object" || Array.isArray(wrapped)) { - return finish([{ type: "error", message: "google-antigravity response missing response wrapper" }]); - } - json = wrapped as Record; - } + const json = raw; const events: AdapterEvent[] = []; const candidates = json.candidates as { content?: { parts?: GoogleResponsePart[] }; finishReason?: string }[] | undefined; @@ -932,10 +985,9 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte if (candidates?.[0]?.content?.parts) { // Non-streaming Google-family response: observe thought signatures for the next turn, // using the same transport-scoped namespace as the streaming path. - const replayModel = provider.googleMode === "cloud-code-assist" ? antigravityModel : vertexReplayModel; - const replaySession = provider.googleMode === "cloud-code-assist" ? antigravitySession : vertexReplaySession; - if ((provider.googleMode === "cloud-code-assist" || provider.googleMode === "vertex") - && replayModel && replaySession) { + const replayModel = vertexReplayModel; + const replaySession = vertexReplaySession; + if (provider.googleMode === "vertex" && replayModel && replaySession) { observeAntigravityReplay(replayModel, replaySession, candidates[0].content.parts as unknown[]); } let pendingThoughtSig: string | undefined; @@ -977,7 +1029,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte // Fail-closed truncation, same as the stream path: a non-stream turn cut off mid tool call // (MAX_TOKENS / MALFORMED_FUNCTION_CALL) surfaces an error instead of a silent done. - if ((provider.googleMode === "vertex" || provider.googleMode === "cloud-code-assist") + if (provider.googleMode === "vertex" && isVertexTruncatedTurn(candidates?.[0]?.finishReason, toolCallsStarted)) { return finish([{ type: "error", message: vertexTruncationErrorMessage(candidates?.[0]?.finishReason) }]); } diff --git a/src/providers/antigravity-quota.ts b/src/providers/antigravity-quota.ts new file mode 100644 index 0000000000..c4f0cb58f2 --- /dev/null +++ b/src/providers/antigravity-quota.ts @@ -0,0 +1,216 @@ +import { antigravityUserAgent } from "../adapters/client-fingerprint"; +import { antigravityHostCandidates, isAntigravityHttpsHost } from "../adapters/google-antigravity-hosts"; +import { readProviderQuotaJsonForTests } from "./quota"; +import type { ProviderQuota, ProviderQuotaWindow } from "./quota"; + +const LIVE_QUOTA_PATH = "/v1internal:retrieveUserQuota"; +const LIVE_SUMMARY_PATH = "/v1internal:retrieveUserQuotaSummary"; + +type FetchImpl = typeof fetch; + +export interface AntigravityLiveQuotaArgs { + accessToken: string; + projectId: string; + baseUrl: string; + timeoutMs: number; + fetchImpl?: FetchImpl; +} + +interface QuotaCandidate { + record: Record; + path: string[]; +} + +function asRecord(value: unknown): Record | null { + return value && typeof value === "object" && !Array.isArray(value) + ? value as Record + : null; +} + +function finiteNumber(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +function normalizePercent(value: unknown): number | undefined { + const numeric = finiteNumber(value); + return numeric === undefined ? undefined : Math.max(0, Math.min(100, numeric)); +} + +function resetAt(value: unknown): number | undefined { + const numeric = finiteNumber(value); + if (numeric !== undefined && numeric > 0) return numeric > 10_000_000_000 ? numeric : numeric * 1000; + if (typeof value !== "string" || !value.trim()) return undefined; + const parsed = Date.parse(value); + return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined; +} + +function remainingPercent(record: Record): number | undefined { + const fraction = finiteNumber(record.remainingFraction); + if (fraction !== undefined) return normalizePercent(fraction * 100); + const percentage = finiteNumber( + record.remainingPercentage + ?? record.remainingPercent + ?? record.remaining_percent, + ); + if (percentage !== undefined) return normalizePercent(percentage); + return undefined; +} + +function usedPercent(record: Record): number | undefined { + const remaining = remainingPercent(record); + return remaining === undefined ? undefined : normalizePercent(100 - remaining); +} + +function recordResetAt(record: Record): number | undefined { + return resetAt(record.resetTime ?? record.resetAt ?? record.resetsAt ?? record.reset_time ?? record.nextReset); +} + +function collectCandidates(value: unknown, path: string[] = [], output: QuotaCandidate[] = []): QuotaCandidate[] { + if (Array.isArray(value)) { + for (const [index, item] of value.entries()) collectCandidates(item, [...path, String(index)], output); + return output; + } + const record = asRecord(value); + if (!record) return output; + output.push({ record, path }); + for (const [key, child] of Object.entries(record)) { + if (child && typeof child === "object") collectCandidates(child, [...path, key], output); + } + return output; +} + +function candidateModelName(record: Record): string { + const explicit = record.modelId + ?? record.model_id + ?? record.modelName + ?? record.model + ?? record.name + ?? record.displayName; + return typeof explicit === "string" ? explicit.toLowerCase() : ""; +} + +function parseGeminiWindow(payload: unknown): ProviderQuotaWindow | undefined { + for (const candidate of collectCandidates(payload)) { + if (!candidateModelName(candidate.record).includes("gemini")) continue; + const percent = usedPercent(candidate.record); + if (percent === undefined) continue; + const reset = recordResetAt(candidate.record); + return { + label: "Gem", + percent, + ...(reset !== undefined ? { resetAt: reset } : {}), + }; + } + return undefined; +} + +function isWeeklyPath(path: string[]): boolean { + const leaf = path.at(-1); + return typeof leaf === "string" && /weekly|week|seven[_-]?day/i.test(leaf); +} + +function parseWeeklyWindow(payload: unknown): { percent: number; resetAt?: number } | undefined { + const candidates = collectCandidates(payload); + for (const candidate of candidates) { + if (!isWeeklyPath(candidate.path)) continue; + const percent = usedPercent(candidate.record); + if (percent === undefined) continue; + const reset = recordResetAt(candidate.record); + return { percent, ...(reset !== undefined ? { resetAt: reset } : {}) }; + } + return undefined; +} + +async function readJson(response: Response, timeoutMs: number): Promise { + const payload = await readProviderQuotaJsonForTests(response, timeoutMs); + if (payload === null) throw new Error("Antigravity quota RPC returned unreadable JSON"); + return payload; +} + +export class AntigravityQuotaRpcError extends Error { + constructor(readonly status: number) { + super(`Antigravity quota RPC failed: ${status}`); + } +} + +export function isTerminalAntigravityQuotaStatus(status: number): boolean { + return status === 401 || status === 403 || status === 429; +} + +function terminalRpcError(result: PromiseSettledResult): AntigravityQuotaRpcError | null { + if (result.status !== "rejected" || !(result.reason instanceof AntigravityQuotaRpcError)) return null; + return isTerminalAntigravityQuotaStatus(result.reason.status) ? result.reason : null; +} + +async function fetchRpc( + fetchImpl: FetchImpl, + host: string, + method: "retrieveUserQuota" | "retrieveUserQuotaSummary", + args: AntigravityLiveQuotaArgs, +): Promise { + const path = method === "retrieveUserQuota" ? LIVE_QUOTA_PATH : LIVE_SUMMARY_PATH; + const response = await fetchImpl(`${host}${path}`, { + method: "POST", + headers: { + Accept: "application/json", + "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), + Authorization: `Bearer ${args.accessToken}`, + }, + body: JSON.stringify({ project: args.projectId }), + redirect: "error", + signal: AbortSignal.timeout(args.timeoutMs), + }); + if (!response.ok) throw new AntigravityQuotaRpcError(response.status); + return readJson(response, args.timeoutMs); +} + +async function fetchHostQuota( + fetchImpl: FetchImpl, + host: string, + args: AntigravityLiveQuotaArgs, +): Promise { + const [quotaResult, summaryResult] = await Promise.allSettled([ + fetchRpc(fetchImpl, host, "retrieveUserQuota", args), + fetchRpc(fetchImpl, host, "retrieveUserQuotaSummary", args), + ]); + const terminalError = terminalRpcError(quotaResult) ?? terminalRpcError(summaryResult); + if (terminalError) throw terminalError; + if (quotaResult.status === "rejected") return null; + const quotaPayload = quotaResult.value; + const summaryPayload = summaryResult.status === "fulfilled" ? summaryResult.value : null; + const gem = parseGeminiWindow(quotaPayload); + const weekly = parseWeeklyWindow(summaryPayload); + if (!gem && !weekly) return null; + return { + ...(gem ? { customWindows: [gem] } : {}), + ...(weekly ? { + weeklyPercent: weekly.percent, + ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), + } : {}), + updatedAt: Date.now(), + }; +} + +export async function fetchAntigravityLiveQuota( + args: AntigravityLiveQuotaArgs, +): Promise { + const fetchImpl = args.fetchImpl ?? fetch; + for (const host of antigravityHostCandidates(args.baseUrl)) { + if (!isAntigravityHttpsHost(host)) continue; + try { + const quota = await fetchHostQuota(fetchImpl, host, args); + if (quota) return quota; + } catch (error) { + if (error instanceof AntigravityQuotaRpcError && isTerminalAntigravityQuotaStatus(error.status)) { + throw error; + } + } + } + return null; +} diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 24b06ef7a9..c03e0b804a 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -28,6 +28,12 @@ import { type CodexCapacityAggregation, type CodexCapacityQuota, } from "./codex-capacity"; +import { + AntigravityQuotaRpcError, + fetchAntigravityLiveQuota, + isTerminalAntigravityQuotaStatus, +} from "./antigravity-quota"; +import { antigravityHostCandidates, isAntigravityHttpsHost } from "../adapters/google-antigravity-hosts"; /** Match oauth/index REFRESH_SKEW_MS — use stored access without refresh when still fresh. */ const ACCOUNT_TOKEN_SKEW_MS = 60_000; @@ -1993,13 +1999,13 @@ function antigravityUsedPercent(quotaInfo: Record): number | un const remaining = normalizePercent(toFiniteNumber(quotaInfo.remainingFraction) !== undefined ? toFiniteNumber(quotaInfo.remainingFraction)! * 100 : toFiniteNumber(quotaInfo.remainingPercentage) !== undefined - ? toFiniteNumber(quotaInfo.remainingPercentage)! * 100 + ? toFiniteNumber(quotaInfo.remainingPercentage)! : undefined); if (remaining === undefined) return undefined; return normalizePercent(100 - remaining); } -async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise { +async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise { const credential = getCredential("google-antigravity"); if (!credential?.projectId) return null; let accessToken: string; @@ -2009,39 +2015,81 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig return null; } const baseUrl = (config.baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); - const response = await fetch(`${baseUrl}/v1internal:fetchAvailableModels`, { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - "User-Agent": antigravityUserAgent(), - Authorization: `Bearer ${accessToken}`, - }, - body: JSON.stringify({ project: credential.projectId }), - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - const body = asRecord(await readQuotaJson(response)); - const models = asRecord(body?.models); - if (!models) return null; + let liveQuota: ProviderQuota | null; + try { + liveQuota = await fetchAntigravityLiveQuota({ + accessToken, + projectId: credential.projectId, + baseUrl, + timeoutMs: REQUEST_TIMEOUT_MS, + }); + } catch (error) { + if (error instanceof AntigravityQuotaRpcError && isTerminalAntigravityQuotaStatus(error.status)) { + return TERMINAL_QUOTA_FAILURE; + } + liveQuota = null; + } const windows = new Map(); - for (const [modelId, rawModelInfo] of Object.entries(models)) { - const modelInfo = asRecord(rawModelInfo); - if (!modelInfo) continue; - for (const quotaInfo of quotaInfoEntries(modelInfo)) { - const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo); - if (!label || windows.has(label)) continue; - const percent = antigravityUsedPercent(quotaInfo); - if (percent === undefined) continue; - windows.set(label, { - label, - percent, - ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}), + for (const [index, host] of antigravityHostCandidates(baseUrl).entries()) { + if (!isAntigravityHttpsHost(host)) continue; + try { + const response = await fetch(`${host}/v1internal:fetchAvailableModels`, { + method: "POST", + headers: { + Accept: "application/json", + "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), + Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ project: credential.projectId }), + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); + if (!response.ok) { + if (index === 0 && (response.status === 404 || response.status === 503)) continue; + break; + } + const body = asRecord(await readQuotaJson(response)); + const models = asRecord(body?.models); + if (models) { + for (const [modelId, rawModelInfo] of Object.entries(models)) { + const modelInfo = asRecord(rawModelInfo); + if (!modelInfo) continue; + for (const quotaInfo of quotaInfoEntries(modelInfo)) { + const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo); + if (!label || windows.has(label)) continue; + const percent = antigravityUsedPercent(quotaInfo); + if (percent === undefined) continue; + windows.set(label, { + label, + percent, + ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}), + }); + } + } + } + break; + } catch { + if (index === 0) continue; + break; } } + if (liveQuota) { + const liveWindows = liveQuota.customWindows ?? []; + const catalogClaude = windows.get("Cla"); + const customWindows = [ + ...liveWindows, + ...(liveWindows.some(window => window.label === "Cla") || !catalogClaude ? [] : [catalogClaude]), + ]; + return report(provider, "google-antigravity:retrieveUserQuota", { + ...liveQuota, + ...(customWindows.length > 0 ? { customWindows } : {}), + updatedAt: Date.now(), + }); + } + const customWindows = ["Gem", "Cla"].flatMap(label => { const window = windows.get(label); return window ? [window] : []; diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts new file mode 100644 index 0000000000..0a910d959d --- /dev/null +++ b/tests/antigravity-quota.test.ts @@ -0,0 +1,523 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fetchAntigravityLiveQuota } from "../src/providers/antigravity-quota"; +import { + clearProviderQuotaCache, + fetchProviderQuotaReports, + QUOTA_RESPONSE_MAX_BYTES, +} from "../src/providers/quota"; +import { saveCredential } from "../src/oauth/store"; +import type { OcxConfig } from "../src/types"; + +const originalFetch = globalThis.fetch; +const previousOpencodexHome = process.env.OPENCODEX_HOME; +let opencodexHome: string; + +const DAILY_HOST = "https://daily-cloudcode-pa.googleapis.com"; +const PROD_HOST = "https://cloudcode-pa.googleapis.com"; +const TOKEN = "antigravity-access-token"; +const PROJECT = "antigravity-project"; + +function liveGeminiQuota(): Response { + return jsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.4, resetTime: "2026-08-19T12:00:00Z" }, + ], + }); +} + +function liveWeeklySummary(): Response { + return jsonResponse({ + weekly: { remainingPercentage: 75, resetTime: "2026-08-25T00:00:00Z" }, + }); +} + +function config(baseUrl = DAILY_HOST): OcxConfig { + return { + defaultProvider: "google-antigravity", + providers: { + "google-antigravity": { adapter: "google", authMode: "oauth", baseUrl }, + }, + } as OcxConfig; +} + +function jsonResponse(value: unknown, status = 200): Response { + return new Response(JSON.stringify(value), { + status, + headers: { "content-type": "application/json" }, + }); +} + +function catalogResponse(): Response { + return jsonResponse({ + models: { + "gemini-3.6-flash-medium": { + displayName: "Gemini 3.6 Flash (Medium)", + quotaInfo: { remainingFraction: 0.64, resetTime: "2026-08-20T14:00:00Z" }, + }, + "claude-sonnet-4.6": { + displayName: "Claude Sonnet", + quotaInfo: { remainingFraction: 0.21, resetTime: "2026-08-21T15:00:00Z" }, + }, + }, + }); +} + +function oversizedJsonResponse(value: Record): Response { + return new Response(JSON.stringify({ + ...value, + padding: "x".repeat(QUOTA_RESPONSE_MAX_BYTES), + }), { + status: 200, + headers: { "content-type": "application/json" }, + }); +} + +beforeEach(async () => { + opencodexHome = mkdtempSync(join(tmpdir(), "ocx-antigravity-quota-")); + process.env.OPENCODEX_HOME = opencodexHome; + await saveCredential("google-antigravity", { + access: TOKEN, + refresh: "antigravity-refresh-token", + expires: Date.now() + 3_600_000, + projectId: PROJECT, + }); + clearProviderQuotaCache(); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + clearProviderQuotaCache(); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + rmSync(opencodexHome, { recursive: true, force: true }); +}); + +test("does not classify a daily bucket nested under a weekly ancestor as weekly", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: TOKEN, + projectId: PROJECT, + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + weekly: { daily: { remainingPercentage: 90 } }, + }); + } + return jsonResponse({}, 404); + }, + }); + + expect(result?.weeklyPercent).toBeUndefined(); +}); + +test("does not classify an unlabelled daily summary window as weekly", async () => { + const requestOptions: Array<{ url: string; init?: RequestInit }> = []; + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + requestOptions.push({ url, init }); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + daily: { remainingFraction: 0.75 }, + }); + } + return jsonResponse({}, 404); + }, + }); + + expect(result?.customWindows?.[0]?.label).toBe("Gem"); + expect(result?.weeklyPercent).toBeUndefined(); + expect( + requestOptions + .filter(({ url }) => url.includes(":retrieveUserQuota")) + .map(({ init }) => init?.redirect), + ).toEqual(["error", "error"]); +}); + +test("keeps the daily quota when the summary RPC fails", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) return jsonResponse({}, 503); + return jsonResponse({}, 404); + }, + }); + + expect(result?.customWindows).toEqual([{ label: "Gem", percent: 50 }]); + expect(result?.weeklyPercent).toBeUndefined(); +}); + +test("treats a daily quota JSON read failure as an RPC failure", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return oversizedJsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.01 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + weekly: { remainingPercentage: 75 }, + }); + } + return jsonResponse({}, 404); + }, + }); + + expect(result).toBeNull(); +}); + +test("does not parse gemini from ancestor JSON path keys", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: TOKEN, + projectId: PROJECT, + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + "gemini-quotas": { + items: [{ remainingFraction: 0.5 }], + }, + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) return jsonResponse({}, 404); + return jsonResponse({}, 404); + }, + }); + + expect(result?.customWindows).toBeUndefined(); +}); + +describe("Antigravity live quota", () => { + test("merges live Gemini and weekly quota with catalog-only Claude windows", async () => { + const requestOptions: Array<{ url: string; init?: RequestInit }> = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + requestOptions.push({ url, init }); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.4, resetTime: "2026-08-19T12:00:00Z" }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + weekly: { remainingPercentage: 75, resetTime: "2026-08-25T00:00:00Z" }, + }); + } + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + const report = result.reports[0]; + + expect(report?.source).toBe("google-antigravity:retrieveUserQuota"); + expect(report?.quota.customWindows).toEqual([ + { label: "Gem", percent: 60, resetAt: Date.parse("2026-08-19T12:00:00Z") }, + { label: "Cla", percent: 79, resetAt: Date.parse("2026-08-21T15:00:00Z") }, + ]); + expect(report?.quota.weeklyPercent).toBe(25); + expect(report?.quota.weeklyResetAt).toBe(Date.parse("2026-08-25T00:00:00Z")); + expect(requestOptions.every(({ init }) => init?.redirect === "error")).toBe(true); + }); + + test("retries the production host after daily retrieveUserQuota returns 404", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.startsWith(DAILY_HOST) && url.includes(":retrieveUserQuota")) return jsonResponse({}, 404); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested).toContain(`${PROD_HOST}/v1internal:retrieveUserQuota`); + expect(result.reports[0]?.source).toBe("google-antigravity:retrieveUserQuota"); + }); + + test("falls back to the catalog when both live RPCs return 404", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes(":retrieveUserQuota")) return jsonResponse({}, 404); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + const report = result.reports[0]; + + expect(report?.source).toBe("google-antigravity:fetchAvailableModels"); + expect(report?.quota.customWindows).toEqual([ + { label: "Gem", percent: 36, resetAt: Date.parse("2026-08-20T14:00:00Z") }, + { label: "Cla", percent: 79, resetAt: Date.parse("2026-08-21T15:00:00Z") }, + ]); + expect(report?.quota.weeklyPercent).toBeUndefined(); + }); + + test("falls back to the catalog when live RPC fetch throws", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes(":retrieveUserQuota")) throw new Error("simulated timeout"); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + await expect(fetchProviderQuotaReports(config(), true)).resolves.toMatchObject({ + reports: [{ + source: "google-antigravity:fetchAvailableModels", + quota: { customWindows: expect.any(Array) }, + }], + }); + }); + + test("fails open to the catalog when live RPC bodies exceed the quota JSON limit", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return oversizedJsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.01 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return oversizedJsonResponse({ + weekly: { remainingPercentage: 1 }, + }); + } + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + const report = result.reports[0]; + + expect(report?.source).toBe("google-antigravity:fetchAvailableModels"); + expect(report?.quota.customWindows).toEqual([ + { label: "Gem", percent: 36, resetAt: Date.parse("2026-08-20T14:00:00Z") }, + { label: "Cla", percent: 79, resetAt: Date.parse("2026-08-21T15:00:00Z") }, + ]); + expect(report?.quota.weeklyPercent).toBeUndefined(); + }); + + test("does not fetch production or catalog after daily retrieveUserQuota returns 401", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 401); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return jsonResponse({}, 404); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); + }); + + test("drops last-good Antigravity quota after a terminal 401 refresh", async () => { + let rejected = false; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (rejected && url.includes(":retrieveUserQuota") && !url.includes("Summary")) { + return jsonResponse({}, 401); + } + if (url.endsWith(":retrieveUserQuota") && !url.includes("Summary")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const valid = await fetchProviderQuotaReports(config(), true); + rejected = true; + const rejectedRefresh = await fetchProviderQuotaReports(config(), true); + + expect(valid.reports).toHaveLength(1); + expect(rejectedRefresh.reports).toEqual([]); + }); + + test("does not fetch production or catalog when daily retrieveUserQuota 401 races a 404 summary", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) { + await Bun.sleep(20); + return jsonResponse({}, 401); + } + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuotaSummary`) return jsonResponse({}, 404); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return jsonResponse({}, 404); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); + }); + + test("does not fetch production or catalog after daily retrieveUserQuota returns 429", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 429); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return jsonResponse({}, 503); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); + }); + + test("does not fetch production or catalog after daily retrieveUserQuota returns 403", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 403); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); + }); + + test("does not fail over to daily or prod for a custom baseUrl on 404/503", async () => { + const customHost = "https://custom-proxy.example.com"; + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.startsWith(customHost) && url.includes(":retrieveUserQuota")) return jsonResponse({}, 404); + if (url.startsWith(customHost) && url.includes(":retrieveUserQuotaSummary")) return jsonResponse({}, 503); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(customHost), true); + + expect(requested.filter(url => url.startsWith(DAILY_HOST) || url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.startsWith(customHost))).toEqual([ + `${customHost}/v1internal:retrieveUserQuota`, + `${customHost}/v1internal:retrieveUserQuotaSummary`, + `${customHost}/v1internal:fetchAvailableModels`, + ]); + expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); + }); + + test("does not POST retrieveUserQuota or retrieveUserQuotaSummary to an http host", async () => { + const httpHost = "http://daily-cloudcode-pa.googleapis.com"; + const requested: string[] = []; + const fetchImpl = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const quota = await fetchAntigravityLiveQuota({ + accessToken: TOKEN, + projectId: PROJECT, + baseUrl: httpHost, + timeoutMs: 8_000, + fetchImpl, + }); + + expect(requested.filter(url => url.startsWith("http://"))).toEqual([]); + expect(requested).not.toContain(`${httpHost}/v1internal:retrieveUserQuota`); + expect(requested).not.toContain(`${httpHost}/v1internal:retrieveUserQuotaSummary`); + expect(quota).toBeNull(); + }); + + test("does not POST fetchAvailableModels to an http host", async () => { + const httpHost = "http://daily-cloudcode-pa.googleapis.com"; + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + await fetchProviderQuotaReports(config(httpHost), true); + + expect(requested.filter(url => url.startsWith("http://"))).toEqual([]); + expect(requested).not.toContain(`${httpHost}/v1internal:fetchAvailableModels`); + }); +}); diff --git a/tests/google-adapter.test.ts b/tests/google-adapter.test.ts index 15af3a264a..4d65b99740 100644 --- a/tests/google-adapter.test.ts +++ b/tests/google-adapter.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test"; import { createGoogleAdapter } from "../src/adapters/google"; +import { anthropicToolCallId } from "../src/adapters/tool-call-id"; import type { OcxParsedRequest } from "../src/types"; const provider = { adapter: "google", baseUrl: "https://generativelanguage.googleapis.com", apiKey: "key" }; @@ -133,6 +134,30 @@ describe("google adapter — tool-call ids on the wire", () => { expect(frPart.functionResponse.id).toBe("call_abc"); }); + test("orphan tool results are omitted instead of emitting an unmatched functionResponse", async () => { + const contents = await geminiContents(parsedWith([ + { role: "toolResult", toolCallId: "orphan", toolName: "missing", content: "discard", isError: false }, + { role: "user", content: "continue" }, + ])); + + expect(contents.flatMap(content => content.parts).some(part => "functionResponse" in part)).toBe(false); + expect(JSON.stringify(contents)).not.toContain("orphan"); + }); + + test("orphan result ids do not reserve allocator slots", async () => { + const rawId = "call:a"; + const normalizedId = anthropicToolCallId(rawId)!; + const contents = await geminiContents(parsedWith([ + { role: "assistant", content: [{ type: "toolCall", id: rawId, name: "bash", arguments: {} }] }, + { role: "toolResult", toolCallId: rawId, toolName: "bash", content: "ok", isError: false }, + { role: "toolResult", toolCallId: normalizedId, toolName: "missing", content: "discard", isError: false }, + ])); + + const functionCall = contents.find(content => content.role === "model")!.parts + .find(part => "functionCall" in part) as { functionCall: { id?: string } }; + expect(functionCall.functionCall.id).toBe(normalizedId); + }); + test("ids are normalized to Anthropic's tool_use.id charset, preserving call/response pairing", async () => { const contents = await geminiContents(parsedWith([ { role: "assistant", content: [{ type: "toolCall", id: "fc:weird/id#1", name: "bash", arguments: {} }] }, @@ -153,6 +178,8 @@ describe("google adapter — tool-call ids on the wire", () => { { type: "toolCall", id: "call:a", name: "bash", arguments: {} }, { type: "toolCall", id: "call/a", name: "bash", arguments: {} }, ] }, + { role: "toolResult", toolCallId: "call:a", toolName: "bash", content: "one", isError: false }, + { role: "toolResult", toolCallId: "call/a", toolName: "bash", content: "two", isError: false }, ])); const ids = contents.find(c => c.role === "model")!.parts .filter(p => "functionCall" in p) diff --git a/tests/google-antigravity-errors.test.ts b/tests/google-antigravity-errors.test.ts new file mode 100644 index 0000000000..fccbbf63ad --- /dev/null +++ b/tests/google-antigravity-errors.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, test } from "bun:test"; +import { + isAntigravityGeoBlockedBody, + isQuotaExhaustedBody, + retryableGoogleStatus, + safeAntigravityHttpErrorMessage, +} from "../src/adapters/google-errors"; + +const GEO_BLOCKED_DETAIL = "User location is not supported for the API use"; + +describe("Antigravity Google error classification", () => { + test("classifies geo-blocked 403 responses and redacts echoed credentials", () => { + const leakedToken = "geo-access-token-value-123456"; + const payload = JSON.stringify({ + error: { + status: "PERMISSION_DENIED", + message: `${GEO_BLOCKED_DETAIL}; accessToken=${leakedToken}`, + }, + }); + + expect(isAntigravityGeoBlockedBody(payload)).toBe(true); + const message = safeAntigravityHttpErrorMessage(403, payload); + expect(message).toContain("Antigravity location not supported"); + expect(message).not.toContain(leakedToken); + }); + + test("keeps ordinary permission-denied 403 responses as access denied", () => { + const payload = JSON.stringify({ + error: { + status: "PERMISSION_DENIED", + message: "The caller does not have permission to use this resource", + }, + }); + + expect(isAntigravityGeoBlockedBody(payload)).toBe(false); + expect(safeAntigravityHttpErrorMessage(403, payload)).toContain("Antigravity access denied"); + }); + + test("preserves quota and rate-limit classification for 429 responses", () => { + const quotaPayload = JSON.stringify({ + error: { + status: "RESOURCE_EXHAUSTED", + message: "Quota exceeded for this project", + }, + }); + const rateLimitPayload = JSON.stringify({ + error: { + status: "RESOURCE_EXHAUSTED", + message: "Rate limit exceeded", + }, + }); + + expect(safeAntigravityHttpErrorMessage(429, quotaPayload)).toContain("Antigravity quota exhausted"); + expect(safeAntigravityHttpErrorMessage(429, rateLimitPayload)).toContain("Antigravity rate limit exceeded"); + expect(isQuotaExhaustedBody(quotaPayload)).toBe(true); + expect(isQuotaExhaustedBody(rateLimitPayload)).toBe(false); + expect(retryableGoogleStatus(403)).toBe(false); + }); +}); diff --git a/tests/google-antigravity-wire.test.ts b/tests/google-antigravity-wire.test.ts index 8aefb92d46..b04b2560d9 100644 --- a/tests/google-antigravity-wire.test.ts +++ b/tests/google-antigravity-wire.test.ts @@ -1,10 +1,12 @@ import { describe, expect, test } from "bun:test"; import { createGoogleAdapter as createGoogleAdapterProduction } from "../src/adapters/google"; import { antigravitySessionId, isLikelyRealThoughtSignature } from "../src/adapters/google-antigravity-wire"; +import { antigravityHostCandidates } from "../src/adapters/google-antigravity-hosts"; +import { repairGoogleToolPairs, stripTrailingClaudePrefill } from "../src/adapters/google-antigravity-tools"; import { ANTIGRAVITY_MODELS, ANTIGRAVITY_MODEL_EFFORTS, canonicalAntigravityUsageModel, parseAntigravityAvailableModels, registerAntigravityDiscoveredWireModels, resolveAntigravityEffortWireModel, resolveAntigravityWireModelId } from "../src/providers/antigravity-models"; import { MODEL_DISCOVERY_MAX_MODEL_ID_LENGTH, MODEL_DISCOVERY_MAX_MODELS } from "../src/providers/model-discovery"; import type { AdapterEvent, OcxParsedRequest, OcxProviderConfig } from "../src/types"; -import { withTestTranslatorBudget } from "./helpers/translator-budget"; +import { createTestTranslatorBudget, withTestTranslatorBudget } from "./helpers/translator-budget"; const createGoogleAdapter = (...args: Parameters) => withTestTranslatorBudget(createGoogleAdapterProduction(...args)); @@ -44,7 +46,7 @@ describe("antigravity CCA envelope", () => { test("wraps the gemini body in the CCA envelope with project/userAgent/requestType/requestId/sessionId", async () => { const req = await createGoogleAdapter(provider).buildRequest(parsed()); const env = JSON.parse(req.body); - expect(req.url).toBe("https://daily-cloudcode-pa.googleapis.com/v1internal:generateContent"); + expect(req.url).toBe("https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"); expect(env.model).toBe("gemini-3-pro"); // The envelope BODY userAgent is the protocol constant; the versioned CLI UA rides in the header. expect(env.userAgent).toBe("antigravity"); @@ -76,6 +78,56 @@ describe("antigravity CCA envelope", () => { expect(req.url).toBe("https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"); }); + test("host candidates keep the configured host first and use only daily/prod", () => { + expect(antigravityHostCandidates("https://daily-cloudcode-pa.googleapis.com")).toEqual([ + "https://daily-cloudcode-pa.googleapis.com", + "https://cloudcode-pa.googleapis.com", + ]); + expect(antigravityHostCandidates("https://cloudcode-pa.googleapis.com/")).toEqual([ + "https://cloudcode-pa.googleapis.com", + "https://daily-cloudcode-pa.googleapis.com", + ]); + }); + + test("Claude CCA adds the interleaved-thinking beta header and preamble mode", async () => { + const req = await createGoogleAdapter(provider).buildRequest(parsed("x", false, "claude-sonnet-4-6")); + const env = JSON.parse(req.body); + + expect(req.headers["anthropic-beta"]).toBe("interleaved-thinking-2025-05-14"); + expect(env.request.preambleConfig).toEqual({ mode: "SYSTEM_INSTRUCTION_MODE_REPLACE" }); + }); + + test("Gemini CCA does not receive the Claude beta header", async () => { + const req = await createGoogleAdapter(provider).buildRequest(parsed()); + expect(req.headers["anthropic-beta"]).toBeUndefined(); + }); + + test("Claude CCA strips a trailing prefill model turn but keeps a lone model turn", async () => { + const withPrefill = { + ...parsed("x", false, "claude-sonnet-4-6"), + context: { + messages: [ + { role: "user", content: "question" }, + { role: "assistant", content: [{ type: "text", text: "prefill" }] }, + ], + systemPrompt: [], + tools: [], + }, + } as unknown as OcxParsedRequest; + const prefillEnv = JSON.parse((await createGoogleAdapter(provider).buildRequest(withPrefill)).body); + expect(prefillEnv.request.contents.map((content: { role: string }) => content.role)).toEqual(["user", "user"]); + + const loneModel = { + ...withPrefill, + context: { + ...withPrefill.context, + messages: [{ role: "assistant", content: [{ type: "text", text: "only turn" }] }], + }, + } as unknown as OcxParsedRequest; + const loneEnv = JSON.parse((await createGoogleAdapter(provider).buildRequest(loneModel)).body); + expect(loneEnv.request.contents.map((content: { role: string }) => content.role)).toEqual(["model", "user"]); + }); + test("exposes Gemini 3.7 Flash while retired Flash ids resolve to it", async () => { // Collapsed picker: base models only. expect(ANTIGRAVITY_MODELS).toEqual([ @@ -658,6 +710,84 @@ describe("antigravity CCA envelope", () => { }); }); +describe("Google Antigravity history repair", () => { + test("drops orphan tool results and unmatched trailing calls", () => { + const messages = [ + { role: "user", content: "run tools" }, + { + role: "assistant", + content: [ + { type: "toolCall", id: "call-1", name: "one", arguments: {} }, + { type: "toolCall", id: "call-2", name: "two", arguments: {} }, + ], + }, + { role: "toolResult", toolCallId: "call-1", toolName: "one", content: "ok", isError: false }, + { role: "toolResult", toolCallId: "orphan", toolName: "missing", content: "discard", isError: false }, + ] as unknown as Parameters[0]; + + const repaired = repairGoogleToolPairs(messages); + expect(repaired).toHaveLength(3); + expect((repaired[1] as { content: { id: string }[] }).content.map(part => part.id)).toEqual(["call-1"]); + expect((repaired[2] as { toolCallId: string }).toolCallId).toBe("call-1"); + }); + + test("keeps parallel calls when every call has a later result", () => { + const messages = [ + { role: "assistant", content: [ + { type: "toolCall", id: "call-1", name: "one", arguments: {} }, + { type: "toolCall", id: "call-2", name: "two", arguments: {} }, + ] }, + { role: "toolResult", toolCallId: "call-1", toolName: "one", content: "one", isError: false }, + { role: "toolResult", toolCallId: "call-2", toolName: "two", content: "two", isError: false }, + ] as unknown as Parameters[0]; + + const repaired = repairGoogleToolPairs(messages); + expect((repaired[0] as { content: { id: string }[] }).content.map(part => part.id)).toEqual(["call-1", "call-2"]); + expect(repaired).toHaveLength(3); + }); + + test("pairs duplicate tool-call ids by occurrence", () => { + const messages = [ + { role: "assistant", content: [ + { type: "toolCall", id: "dup", name: "one", arguments: {} }, + { type: "toolCall", id: "dup", name: "one", arguments: {} }, + ] }, + { role: "toolResult", toolCallId: "dup", toolName: "one", content: "first", isError: false }, + ] as unknown as Parameters[0]; + + const repaired = repairGoogleToolPairs(messages); + expect((repaired[0] as { content: { id: string }[] }).content).toHaveLength(1); + expect(repaired).toHaveLength(2); + }); + + test("keeps only the first complete exchange when duplicate ids have matching results", () => { + const messages = [ + { role: "assistant", content: [ + { type: "toolCall", id: "dup", name: "one", arguments: { n: 1 } }, + { type: "toolCall", id: "dup", name: "one", arguments: { n: 2 } }, + ] }, + { role: "toolResult", toolCallId: "dup", toolName: "one", content: "first", isError: false }, + { role: "toolResult", toolCallId: "dup", toolName: "one", content: "second", isError: false }, + ] as unknown as Parameters[0]; + + const repaired = repairGoogleToolPairs(messages); + expect((repaired[0] as { content: { id: string; arguments: { n: number } }[] }).content).toEqual([ + { type: "toolCall", id: "dup", name: "one", arguments: { n: 1 } }, + ]); + expect(repaired).toHaveLength(2); + expect((repaired[1] as { content: string }).content).toBe("first"); + }); + + test("strips only trailing model turns when another content turn remains", () => { + const contents = [{ role: "user" }, { role: "model" }, { role: "model" }]; + expect(stripTrailingClaudePrefill(contents)).toBe(true); + expect(contents).toEqual([{ role: "user" }]); + const soloModel = [{ role: "model" }]; + expect(stripTrailingClaudePrefill(soloModel)).toBe(false); + expect(soloModel).toEqual([{ role: "model" }]); + }); +}); + function sseResponse(chunks: unknown[]): Response { const body = chunks.map(c => `data: ${JSON.stringify(c)}\n`).join("\n") + "\n"; return new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }); @@ -680,10 +810,32 @@ describe("antigravity parseStream unwraps response", () => { }); describe("antigravity parseResponse unwraps response (non-streaming)", () => { + test("buffers CCA SSE frames for unary callers", async () => { + const adapter = createGoogleAdapter(provider); + const events = await adapter.parseResponse!(sseResponse([ + { response: { candidates: [{ content: { parts: [{ text: "hello" }] } }] } }, + { response: { candidates: [{ finishReason: "STOP" }] } }, + ])); + expect(events).toContainEqual({ type: "text_delta", text: "hello" }); + expect(events.at(-1)?.type).toBe("done"); + }); + + test("unary CCA responses retain the translated event batch in the translator budget", async () => { + const adapter = createGoogleAdapter(provider); + const budget = createTestTranslatorBudget({ maxTurnBytes: 1024 }); + const events = await adapter.parseResponse!(sseResponse([ + { response: { candidates: [{ content: { parts: [{ text: "hello" }] } }] } }, + { response: { candidates: [{ finishReason: "STOP" }] } }, + ]), budget); + + expect(events).toContainEqual({ type: "text_delta", text: "hello" }); + expect(budget.snapshot().currentBytes).toBeGreaterThan(0); + }); + test("reads response.candidates + response.usageMetadata from the CCA envelope", async () => { const adapter = createGoogleAdapter(provider); - const body = JSON.stringify({ response: { candidates: [{ content: { parts: [{ text: "hello" }] } }], usageMetadata: { promptTokenCount: 9, candidatesTokenCount: 2, cachedContentTokenCount: 7 } } }); - const events = await adapter.parseResponse!(new Response(body, { status: 200 })); + const body = { response: { candidates: [{ content: { parts: [{ text: "hello" }] } }], usageMetadata: { promptTokenCount: 9, candidatesTokenCount: 2, cachedContentTokenCount: 7 } } }; + const events = await adapter.parseResponse!(sseResponse([body])); expect(events.some(e => e.type === "text_delta" && e.text === "hello")).toBe(true); const done = events.find(e => e.type === "done"); expect((done as Extract).usage?.inputTokens).toBe(9); @@ -696,8 +848,12 @@ describe("antigravity parseResponse unwraps response (non-streaming)", () => { const adapter = createGoogleAdapter(provider); // buildRequest first to set the per-adapter model/session, then parseResponse to observe. await adapter.buildRequest(parsed("hello world")); - const body = JSON.stringify({ response: { candidates: [{ content: { parts: [{ functionCall: { name: "do_x", args: { a: 1 } }, thoughtSignature: "sig-nonstream0000000" } ] } }] } }); - await adapter.parseResponse!(new Response(body, { status: 200 })); + const body = { response: { candidates: [{ content: { parts: [{ functionCall: { name: "do_x", args: { a: 1 } }, thoughtSignature: "sig-nonstream0000000" } ] } }] } }; + const events = await adapter.parseResponse!(sseResponse([ + body, + { response: { candidates: [{ finishReason: "STOP" }] } }, + ])); + expect(events.at(-1)?.type).toBe("done"); // A follow-up request's history should now get the signature re-injected. const followup = parsed("hello world"); const contents = [{ role: "model", parts: [{ functionCall: { name: "do_x", args: { a: 1 } } }] }]; @@ -716,7 +872,7 @@ describe("antigravity parseResponse unwraps response (non-streaming)", () => { __resetAntigravityReplayCache(); const adapter = createGoogleAdapter(provider); await adapter.buildRequest(parsed("hello world")); - const body = JSON.stringify({ + const body = { response: { candidates: [{ content: { @@ -727,8 +883,8 @@ describe("antigravity parseResponse unwraps response (non-streaming)", () => { }, }], }, - }); - const events = await adapter.parseResponse!(new Response(body, { status: 200 })); + }; + const events = await adapter.parseResponse!(sseResponse([body])); expect(events).not.toContainEqual({ type: "text_delta", text: "deciding which tool to call" }); @@ -748,6 +904,7 @@ describe("antigravity history preserves tool-call thoughtSignature", () => { messages: [ { role: "user", content: "go" }, { role: "assistant", content: [{ type: "toolCall", id: "c1", name: "get_x", namespace: "mcp__t", arguments: { a: 1 }, thoughtSignature: "sig-abcdef0123456789" }] }, + { role: "toolResult", toolCallId: "c1", toolName: "get_x", content: "ok", isError: false }, ], systemPrompt: [], tools: [], }, @@ -768,6 +925,7 @@ describe("antigravity history preserves tool-call thoughtSignature", () => { messages: [ { role: "user", content: "go" }, { role: "assistant", content: [{ type: "toolCall", id: "c1", name: "get_x", namespace: "mcp__t", arguments: {}, thoughtSignature: "fc_d8df7548e31a4130b7624f3d27571cdd" }] }, + { role: "toolResult", toolCallId: "c1", toolName: "get_x", content: "ok", isError: false }, ], systemPrompt: [], tools: [], }, @@ -788,6 +946,7 @@ describe("antigravity history preserves tool-call thoughtSignature", () => { messages: [ { role: "user", content: "go" }, { role: "assistant", content: [{ type: "toolCall", id: "c1", name: "get_x", namespace: "mcp__t", arguments: {}, thoughtSignature: "ctc_038f26d3f20962bc016a54f0fcfa208190a8ec0f289c2ba211" }] }, + { role: "toolResult", toolCallId: "c1", toolName: "get_x", content: "ok", isError: false }, ], systemPrompt: [], tools: [], }, diff --git a/tests/google-hardening.test.ts b/tests/google-hardening.test.ts index 207dddd8bc..fcca2ca77d 100644 --- a/tests/google-hardening.test.ts +++ b/tests/google-hardening.test.ts @@ -84,21 +84,36 @@ describe("google provider hardening", () => { ); }); + test("Antigravity rejects a cleartext http baseUrl before dispatch", async () => { + const adapter = createGoogleAdapter(antigravityProvider({ + baseUrl: "http://daily-cloudcode-pa.googleapis.com", + })); + + await expect(adapter.buildRequest(parsed())).rejects.toThrow( + "google-antigravity requires an HTTPS baseUrl", + ); + }); + test("Antigravity rejects flat Gemini payloads without the response wrapper", async () => { const adapter = createGoogleAdapter(antigravityProvider()); const flatPayload = { candidates: [{ content: { parts: [{ text: "unexpected" }] } }] }; + // SSE-framed flat payload: parseStream unwraps each data frame and rejects missing `response`. const streamEvents = await collect(adapter.parseStream(sseResponse([flatPayload]))); + expect(streamEvents).toEqual([{ + type: "error", + message: "google-antigravity response missing response wrapper", + }]); + + // Plain JSON without SSE framing: CCA parseResponse delegates to parseStream, which reads until + // EOF without finding a `data:` frame and fails closed as truncated SSE transport. const responseEvents = await adapter.parseResponse!( new Response(JSON.stringify(flatPayload), { status: 200 }), ); - - const expected = [{ + expect(responseEvents).toEqual([{ type: "error", - message: "google-antigravity response missing response wrapper", - }]; - expect(streamEvents).toEqual(expected); - expect(responseEvents).toEqual(expected); + message: "upstream stream ended with an incomplete SSE frame — possible truncation", + }]); }); test("truncated final JSON is a terminal stream error", async () => { diff --git a/tests/google-sse-frame-cap.test.ts b/tests/google-sse-frame-cap.test.ts new file mode 100644 index 0000000000..d06d9b55f8 --- /dev/null +++ b/tests/google-sse-frame-cap.test.ts @@ -0,0 +1,204 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { + createGoogleAdapter as createGoogleAdapterProduction, + setGoogleSseFrameMaxBytesForTests, +} from "../src/adapters/google"; +import type { AdapterEvent, OcxProviderConfig } from "../src/types"; +import { withTestTranslatorBudget } from "./helpers/translator-budget"; + +const createGoogleAdapter = (...args: Parameters) => + withTestTranslatorBudget(createGoogleAdapterProduction(...args)); + +const CAP = 32; +const originalDecode = TextDecoder.prototype.decode; +let decodedOverflowByteLength = 0; + +function installDecodeProbe(): void { + decodedOverflowByteLength = 0; + TextDecoder.prototype.decode = function ( + this: TextDecoder, + input?: AllowSharedBufferSource, + options?: TextDecodeOptions, + ): string { + const size = input && typeof (input as ArrayBufferView).byteLength === "number" + ? (input as ArrayBufferView).byteLength + : 0; + if (size > CAP) decodedOverflowByteLength = size; + return originalDecode.call(this, input as ArrayBuffer, options); + }; +} + +afterEach(() => { + setGoogleSseFrameMaxBytesForTests(); + TextDecoder.prototype.decode = originalDecode; + decodedOverflowByteLength = 0; +}); + +function googleProvider(): OcxProviderConfig { + return { + adapter: "google", + baseUrl: "https://generativelanguage.googleapis.com", + apiKey: "google-test-key", + authMode: "key", + }; +} + +function ccaProvider(): OcxProviderConfig { + return { + adapter: "google", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + apiKey: "antigravity-test-token", + authMode: "oauth", + googleMode: "cloud-code-assist", + project: "project-test", + }; +} + +/** 20 × U+4E2D: 60 UTF-8 bytes, 20 UTF-16 units — over a 32-byte cap, under it as string length. */ +function oversizedMultibyteChunk(): Uint8Array { + const charUtf8 = new TextEncoder().encode("中"); + const repeats = 20; + const chunk = new Uint8Array(charUtf8.byteLength * repeats); + for (let i = 0; i < repeats; i++) chunk.set(charUtf8, i * charUtf8.byteLength); + return chunk; +} + +function byteStreamResponse(chunks: Uint8Array[]): Response { + return new Response(new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(chunk); + controller.close(); + }, + }), { status: 200, headers: { "content-type": "text/event-stream" } }); +} + +async function collect(events: AsyncGenerator): Promise { + const collected: AdapterEvent[] = []; + for await (const event of events) collected.push(event); + return collected; +} + +describe("google SSE frame byte cap", () => { + test("rejects an oversize UTF-8 chunk before TextDecoder.decode", async () => { + const chunk = oversizedMultibyteChunk(); + expect(chunk.byteLength).toBeGreaterThan(CAP); + expect(new TextDecoder().decode(chunk).length).toBeLessThan(CAP); + + setGoogleSseFrameMaxBytesForTests(CAP); + installDecodeProbe(); + + const events = await collect( + createGoogleAdapter(googleProvider()).parseStream(byteStreamResponse([chunk])), + ); + + expect(decodedOverflowByteLength).toBe(0); + expect(events).toContainEqual({ + type: "error", + message: `upstream SSE data frame exceeds ${CAP} bytes`, + }); + }); + + test("CCA unary parseResponse applies the same SSE byte cap", async () => { + const chunk = oversizedMultibyteChunk(); + setGoogleSseFrameMaxBytesForTests(CAP); + installDecodeProbe(); + + const events = await createGoogleAdapter(ccaProvider()).parseResponse!( + byteStreamResponse([chunk]), + ); + + expect(decodedOverflowByteLength).toBe(0); + expect(events).toContainEqual({ + type: "error", + message: `upstream SSE data frame exceeds ${CAP} bytes`, + }); + }); + + test("rejects a split oversized multibyte line before decoding the completing chunk", async () => { + const zhong = new TextEncoder().encode("中"); + expect(zhong.byteLength).toBe(3); + const prefix = new TextEncoder().encode("data: "); + const fill = new Uint8Array(CAP + 1 - prefix.byteLength - zhong.byteLength).fill(0x61); + const first = new Uint8Array(prefix.byteLength + fill.byteLength + 1); + first.set(prefix, 0); + first.set(fill, prefix.byteLength); + first.set(zhong.subarray(0, 1), prefix.byteLength + fill.byteLength); + const completing = zhong.subarray(1); + expect(first.byteLength).toBe(CAP - 1); + expect(first.byteLength + completing.byteLength).toBe(CAP + 1); + + setGoogleSseFrameMaxBytesForTests(CAP); + let decodedCompletingChunk = false; + TextDecoder.prototype.decode = function ( + this: TextDecoder, + input?: AllowSharedBufferSource, + options?: TextDecodeOptions, + ): string { + const view = input instanceof Uint8Array + ? input + : input && typeof (input as ArrayBufferView).byteLength === "number" + ? new Uint8Array(input as ArrayBufferView) + : null; + if ( + view + && view.byteLength === completing.byteLength + && completing.every((byte, index) => view[index] === byte) + ) { + decodedCompletingChunk = true; + } + return originalDecode.call(this, input as ArrayBuffer, options); + }; + + const events = await collect( + createGoogleAdapter(googleProvider()).parseStream(byteStreamResponse([first, completing])), + ); + + expect(decodedCompletingChunk).toBe(false); + expect(events).toContainEqual({ + type: "error", + message: `upstream SSE data frame exceeds ${CAP} bytes`, + }); + }); + + test("accepts a data line exactly at the cap before its LF delimiter", async () => { + const cap = 128; + const envelope = (text: string) => ({ + response: { candidates: [{ content: { parts: [{ text }] } }] }, + }); + const encoder = new TextEncoder(); + const emptyLine = `data: ${JSON.stringify(envelope(""))}`; + const line = `data: ${JSON.stringify(envelope("a".repeat(cap - encoder.encode(emptyLine).byteLength)))}`; + expect(encoder.encode(line).byteLength).toBe(cap); + + setGoogleSseFrameMaxBytesForTests(cap); + const events = await collect( + createGoogleAdapter(googleProvider()).parseStream( + byteStreamResponse([encoder.encode(`${line}\n\n`)]), + ), + ); + + expect(events.some(event => event.type === "error" && event.message.includes("exceeds"))).toBe(false); + }); + + test("accepts multiple sub-cap data frames delivered in one oversized chunk", async () => { + const cap = 96; + const body = [ + `data: ${JSON.stringify({ response: { candidates: [{ content: { parts: [{ text: "a" }] } }] } })}\n`, + `data: ${JSON.stringify({ response: { candidates: [{ finishReason: "STOP" }] } })}\n`, + ].join("\n"); + const lines = body.split("\n").filter(Boolean); + expect(new TextEncoder().encode(body).byteLength).toBeGreaterThan(cap); + expect(Math.max(...lines.map(line => new TextEncoder().encode(line).byteLength))).toBeLessThanOrEqual(cap); + + setGoogleSseFrameMaxBytesForTests(cap); + const events = await collect( + createGoogleAdapter(ccaProvider()).parseStream(byteStreamResponse([ + new TextEncoder().encode(body), + ])), + ); + + expect(events).toContainEqual({ type: "text_delta", text: "a" }); + expect(events).toContainEqual({ type: "done", usage: undefined }); + expect(events.some(event => event.type === "error")).toBe(false); + }); +}); diff --git a/tests/provider-quota.test.ts b/tests/provider-quota.test.ts index 1c4e733683..db1bad354a 100644 --- a/tests/provider-quota.test.ts +++ b/tests/provider-quota.test.ts @@ -121,11 +121,16 @@ describe("fetchProviderQuotaReports", () => { await saveCredential("google-antigravity", { access: "agy-access-secret", refresh: "agy-refresh-secret", expires: Date.now() + 3600_000, projectId: "agy-project-secret" }); await saveCredential("kimi", { access: "kimi-access-secret", refresh: "kimi-refresh-secret", expires: Date.now() + 3600_000 }); - const seen: { url: string; authorization?: string; body?: string }[] = []; + const seen: { url: string; authorization?: string; body?: string; redirect?: RequestRedirect }[] = []; globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { const url = String(input); const headers = init?.headers as Record | undefined; - seen.push({ url, authorization: headers?.Authorization, body: typeof init?.body === "string" ? init.body : undefined }); + seen.push({ + url, + authorization: headers?.Authorization, + body: typeof init?.body === "string" ? init.body : undefined, + redirect: init?.redirect, + }); if (url === "https://chatgpt.com/backend-api/wham/usage") { return new Response(JSON.stringify({ email: "person@example.com", @@ -255,9 +260,48 @@ describe("fetchProviderQuotaReports", () => { expect(seen.find(row => row.url.includes("anthropic.com"))?.authorization).toBe("Bearer claude-access-secret"); expect(seen.find(row => row.url.includes("cloudcode-pa.googleapis.com"))?.authorization).toBe("Bearer agy-access-secret"); expect(seen.find(row => row.url.includes("cloudcode-pa.googleapis.com"))?.body).toBe(JSON.stringify({ project: "agy-project-secret" })); + expect(seen.find(row => row.url.endsWith("/v1internal:fetchAvailableModels"))?.redirect).toBe("error"); expect(seen.find(row => row.url === "https://api.kimi.com/coding/v1/usages")?.authorization).toBe("Bearer kimi-access-secret"); }); + test("treats remainingPercentage as a percentage at values one and below", async () => { + await saveCredential("google-antigravity", { + access: "agy-access-secret", + refresh: "agy-refresh-secret", + expires: Date.now() + 3600_000, + projectId: "agy-project-secret", + }); + const config = { + defaultProvider: "google-antigravity", + providers: { + "google-antigravity": { + adapter: "google", + authMode: "oauth", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + }, + }, + } as OcxConfig; + + for (const [remainingPercentage, expectedUsedPercentage] of [[1, 99], [0.75, 99.25], [75, 25]]) { + clearProviderQuotaCache(); + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes(":retrieveUserQuota")) return new Response("not found", { status: 404 }); + if (url.endsWith("/v1internal:fetchAvailableModels")) { + return new Response(JSON.stringify({ + models: { + "gemini-test": { quotaInfo: { remainingPercentage } }, + }, + }), { status: 200 }); + } + return new Response("not found", { status: 404 }); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config, true); + expect(result.reports[0]?.quota.customWindows?.[0]?.percent).toBe(expectedUsedPercentage); + } + }); + function kimiOnlyConfig(baseUrl = "https://api.kimi.com/coding/v1"): OcxConfig { return { defaultProvider: "kimi",