diff --git a/.env.example b/.env.example index f10df5cff4..fdb38b67f7 100644 --- a/.env.example +++ b/.env.example @@ -222,7 +222,7 @@ CONTAINER_HOST=docker # - orbstack: OrbStack (high-perf Linux VM + docker shim on macOS) # - podman: Podman (rootless, daemonless) # - docker: Docker (default fallback) -SKILLS_SANDBOX_RUNTIME=auto +# (defined under SKILLS & SANDBOXING section below) # ═══════════════════════════════════════════════════════════════════════════════ # 4. SECURITY & AUTHENTICATION diff --git a/.gitignore b/.gitignore index b7406dbfca..a507938f26 100644 --- a/.gitignore +++ b/.gitignore @@ -233,7 +233,7 @@ omniroute.md # mise configuration mise.toml -_artifacts/ +_artifacts/ # release-green artifacts .claude-flow/ # ESLint file cache (npm run lint --cache / complexity ratchets) @@ -241,7 +241,7 @@ _artifacts/ .eslintcache-complexity -# CI/local quality artifacts (eslint-results.json, etc.) +# CI/local quality artifacts (eslint-results.json, quality-ratchet.md, etc.) .artifacts/ # Homologation E2E suite (npm run homolog) — real-environment credentials + report output diff --git a/config/quality/eslint-suppressions.json b/config/quality/eslint-suppressions.json index 6f94816861..e0c1a0a73d 100644 --- a/config/quality/eslint-suppressions.json +++ b/config/quality/eslint-suppressions.json @@ -2047,11 +2047,6 @@ "count": 13 } }, - "tests/unit/sse-parser.test.ts": { - "@typescript-eslint/no-explicit-any": { - "count": 15 - } - }, "tests/unit/startup-stale-cooldown-recovery.test.ts": { "@typescript-eslint/no-explicit-any": { "count": 1 diff --git a/open-sse/executors/antigravity.ts b/open-sse/executors/antigravity.ts index e9f13e5bb5..e56ea92164 100644 --- a/open-sse/executors/antigravity.ts +++ b/open-sse/executors/antigravity.ts @@ -1,20 +1,16 @@ import crypto, { randomUUID } from "crypto"; import { BaseExecutor, - mergeAbortSignals, mergeUpstreamExtraHeaders, type ExecuteInput, type ExecutorLog, type ProviderCredentials, } from "./base.ts"; -import { applyFingerprint, isCliCompatEnabled } from "../config/cliFingerprints.ts"; -import { buildAntigravityUpstreamError } from "./antigravityUpstreamError.ts"; import { PROVIDERS, OAUTH_ENDPOINTS, HTTP_STATUS, - STREAM_READINESS_TIMEOUT_MS, - ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, + FETCH_TIMEOUT_MS, } from "../config/constants.ts"; import { scrubProxyAndFingerprintHeaders } from "../services/antigravityHeaderScrub.ts"; import { @@ -23,7 +19,6 @@ import { } from "../services/antigravityHeaders.ts"; import { classify429, decide429, type Decision } from "../services/antigravity429Engine.ts"; import { - injectCreditsField, shouldRetryWithCredits, shouldUseCreditsFirst, getCreditsMode, @@ -39,7 +34,6 @@ import { resolveAntigravityModelId, getAntigravityModelFallbacks, } from "../config/antigravityModelAliases.ts"; -import { cloakAntigravityToolPayload } from "../config/toolCloaking.ts"; import { shouldStripCloudCodeThinking, stripCloudCodeThinkingConfig, @@ -53,27 +47,39 @@ import { } from "./antigravity/sseCollect.ts"; // processAntigravitySSEPayload re-exported for external importers (tests). export { processAntigravitySSEPayload } from "./antigravity/sseCollect.ts"; +import { + createCreditsExtractionTransform as createCreditsExtractionTransformImpl, + type SsePassthroughResult, +} from "./antigravity/streamingPassthrough.ts"; +import { + toSafeAntigravityLog, + finalizeAntigravityRequestBody, + sendAntigravityRequest, + tryCreditsRetry, + tryEmbedLongRetryAfter, + buildFinalAntigravityResult, + buildAntigravity429ErrorMessage, + markCreditsExhausted, + type SafeAntigravityLog, +} from "./antigravity/executeAttempt.ts"; import { handleAntigravityFallbackChainError, handleAntigravityFallback400, } from "./antigravity/proFallbackChain.ts"; -import { - applyAntigravityClientProfileHeaders, - removeHeaderCaseInsensitive, -} from "../services/antigravityClientProfile.ts"; import { generateAntigravityRequestId, getAntigravityEnvelopeUserAgent, getAntigravitySessionId, } from "../services/antigravityIdentity.ts"; -import * as prl from "../utils/providerRequestLogging.ts"; const MAX_RETRY_AFTER_MS = 60_000; const LONG_RETRY_THRESHOLD_MS = 60_000; -const CREDITS_EXHAUSTED_TTL_MS = 5 * 60 * 60 * 1000; // 5 hours // Cap for transient 5xx backoff — shorter than the 429 cap to avoid long stalls on // infra hiccups ("Agent execution terminated", "high traffic", capacity errors). const ANTIGRAVITY_TRANSIENT_RETRY_MAX_MS = 15_000; +// Bounded per-URL auto-retry count for both the Retry-After-driven short retry and +// the no-Retry-After transient/429 backoff loop in executeOnce(). +const MAX_AUTO_RETRIES = 3; const ANTIGRAVITY_TRANSIENT_ERROR_PATTERNS: RegExp[] = [ /high\s+traffic/i, @@ -105,7 +111,7 @@ interface AntigravityContent { [key: string]: unknown; } -type AntigravityCredentials = ProviderCredentials & { +export type AntigravityCredentials = ProviderCredentials & { projectId?: string | null; expiresIn?: number; }; @@ -123,51 +129,6 @@ type AntigravityChunkContent = Record & { >; }; -type AntigravityCreditEntry = { - creditType?: string; - creditAmount?: string; -}; - -function getChunkedOrFixedBody(bodyStr: string, stream: boolean): BodyInit { - if (stream) { - return new ReadableStream( - { - async start(controller) { - controller.enqueue(new TextEncoder().encode(bodyStr)); - controller.close(); - }, - }, - { highWaterMark: 16384 } - ); - } - return bodyStr; -} - -function cloneAntigravityRequestBody(body: unknown): unknown { - if (!body || typeof body !== "object") { - return body; - } - - try { - return structuredClone(body); - } catch { - return JSON.parse(JSON.stringify(body)); - } -} - -function serializeAntigravityRequest( - provider: string, - headers: Record, - body: unknown -): { headers: Record; bodyString: string } { - const serializedBody = cloneAntigravityRequestBody(body); - - if (!isCliCompatEnabled(provider)) { - return { headers, bodyString: JSON.stringify(serializedBody) }; - } - return applyFingerprint(provider, { ...headers }, serializedBody); -} - type AntigravityRequestEnvelope = Record & { project: string; model?: string; @@ -178,44 +139,6 @@ type AntigravityRequestEnvelope = Record & { enabledCreditTypes?: string[]; }; -class AntigravityPreResponseTimeoutError extends Error { - code = ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; - status = HTTP_STATUS.GATEWAY_TIMEOUT; - - constructor(timeoutMs: number, url: string) { - super(`Antigravity upstream did not return response headers within ${timeoutMs}ms: ${url}`); - this.name = "TimeoutError"; - } -} - -function getAbortErrorCode(error: unknown): string | null { - if (!error || typeof error !== "object") return null; - const value = (error as { code?: unknown }).code; - return typeof value === "string" ? value : null; -} - -function isAntigravityPreResponseTimeout(error: unknown): boolean { - return getAbortErrorCode(error) === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; -} - -/** - * Per-account GOOGLE_ONE_AI credits-exhausted tracker. - * Key: accountId (OAuth subject / email). Value: expiry timestamp. - * When credits hit 0 we skip the credit retry for CREDITS_EXHAUSTED_TTL_MS. - */ -const MAX_CREDITS_EXHAUSTED_ENTRIES = 50; -const creditsExhaustedUntil = new Map(); - -const _creditsExhaustedSweep = setInterval(() => { - const now = Date.now(); - for (const [key, until] of creditsExhaustedUntil) { - if (now >= until) creditsExhaustedUntil.delete(key); - } -}, 60_000); -if (typeof _creditsExhaustedSweep === "object" && "unref" in _creditsExhaustedSweep) { - (_creditsExhaustedSweep as { unref?: () => void }).unref?.(); -} - const MAX_CREDIT_BALANCE_ENTRIES = 50; const CREDIT_BALANCE_TTL_MS = 5 * 60 * 1000; const creditBalanceCache = new Map(); @@ -275,33 +198,24 @@ export function updateAntigravityRemainingCredits(accountId: string, balance: nu } catch {} } -function isCreditsExhausted(accountId: string): boolean { - const until = creditsExhaustedUntil.get(accountId); - if (!until) return false; - if (Date.now() >= until) { - creditsExhaustedUntil.delete(accountId); - return false; - } - return true; -} - -function markCreditsExhausted(accountId: string): void { - if ( - creditsExhaustedUntil.size >= MAX_CREDITS_EXHAUSTED_ENTRIES && - !creditsExhaustedUntil.has(accountId) - ) { - const now = Date.now(); - for (const [key, until] of creditsExhaustedUntil) { - if (now >= until) { - creditsExhaustedUntil.delete(key); - } - } - if (creditsExhaustedUntil.size >= MAX_CREDITS_EXHAUSTED_ENTRIES) { - const oldestKey = creditsExhaustedUntil.keys().next().value; - if (oldestKey !== undefined) creditsExhaustedUntil.delete(oldestKey); - } - } - creditsExhaustedUntil.set(accountId, Date.now() + CREDITS_EXHAUSTED_TTL_MS); +/** + * Pass-through TransformStream that extracts `remainingCredits` from SSE + * data without consuming the stream (the downstream client receives the + * unmodified bytes). Thin wrapper around the pure implementation in + * streamingPassthrough.ts, injecting this executor's credit-balance cache + * writer so the two modules don't import each other. See that module's + * doc comment for the full parameter behavior. + * @internal Exported for unit testing only. + */ +export function createCreditsExtractionTransform( + accountId: string, + bufferSize = 0 +): TransformStream { + return createCreditsExtractionTransformImpl( + accountId, + updateAntigravityRemainingCredits, + bufferSize + ); } /** @@ -366,26 +280,6 @@ async function cleanModelName(model: string, modelIdOverride?: string): Promise< return clean; } -function attachToolNameMap(payload: T, toolNameMap: Map | null): T { - if (!toolNameMap?.size || !payload || typeof payload !== "object") { - return payload; - } - - const copy = Array.isArray(payload) ? ([...payload] as T) : ({ ...(payload as object) } as T); - Object.defineProperty(copy, "_toolNameMap", { - value: toolNameMap, - enumerable: false, - configurable: true, - writable: true, - }); - return copy; -} - -function getRequestTargetModel(body: Record): string { - const target = body.model; - return typeof target === "string" && target.length > 0 ? target : "unknown"; -} - /** * Hard ceiling on `generationConfig.maxOutputTokens` for Antigravity Cloud Code. * @@ -546,6 +440,49 @@ function stripTrailingAntigravityAssistantTurn( // Test-only export so the unit suite can exercise the strip logic directly. export const __test_stripTrailingAntigravityAssistantTurn = stripTrailingAntigravityAssistantTurn; +/** Base per-url-index attempt context, before the request has been sent. */ +type AntigravityAttemptContext = { + url: string; + model: string; + /** Pre-serialization headers (built by buildHeaders + mergeUpstreamExtraHeaders) — the + * credits-retry re-serializes from these, NOT from `finalHeaders` (already fingerprinted). */ + headers: Record; + transformedBody: Record; + requestToolNameMap: Map | null; + credentials: AntigravityCredentials; + stream: boolean; + signal: AbortSignal | null | undefined; + log: SafeAntigravityLog; + accountId: string; + creditsMode: ReturnType; + urlIndex: number; + retryAttemptsByUrl: Record; + fallbackCount: number; +}; + +/** Context threaded through the 429/503 handling helpers — adds the sent response. */ +type AntigravityRateLimitContext = AntigravityAttemptContext & { + response: Response; + finalHeaders: Record; +}; + +/** + * Outcome of handling a 429/503 response — tells executeOnce()'s loop what to do next. + * `lastStatus` mirrors the original inline code, which only updated the outer + * `lastStatus` variable when NOT retrying the same url (i.e. on retryNextUrl/fallthrough, + * never on the bounded-short-retry or transient-auto-retry same-url paths). + */ +type AntigravityRateLimitOutcome = + | { action: "return"; result: SsePassthroughResult } + | { action: "retrySameUrl" } + | { action: "retryNextUrl"; lastStatus: number } + | { action: "fallthrough"; retryMs: number | null; lastStatus: number }; + +/** Outcome of one full per-url attempt in executeOnce() — return a result, or retry. */ +type AntigravityAttemptOutcome = + | { action: "return"; result: SsePassthroughResult } + | { action: "retry"; sameUrl: boolean; lastStatus?: number }; + export class AntigravityExecutor extends BaseExecutor { constructor() { super("antigravity", PROVIDERS.antigravity); @@ -945,6 +882,10 @@ export class AntigravityExecutor extends BaseExecutor { * Collect an SSE streaming response into a single non-streaming JSON response. * Parses Gemini-format SSE chunks and assembles text content + usage into one * OpenAI-format chat.completion payload. + * + * @deprecated Use the non-streaming SSE path in chatCore instead, which calls + * parseSSEToGeminiResponse() from sseParser/geminiResponse.ts. This method is + * retained only for backward compatibility and may be removed in a future release. */ collectStreamToResponse( response: Response, @@ -963,7 +904,11 @@ export class AntigravityExecutor extends BaseExecutor { const decoder = new TextDecoder(); const logger = log || undefined; - const SSE_COLLECT_TIMEOUT_MS = 120_000; + // Guard against indefinite hangs when the upstream sends headers but + // stalls on the body. Inherit the global FETCH_TIMEOUT_MS (default 600 s, + // overridable via env) so reasoning-heavy models (gemini-3.1-pro-high on + // large prompts) are not killed by a hardcoded 120 s ceiling. + const SSE_COLLECT_TIMEOUT_MS = FETCH_TIMEOUT_MS; const collect = async () => { const collected: AntigravityCollectedStream = { @@ -1145,9 +1090,9 @@ export class AntigravityExecutor extends BaseExecutor { ) { await resolveAntigravityVersion(); const fallbackCount = this.getFallbackCount(); + const l = toSafeAntigravityLog(log); let lastError = null; let lastStatus = 0; - const MAX_AUTO_RETRIES = 3; const retryAttemptsByUrl: Record = {}; // Track retry attempts per URL // Always stream upstream — buildUrl always returns the streaming endpoint. @@ -1167,44 +1112,6 @@ export class AntigravityExecutor extends BaseExecutor { const creditsMode = getCreditsMode(); const useCreditsFirst = shouldUseCreditsFirst(credentials?.accessToken || "", creditsMode); - const fetchWithReadinessTimeout = async ( - url: string, - init: RequestInit, - timeoutMs = STREAM_READINESS_TIMEOUT_MS - ): Promise => { - const boundedTimeoutMs = Math.max(0, Math.floor(timeoutMs)); - if (boundedTimeoutMs <= 0) { - return fetch(url, init); - } - - const timeoutController = new AbortController(); - let timeoutId: ReturnType | null = setTimeout(() => { - timeoutController.abort(new AntigravityPreResponseTimeoutError(boundedTimeoutMs, url)); - }, boundedTimeoutMs); - - const existingSignal = init.signal instanceof AbortSignal ? init.signal : null; - const combinedSignal = existingSignal - ? mergeAbortSignals(existingSignal, timeoutController.signal) - : timeoutController.signal; - - try { - return await fetch(url, { ...init, signal: combinedSignal }); - } catch (error) { - if ( - timeoutController.signal.aborted && - isAntigravityPreResponseTimeout(timeoutController.signal.reason) - ) { - throw timeoutController.signal.reason; - } - throw error; - } finally { - if (timeoutId) { - clearTimeout(timeoutId); - timeoutId = null; - } - } - }; - for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) { const url = this.buildUrl(model, upstreamStream, urlIndex); const headers = this.buildHeaders(credentials, upstreamStream); @@ -1216,27 +1123,16 @@ export class AntigravityExecutor extends BaseExecutor { credentials, modelIdOverride ); - let requestToolNameMap: Map | null = null; if (transformed instanceof Response) { return { response: transformed, url, headers, transformedBody: body }; } - let transformedBody: Record = transformed; - - if (transformedBody && typeof transformedBody === "object") { - const cloaked = cloakAntigravityToolPayload(transformedBody); - transformedBody = cloaked.body; - requestToolNameMap = cloaked.toolNameMap; - } - - // Credits-first: inject GOOGLE_ONE_AI upfront so we never try the normal - // quota path. If credits are exhausted / disabled shouldUseCreditsFirst() - // returns false and we fall back to the legacy retry-on-429 flow. - if (useCreditsFirst) { - transformedBody = injectCreditsField(transformedBody); - log?.debug?.("AG_CREDITS", "Credits-first enabled (ANTIGRAVITY_CREDITS=always)"); - } + const { transformedBody, requestToolNameMap } = finalizeAntigravityRequestBody( + transformed, + useCreditsFirst, + l + ); // Initialize retry counter for this URL if (!retryAttemptsByUrl[urlIndex]) { @@ -1244,535 +1140,35 @@ export class AntigravityExecutor extends BaseExecutor { } try { - const serializedRequest = serializeAntigravityRequest( - this.provider, + const outcome = await this.runAntigravityAttempt({ + url, + model, headers, - transformedBody - ); - let finalHeaders = serializedRequest.headers; - const capture = (h: Record, s: string) => - prl.captureCurrentProviderBody(url, h, s, log); - const clientProfile = applyAntigravityClientProfileHeaders( - finalHeaders, + transformedBody, + requestToolNameMap, credentials, - transformedBody - ); - - log?.debug?.( - "TELEMETRY", - `[Antigravity] Execute - URL: ${url}, Model: ${model}, Target: ${getRequestTargetModel(transformedBody)}, RetryAttempt: ${retryAttemptsByUrl[urlIndex]}` - ); - - // Dump outgoing headers (mask Authorization) and envelope shape for debugging - if (log?.debug) { - const safeHeaders = { ...finalHeaders }; - if (safeHeaders["Authorization"]) safeHeaders["Authorization"] = "Bearer ***"; - log.debug("AG_REQUEST_HEADERS", JSON.stringify(safeHeaders)); - - const envelope = transformedBody as Record; - const requestInner = envelope.request as Record | undefined; - log.debug( - "AG_REQUEST_ENVELOPE", - JSON.stringify({ - fieldOrder: Object.keys(envelope), - project: envelope.project, - requestId: envelope.requestId, - model: envelope.model, - userAgent: envelope.userAgent, - requestType: envelope.requestType, - enabledCreditTypes: envelope.enabledCreditTypes, - clientProfile, - sessionId: requestInner?.sessionId, - generationConfig: requestInner?.generationConfig, - }) - ); - } - - await capture(finalHeaders, serializedRequest.bodyString); - let response = await fetchWithReadinessTimeout(url, { - method: "POST", - headers: finalHeaders, - body: getChunkedOrFixedBody(serializedRequest.bodyString, stream), - ...(stream ? { duplex: "half" } : {}), + stream, signal, + log: l, + accountId, + creditsMode, + urlIndex, + retryAttemptsByUrl, + fallbackCount, }); - if (response.status === HTTP_STATUS.FORBIDDEN && finalHeaders["x-goog-user-project"]) { - const retryHeaders = { ...finalHeaders }; - removeHeaderCaseInsensitive(retryHeaders, "x-goog-user-project"); - log?.debug?.("RETRY", "403 with x-goog-user-project, retrying once without it"); - await capture(retryHeaders, serializedRequest.bodyString); - response = await fetchWithReadinessTimeout(url, { - method: "POST", - headers: retryHeaders, - body: getChunkedOrFixedBody(serializedRequest.bodyString, stream), - ...(stream ? { duplex: "half" } : {}), - signal, - }); - finalHeaders = retryHeaders; - } - - if (!response.ok) { - log?.warn?.( - "TELEMETRY", - `[Antigravity] Error Response - URL: ${url}, Status: ${response.status}, Model: ${model}` - ); - } - - // Parse retry time for 429/503 responses - let retryMs: number | null = null; - - if ( - response.status === HTTP_STATUS.RATE_LIMITED || - response.status === HTTP_STATUS.SERVICE_UNAVAILABLE - ) { - // Try to get retry time from headers first - retryMs = this.parseRetryHeaders(response.headers); - - // If no retry time in headers, try to parse from error message body - if (!retryMs) { - try { - const errorBody = await response.clone().text(); - const errorJson = JSON.parse(errorBody); - let errorMessage = errorJson?.error?.message || errorJson?.message || ""; - if (errorJson?.error?.details && Array.isArray(errorJson.error.details)) { - for (const detail of errorJson.error.details) { - if (detail?.reason) { - errorMessage += ` ${detail.reason}`; - } - } - } - - // 1. Try to parse explicit retry time from message - const parsedRetryMs = this.parseRetryFromErrorMessage(errorMessage); - - // 2. Classify 429 (pass header-parsed retry hint as fallback - // signal — multi-hour Retry-After upgrades rate_limited to - // quota_exhausted so the GOOGLE_ONE_AI credits retry fires). - const effectiveRetryHintMs = retryMs ?? parsedRetryMs ?? null; - const category = classify429(errorMessage); - - // 3. Decide final retry time BEFORE the credits retry so that - // full_quota_exhausted can skip the credits attempt entirely - // (avoids ~41s hold on an already-exhausted account) and - // persist the cooldown to DB for post-restart routing. - const decision: Decision = decide429(category, parsedRetryMs); - retryMs = decision.retryAfterMs; - log?.debug?.( - "AG_429", - `Category: ${category}, Decision: ${decision.kind} — ${decision.reason}` - ); - - if (decision.kind === "full_quota_exhausted" && retryMs) { - markConnectionQuotaExhausted(accountId, retryMs); - } - - const creditsAlreadyInjected = - (transformedBody as { enabledCreditTypes?: unknown }).enabledCreditTypes != null; - - if (category === "quota_exhausted" && creditsAlreadyInjected) { - handleCreditsFailure(credentials?.accessToken || ""); - log?.warn?.("AG_CREDITS", "Credits-first request 429'd — credits likely exhausted"); - markCreditsExhausted(accountId); - } - - if ( - category === "quota_exhausted" && - decision.kind !== "full_quota_exhausted" && - !creditsAlreadyInjected && - shouldRetryWithCredits(credentials?.accessToken || "", creditsMode !== "off") - ) { - log?.info?.("AG_CREDITS", "Retrying with Google One AI credits"); - const creditsBody = injectCreditsField(transformedBody); - const serializedCreditsRequest = serializeAntigravityRequest( - this.provider, - headers, - creditsBody - ); - const finalCreditsHeaders = serializedCreditsRequest.headers; - try { - await capture(finalCreditsHeaders, serializedCreditsRequest.bodyString); - const creditsResp = await fetchWithReadinessTimeout(url, { - method: "POST", - headers: finalCreditsHeaders, - body: getChunkedOrFixedBody(serializedCreditsRequest.bodyString, stream), - ...(stream ? { duplex: "half" } : {}), - signal, - }); - if (creditsResp.ok || creditsResp.status !== HTTP_STATUS.RATE_LIMITED) { - log?.info?.("AG_CREDITS", `Credits retry succeeded: ${creditsResp.status}`); - if (!stream) { - const collected = await this.collectStreamToResponse( - creditsResp, - model, - url, - finalCreditsHeaders, - creditsBody, - log, - signal - ); - // Parse _remainingCredits from the synthetic response and cache - try { - const syntheticJson = await collected.response.clone().json(); - const rc = syntheticJson?._remainingCredits; - if (Array.isArray(rc)) { - const googleCredit = rc.find((c) => c.creditType === "GOOGLE_ONE_AI"); - if (googleCredit) { - const balance = parseInt(googleCredit.creditAmount, 10); - if (!isNaN(balance)) - updateAntigravityRemainingCredits(accountId, balance); - } - } - } catch { - /**/ - } - return { - ...collected, - transformedBody: attachToolNameMap(creditsBody, requestToolNameMap), - }; - } - return { - response: creditsResp, - url, - headers: finalCreditsHeaders, - transformedBody: attachToolNameMap(creditsBody, requestToolNameMap), - }; - } - - // Credit retry also 429'd - handleCreditsFailure(credentials?.accessToken || ""); - log?.warn?.("AG_CREDITS", "Credits retry also 429'd"); - - // Also mark in our legacy exhaustion map to avoid retrying other routes - markCreditsExhausted(accountId); - } catch (creditsErr) { - handleCreditsFailure(credentials?.accessToken || ""); - log?.warn?.("AG_CREDITS", `Credits retry failed: ${creditsErr}`); - } - } - } catch (e) { - // Ignore parse errors, will fall back to exponential backoff - } - } - - // Bounded short-retry: a non-null retryAfterMs ≤ 60s covers nearly every - // 429 (decide429 returns 2s/5s/60s defaults), so this branch MUST share the - // per-URL attempt counter. Without the bound a persistent 429 loops forever - // on the same endpoint/account (urlIndex-- cancels the loop's urlIndex++) and - // never returns the 429 to the account-fallback layer in chat.ts. - if ( - retryMs && - retryMs <= LONG_RETRY_THRESHOLD_MS && - retryAttemptsByUrl[urlIndex] < MAX_AUTO_RETRIES - ) { - retryAttemptsByUrl[urlIndex]++; - const effectiveRetryMs = Math.min(retryMs, MAX_RETRY_AFTER_MS); - log?.debug?.( - "RETRY", - `${response.status} retry ${retryAttemptsByUrl[urlIndex]}/${MAX_AUTO_RETRIES} with Retry-After: ${Math.ceil(effectiveRetryMs / 1000)}s, waiting...` - ); - await new Promise((resolve) => setTimeout(resolve, effectiveRetryMs)); - urlIndex--; - continue; - } - - // Auto retry for 429 (no Retry-After) or transient 5xx errors. - // For 5xx we read the body to detect known transient patterns - // ("Agent execution terminated due to error", "high traffic", "capacity"). - if ((!retryMs || retryMs === 0) && retryAttemptsByUrl[urlIndex] < MAX_AUTO_RETRIES) { - let shouldAutoRetry = response.status === HTTP_STATUS.RATE_LIMITED; - if (!shouldAutoRetry && ANTIGRAVITY_TRANSIENT_STATUSES.has(response.status)) { - try { - const errBody = await response.clone().text(); - let errJson: unknown = null; - try { - errJson = errBody ? JSON.parse(errBody) : null; - } catch { - // non-JSON body — fall through to pattern match against raw text - } - const errMsg = this.extractErrorMessage(errJson, errBody); - shouldAutoRetry = this.isTransientAntigravityError(response.status, errMsg); - } catch { - // ignore body read errors - } - } - if (shouldAutoRetry) { - retryAttemptsByUrl[urlIndex]++; - // Exponential backoff: 2s, 4s, 8s… capped per-status - const cap = - response.status === HTTP_STATUS.RATE_LIMITED - ? MAX_RETRY_AFTER_MS - : ANTIGRAVITY_TRANSIENT_RETRY_MAX_MS; - const backoffMs = Math.min(1000 * 2 ** retryAttemptsByUrl[urlIndex], cap); - log?.debug?.( - "RETRY", - `${response.status} transient auto retry ${retryAttemptsByUrl[urlIndex]}/${MAX_AUTO_RETRIES} after ${backoffMs / 1000}s` - ); - await new Promise((resolve) => setTimeout(resolve, backoffMs)); - urlIndex--; - continue; - } - } - - log?.debug?.( - "RETRY", - `${response.status}, Retry-After ${retryMs ? `too long (${Math.ceil(retryMs / 1000)}s)` : "missing"}, trying fallback` - ); - lastStatus = response.status; - - if (urlIndex + 1 < fallbackCount) { - continue; - } - } - - if (this.shouldRetry(response.status, urlIndex)) { - log?.debug?.("RETRY", `${response.status} on ${url}, trying fallback ${urlIndex + 1}`); - lastStatus = response.status; - continue; - } - - // If we have a 429 with long retry time, embed it in response body - if ( - response.status === HTTP_STATUS.RATE_LIMITED && - retryMs && - retryMs > LONG_RETRY_THRESHOLD_MS - ) { - try { - const respBody = await response.clone().text(); - let obj; - try { - obj = JSON.parse(respBody); - } catch { - obj = {}; - } - obj.retryAfterMs = retryMs; - const modifiedBody = JSON.stringify(obj); - const modifiedResponse = new Response(modifiedBody, { - status: response.status, - headers: response.headers, - }); - return { - response: modifiedResponse, - url, - headers: finalHeaders, - transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), - }; - } catch (err) { - log?.warn?.("RETRY", `Failed to embed retryAfterMs: ${err}`); - // Fall back to original response - } - } - - // For non-streaming clients, collect the SSE stream and return a synthetic - // non-streaming Response so chatCore doesn't need to handle SSE conversion. - if (!stream) { - // #3229: surface a real upstream error instead of masking a 4xx/5xx as an - // empty `chat.completion` envelope (collectStreamToResponse synthesizes a - // success-shaped body when the upstream returned no SSE data). - if (!response.ok) { - const rawBody = await response - .clone() - .text() - .catch(() => ""); - const errorBody = buildAntigravityUpstreamError( - response.status, - response.statusText, - rawBody - ); - return { - response: new Response(JSON.stringify(errorBody), { - status: response.status, - headers: { "Content-Type": "application/json" }, - }), - url, - headers: finalHeaders, - transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), - }; - } - const collected = await this.collectStreamToResponse( - response, - model, - url, - finalHeaders, - transformedBody, - log, - signal - ); - // When credits were injected (credits-first or credits-retry), the - // synthetic body contains _remainingCredits — mirror it into the - // balance cache so the dashboard stays fresh. - try { - const syntheticJson = await collected.response.clone().json(); - const rc = syntheticJson?._remainingCredits; - if (Array.isArray(rc)) { - const googleCredit = rc.find( - (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI" - ); - if (googleCredit) { - const balance = parseInt(googleCredit.creditAmount, 10); - if (!isNaN(balance)) updateAntigravityRemainingCredits(accountId, balance); - } - } - } catch { - /* balance cache is best-effort */ - } - return { - ...collected, - transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), - }; - } - - // #2461: a non-ok upstream response (e.g. 403) must never be piped through the - // streaming pass-through below as if it were an SSE body. Google occasionally - // returns non-UTF8/binary error bodies (observed: gzip-magic-byte payloads) for - // 403s on this endpoint; reading/forwarding those raw bytes corrupts the - // client-visible error message. Mirror the non-streaming branch above and build - // a sanitized JSON error via buildAntigravityUpstreamError (hard rule #12) - // instead of streaming unknown bytes straight through. - if (!response.ok) { - const rawBody = await response - .clone() - .text() - .catch(() => ""); - const errorBody = buildAntigravityUpstreamError( - response.status, - response.statusText, - rawBody - ); - return { - response: new Response(JSON.stringify(errorBody), { - status: response.status, - headers: { "Content-Type": "application/json" }, - }), - url, - headers: finalHeaders, - transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), - }; - } - - // Streaming path: wrap the response body in a pass-through TransformStream - // that extracts remainingCredits from the final SSE chunk(s) without - // consuming the stream. The client receives the unmodified SSE data. - if (response.body) { - // If the downstream client aborts, cancel the upstream fetch body immediately - // to release the socket back to the Undici agent pool and prevent memory leaks. - if (signal) { - const abortHandler = () => { - try { - response.body?.cancel().catch(() => {}); - } catch (_) {} - }; - if (signal.aborted) { - abortHandler(); - } else { - signal.addEventListener("abort", abortHandler, { once: true }); - } - } - - let sseBuffer = ""; - const decoder = new TextDecoder(); // Singleton for correct streaming decode - const MAX_BUFFER_SIZE = 16 * 1024; // Limit to prevent OOM on large streams - - const passThrough = new TransformStream( - { - transform(chunk, controller) { - controller.enqueue(chunk); - // Accumulate text to scan for remainingCredits - try { - const text = decoder.decode(chunk, { stream: true }); - sseBuffer += text; - // Limit buffer size to prevent unbounded growth - // Truncate only after a complete newline to avoid splitting SSE lines mid-payload - if (sseBuffer.length > MAX_BUFFER_SIZE) { - const lastNewline = sseBuffer.lastIndexOf( - "\n", - sseBuffer.length - MAX_BUFFER_SIZE - ); - if (lastNewline !== -1) { - sseBuffer = sseBuffer.slice(lastNewline + 1); - } else { - // No newline found in discard region — buffer contains an incomplete SSE line. - // Discard it entirely to avoid returning malformed data; the remainingCredits - // parser won't find valid data in a truncated line anyway. - sseBuffer = ""; - } - } - } catch { - /* decoding best-effort */ - } - }, - flush() { - // Final decode for any remaining bytes - try { - const text = decoder.decode(); // Flush pending bytes - sseBuffer += text; - } catch { - /* decoding best-effort */ - } - - // Parse the accumulated SSE data for remainingCredits - try { - const lines = sseBuffer.split("\n"); - for (const line of lines) { - const trimmed = line.trim(); - if (!trimmed.startsWith("data:")) continue; - const payload = trimmed.slice(5).trim(); - if (!payload || payload === "[DONE]") continue; - try { - const parsed = JSON.parse(payload); - if (Array.isArray(parsed?.remainingCredits)) { - const googleCredit = parsed.remainingCredits.find((c: unknown) => { - const credit = asRecord(c); - return credit?.creditType === "GOOGLE_ONE_AI"; - }) as AntigravityCreditEntry | undefined; - if (googleCredit) { - const balance = parseInt(String(googleCredit.creditAmount ?? ""), 10); - if (!isNaN(balance)) { - updateAntigravityRemainingCredits(accountId, balance); - } - } - } - } catch { - /* skip malformed lines */ - } - } - } catch { - /* credits extraction is best-effort */ - } - sseBuffer = ""; - }, - }, - { highWaterMark: 16384 }, - { highWaterMark: 16384 } - ); - const tappedBody = response.body.pipeThrough(passThrough); - const tappedResponse = new Response(tappedBody, { - status: response.status, - statusText: response.statusText, - headers: response.headers, - }); - return { - response: tappedResponse, - url, - headers: finalHeaders, - transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), - }; - } - - return { - response, - url, - headers: finalHeaders, - transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), - }; + if (outcome.action === "return") return outcome.result; + if (outcome.lastStatus !== undefined) lastStatus = outcome.lastStatus; + if (outcome.sameUrl) urlIndex--; + continue; } catch (error) { lastError = error; - log?.error?.( + l.error( "TELEMETRY", `[Antigravity] Network/Fetch Error - URL: ${url}, Model: ${model}, Error: ${error instanceof Error ? error.message : String(error)}` ); if (urlIndex + 1 < fallbackCount) { - log?.debug?.("RETRY", `Error on ${url}, trying fallback ${urlIndex + 1}`); + l.debug("RETRY", `Error on ${url}, trying fallback ${urlIndex + 1}`); continue; } throw error; @@ -1781,6 +1177,332 @@ export class AntigravityExecutor extends BaseExecutor { throw lastError || new Error(`All ${fallbackCount} URLs failed with status ${lastStatus}`); } + + /** + * Run one full per-url-index attempt: send the request, handle a 429/503 (retry + * same/next url, or a Google One AI credits retry), fall back on other retryable + * statuses, optionally embed a long Retry-After, then build the final non-streaming + * or streaming result. Returns a result to hand back from execute(), or a retry + * instruction for executeOnce()'s loop to act on (continue, optionally urlIndex--). + */ + private async runAntigravityAttempt( + ctx: AntigravityAttemptContext + ): Promise { + const { + url, + model, + headers, + transformedBody, + requestToolNameMap, + credentials, + stream, + signal, + log, + accountId, + urlIndex, + retryAttemptsByUrl, + fallbackCount, + } = ctx; + + const { response, finalHeaders } = await sendAntigravityRequest( + this.provider, + url, + model, + headers, + transformedBody, + credentials, + stream, + signal, + log, + retryAttemptsByUrl[urlIndex] + ); + + let retryMs: number | null = null; + + if ( + response.status === HTTP_STATUS.RATE_LIMITED || + response.status === HTTP_STATUS.SERVICE_UNAVAILABLE + ) { + const rateLimitOutcome = await this.handleAntigravityRateLimit({ + ...ctx, + response, + finalHeaders, + }); + + if (rateLimitOutcome.action === "return") { + return { action: "return", result: rateLimitOutcome.result }; + } + if (rateLimitOutcome.action === "retrySameUrl") return { action: "retry", sameUrl: true }; + if (rateLimitOutcome.action === "retryNextUrl") { + return { action: "retry", sameUrl: false, lastStatus: rateLimitOutcome.lastStatus }; + } + // Only "fallthrough" remains: last url, no more retries — proceed below with + // the resolved retryMs so a long Retry-After can still be embedded in the body. + retryMs = rateLimitOutcome.retryMs; + } + + if (this.shouldRetry(response.status, urlIndex)) { + log.debug("RETRY", `${response.status} on ${url}, trying fallback ${urlIndex + 1}`); + return { action: "retry", sameUrl: false, lastStatus: response.status }; + } + + // If we have a 429 with long retry time, embed it in response body + const embedded = await tryEmbedLongRetryAfter( + response, + retryMs, + url, + finalHeaders, + transformedBody, + requestToolNameMap, + log + ); + if (embedded) return { action: "return", result: embedded }; + + const result = await this.buildAntigravityAttemptResult( + model, + stream, + response, + url, + finalHeaders, + transformedBody, + requestToolNameMap, + accountId, + signal, + log + ); + return { action: "return", result }; + } + + /** + * #3786 — Non-streaming callers (stream: false) keep the buffered + * collect-to-JSON contract: `execute()` (including the Pro-family + * fallback-chain retry loop) inspects `result.response` directly and + * expects a synthesized `chat.completion` JSON body, not a raw SSE + * pass-through. Passthrough is reserved for actual streaming clients + * (buildFinalAntigravityResult's stream:true branch), where the client + * itself drains the SSE bytes — collectStreamToResponse already uses + * FETCH_TIMEOUT_MS (no hardcoded 120s ceiling), so long-thinking models + * are not penalized by buffering here. + */ + private async buildAntigravityAttemptResult( + model: string, + stream: boolean, + response: Response, + url: string, + finalHeaders: Record, + transformedBody: Record, + requestToolNameMap: Map | null, + accountId: string, + signal: AbortSignal | null | undefined, + log: SafeAntigravityLog + ): Promise { + if (!stream && response.ok && response.body) { + return this.collectStreamToResponse( + response, + model, + url, + finalHeaders, + transformedBody, + log, + signal + ); + } + + return buildFinalAntigravityResult( + stream, + response, + url, + finalHeaders, + transformedBody, + requestToolNameMap, + accountId, + signal, + updateAntigravityRemainingCredits + ); + } + + /** + * Handle a 429/503 response for one URL-index attempt: resolve the retry-after + * time (headers, then error-body classification + Google-One-AI credits retry), + * then decide whether to retry the SAME url, fall back to the NEXT url, or (on + * the last url with no more retries left) fall through with the resolved retryMs + * so the caller can still embed a long Retry-After in the final response body. + */ + private async handleAntigravityRateLimit( + ctx: AntigravityRateLimitContext + ): Promise { + const { response, log, urlIndex, retryAttemptsByUrl, fallbackCount } = ctx; + + // Try to get retry time from headers first + let retryMs: number | null = this.parseRetryHeaders(response.headers); + + // If no retry time in headers, try to parse from error message body + if (!retryMs) { + const resolved = await this.tryResolveRetryFromErrorBody(ctx); + if (resolved.kind === "return") return { action: "return", result: resolved.result }; + retryMs = resolved.retryMs; + } + + // Bounded short-retry: a non-null retryAfterMs ≤ 60s covers nearly every + // 429 (decide429 returns 2s/5s/60s defaults), so this branch MUST share the + // per-URL attempt counter. Without the bound a persistent 429 loops forever + // on the same endpoint/account (urlIndex-- cancels the loop's urlIndex++) and + // never returns the 429 to the account-fallback layer in chat.ts. + if ( + retryMs && + retryMs <= LONG_RETRY_THRESHOLD_MS && + retryAttemptsByUrl[urlIndex] < MAX_AUTO_RETRIES + ) { + retryAttemptsByUrl[urlIndex]++; + const effectiveRetryMs = Math.min(retryMs, MAX_RETRY_AFTER_MS); + log.debug( + "RETRY", + `${response.status} retry ${retryAttemptsByUrl[urlIndex]}/${MAX_AUTO_RETRIES} with Retry-After: ${Math.ceil(effectiveRetryMs / 1000)}s, waiting...` + ); + await new Promise((resolve) => setTimeout(resolve, effectiveRetryMs)); + return { action: "retrySameUrl" }; + } + + // Auto retry for 429 (no Retry-After) or transient 5xx errors. + // For 5xx we read the body to detect known transient patterns + // ("Agent execution terminated due to error", "high traffic", "capacity"). + if ((!retryMs || retryMs === 0) && retryAttemptsByUrl[urlIndex] < MAX_AUTO_RETRIES) { + const shouldAutoRetry = await this.shouldAutoRetryTransient(response); + if (shouldAutoRetry) { + retryAttemptsByUrl[urlIndex]++; + // Exponential backoff: 2s, 4s, 8s… capped per-status + const cap = + response.status === HTTP_STATUS.RATE_LIMITED + ? MAX_RETRY_AFTER_MS + : ANTIGRAVITY_TRANSIENT_RETRY_MAX_MS; + const backoffMs = Math.min(1000 * 2 ** retryAttemptsByUrl[urlIndex], cap); + log.debug( + "RETRY", + `${response.status} transient auto retry ${retryAttemptsByUrl[urlIndex]}/${MAX_AUTO_RETRIES} after ${backoffMs / 1000}s` + ); + await new Promise((resolve) => setTimeout(resolve, backoffMs)); + return { action: "retrySameUrl" }; + } + } + + log.debug( + "RETRY", + `${response.status}, Retry-After ${retryMs ? `too long (${Math.ceil(retryMs / 1000)}s)` : "missing"}, trying fallback` + ); + + if (urlIndex + 1 < fallbackCount) { + return { action: "retryNextUrl", lastStatus: response.status }; + } + + return { action: "fallthrough", retryMs, lastStatus: response.status }; + } + + /** + * Parse the 429/503 response body to classify the failure and (for + * quota_exhausted, non-full-exhaustion cases) attempt a Google One AI + * credits retry. Returns the resolved retryMs, or an early "return" result + * when the credits retry itself produced a response to hand back to the client. + */ + private async tryResolveRetryFromErrorBody( + ctx: AntigravityRateLimitContext + ): Promise< + { kind: "return"; result: SsePassthroughResult } | { kind: "resolved"; retryMs: number | null } + > { + const { + response, + url, + headers, + transformedBody, + requestToolNameMap, + credentials, + stream, + signal, + log, + accountId, + creditsMode, + } = ctx; + + try { + const errorBody = await response.clone().text(); + const errorJson = JSON.parse(errorBody); + const errorMessage = buildAntigravity429ErrorMessage(errorJson); + + // 1. Try to parse explicit retry time from message + const parsedRetryMs = this.parseRetryFromErrorMessage(errorMessage); + + // 2. Classify 429, then decide the final retry time BEFORE the credits + // retry so that full_quota_exhausted can skip the credits attempt + // entirely (avoids ~41s hold on an already-exhausted account) and + // persist the cooldown to DB for post-restart routing. + const category = classify429(errorMessage); + const decision: Decision = decide429(category, parsedRetryMs); + const retryMs = decision.retryAfterMs; + log.debug("AG_429", `Category: ${category}, Decision: ${decision.kind} — ${decision.reason}`); + + if (decision.kind === "full_quota_exhausted" && retryMs) { + markConnectionQuotaExhausted(accountId, retryMs); + } + + const creditsAlreadyInjected = + (transformedBody as { enabledCreditTypes?: unknown }).enabledCreditTypes != null; + + if (category === "quota_exhausted" && creditsAlreadyInjected) { + handleCreditsFailure(credentials?.accessToken || ""); + log.warn("AG_CREDITS", "Credits-first request 429'd — credits likely exhausted"); + markCreditsExhausted(accountId); + } + + if ( + category === "quota_exhausted" && + decision.kind !== "full_quota_exhausted" && + !creditsAlreadyInjected && + shouldRetryWithCredits(credentials?.accessToken || "", creditsMode !== "off") + ) { + const creditsResult = await tryCreditsRetry( + this.provider, + url, + headers, + transformedBody, + requestToolNameMap, + credentials, + stream, + signal, + log, + accountId, + updateAntigravityRemainingCredits + ); + if (creditsResult) return { kind: "return", result: creditsResult }; + } + + return { kind: "resolved", retryMs }; + } catch { + // Ignore parse errors, will fall back to exponential backoff + return { kind: "resolved", retryMs: null }; + } + } + + /** + * True for 429 always; for transient 5xx (500/502/503/504) only when the body + * matches a known capacity/traffic/agent-terminated pattern. + */ + private async shouldAutoRetryTransient(response: Response): Promise { + if (response.status === HTTP_STATUS.RATE_LIMITED) return true; + if (!ANTIGRAVITY_TRANSIENT_STATUSES.has(response.status)) return false; + try { + const errBody = await response.clone().text(); + let errJson: unknown = null; + try { + errJson = errBody ? JSON.parse(errBody) : null; + } catch { + // non-JSON body — fall through to pattern match against raw text + } + const errMsg = this.extractErrorMessage(errJson, errBody); + return this.isTransientAntigravityError(response.status, errMsg); + } catch { + // ignore body read errors + return false; + } + } } export default AntigravityExecutor; diff --git a/open-sse/executors/antigravity/executeAttempt.ts b/open-sse/executors/antigravity/executeAttempt.ts new file mode 100644 index 0000000000..566e651713 --- /dev/null +++ b/open-sse/executors/antigravity/executeAttempt.ts @@ -0,0 +1,687 @@ +// Pure-ish per-attempt request/result helpers for the Antigravity executor (#7408 +// complexity-gate decomposition): building + sending one upstream request, and +// building the final non-streaming/streaming result. No host state of their own — +// callers inject `provider` and `onCreditsUpdate` so this module doesn't need to +// import the executor's credit-balance cache. Extracted from antigravity.ts +// (file-size cap), mirroring the existing antigravity/streamingPassthrough.ts and +// antigravity/sseCollect.ts submodule pattern. +import { mergeAbortSignals, type ExecutorLog } from "../base.ts"; +import { applyFingerprint, isCliCompatEnabled } from "../../config/cliFingerprints.ts"; +import { buildAntigravityUpstreamError } from "../antigravityUpstreamError.ts"; +import { + HTTP_STATUS, + STREAM_READINESS_TIMEOUT_MS, + ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, +} from "../../config/constants.ts"; +import { injectCreditsField, handleCreditsFailure } from "../../services/antigravityCredits.ts"; +import { cloakAntigravityToolPayload } from "../../config/toolCloaking.ts"; +import { + applyAntigravityClientProfileHeaders, + removeHeaderCaseInsensitive, +} from "../../services/antigravityClientProfile.ts"; +import * as prl from "../../utils/providerRequestLogging.ts"; +import { + createCreditsExtractionTransform as createCreditsExtractionTransformImpl, + buildSsePassthroughResult, + type SsePassthroughResult, +} from "./streamingPassthrough.ts"; +import type { AntigravityCredentials } from "../antigravity.ts"; + +const LONG_RETRY_THRESHOLD_MS = 60_000; +const CREDITS_EXHAUSTED_TTL_MS = 5 * 60 * 60 * 1000; // 5 hours + +/** Invoked with a fresh GOOGLE_ONE_AI credit balance to persist in the caller's cache. */ +export type OnAntigravityCreditsUpdate = (accountId: string, balance: number) => void; + +/** + * Per-account GOOGLE_ONE_AI credits-exhausted tracker. + * Key: accountId (OAuth subject / email). Value: expiry timestamp. + * When credits hit 0 we skip the credit retry for CREDITS_EXHAUSTED_TTL_MS. + * Lives here (not antigravity.ts) so both this module's tryCreditsRetry and + * antigravity.ts's tryResolveRetryFromErrorBody can share it via a single import + * direction (antigravity.ts -> executeAttempt.ts), avoiding a circular import. + */ +const MAX_CREDITS_EXHAUSTED_ENTRIES = 50; +const creditsExhaustedUntil = new Map(); + +const _creditsExhaustedSweep = setInterval(() => { + const now = Date.now(); + for (const [key, until] of creditsExhaustedUntil) { + if (now >= until) creditsExhaustedUntil.delete(key); + } +}, 60_000); +if (typeof _creditsExhaustedSweep === "object" && "unref" in _creditsExhaustedSweep) { + (_creditsExhaustedSweep as { unref?: () => void }).unref?.(); +} + +/** True while `accountId`'s Google One AI credits are marked exhausted. @internal */ +export function isCreditsExhausted(accountId: string): boolean { + const until = creditsExhaustedUntil.get(accountId); + if (!until) return false; + if (Date.now() >= until) { + creditsExhaustedUntil.delete(accountId); + return false; + } + return true; +} + +/** Mark an account's Google One AI credits as exhausted for CREDITS_EXHAUSTED_TTL_MS. */ +export function markCreditsExhausted(accountId: string): void { + if ( + creditsExhaustedUntil.size >= MAX_CREDITS_EXHAUSTED_ENTRIES && + !creditsExhaustedUntil.has(accountId) + ) { + const now = Date.now(); + for (const [key, until] of creditsExhaustedUntil) { + if (now >= until) { + creditsExhaustedUntil.delete(key); + } + } + if (creditsExhaustedUntil.size >= MAX_CREDITS_EXHAUSTED_ENTRIES) { + const oldestKey = creditsExhaustedUntil.keys().next().value; + if (oldestKey !== undefined) creditsExhaustedUntil.delete(oldestKey); + } + } + creditsExhaustedUntil.set(accountId, Date.now() + CREDITS_EXHAUSTED_TTL_MS); +} + +class AntigravityPreResponseTimeoutError extends Error { + code = ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; + status = HTTP_STATUS.GATEWAY_TIMEOUT; + + constructor(timeoutMs: number, url: string) { + super(`Antigravity upstream did not return response headers within ${timeoutMs}ms: ${url}`); + this.name = "TimeoutError"; + } +} + +function getAbortErrorCode(error: unknown): string | null { + if (!error || typeof error !== "object") return null; + const value = (error as { code?: unknown }).code; + return typeof value === "string" ? value : null; +} + +function isAntigravityPreResponseTimeout(error: unknown): boolean { + return getAbortErrorCode(error) === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE; +} + +/** + * `fetch()` wrapper that aborts if the upstream never returns response headers + * within `timeoutMs` (default STREAM_READINESS_TIMEOUT_MS) — distinct from the + * overall FETCH_TIMEOUT_MS, which bounds the whole request including body streaming. + * Shared by every fetch attempt in executeOnce() (initial, 403-retry, credits-retry). + */ +export async function fetchAntigravityWithReadinessTimeout( + url: string, + init: RequestInit, + timeoutMs = STREAM_READINESS_TIMEOUT_MS +): Promise { + const boundedTimeoutMs = Math.max(0, Math.floor(timeoutMs)); + if (boundedTimeoutMs <= 0) { + return fetch(url, init); + } + + const timeoutController = new AbortController(); + let timeoutId: ReturnType | null = setTimeout(() => { + timeoutController.abort(new AntigravityPreResponseTimeoutError(boundedTimeoutMs, url)); + }, boundedTimeoutMs); + + const existingSignal = init.signal instanceof AbortSignal ? init.signal : null; + const combinedSignal = existingSignal + ? mergeAbortSignals(existingSignal, timeoutController.signal) + : timeoutController.signal; + + try { + return await fetch(url, { ...init, signal: combinedSignal }); + } catch (error) { + if ( + timeoutController.signal.aborted && + isAntigravityPreResponseTimeout(timeoutController.signal.reason) + ) { + throw timeoutController.signal.reason; + } + throw error; + } finally { + if (timeoutId) { + clearTimeout(timeoutId); + timeoutId = null; + } + } +} + +/** ExecutorLog with every method always callable — see toSafeAntigravityLog(). */ +export type SafeAntigravityLog = Required; + +function noopLogFn(): void {} + +/** + * Normalize a possibly-null/undefined ExecutorLog into an object with all four + * methods always callable, so the request/retry helpers below can call + * `l.debug(...)` directly instead of repeating `log?.debug?.(...)` at every call + * site. This isn't just style: the complexity linter (eslint `complexity` rule) + * weighs each `?.` link in a chain as its own branch — a doubly-chained + * `log?.debug?.(...)` costs +2 — so a logging-heavy helper can rack up a large + * complexity score with zero real decision points. Resolving once here keeps + * the actual branch count legible in the functions that matter. + */ +export function toSafeAntigravityLog(log: ExecutorLog | null | undefined): SafeAntigravityLog { + return { + debug: log?.debug ? log.debug.bind(log) : noopLogFn, + info: log?.info ? log.info.bind(log) : noopLogFn, + warn: log?.warn ? log.warn.bind(log) : noopLogFn, + error: log?.error ? log.error.bind(log) : noopLogFn, + }; +} + +/** Flatten a 429/503 error JSON body (message + `error.details[].reason`) into one string. */ +export function buildAntigravity429ErrorMessage(errorJson: unknown): string { + const obj = errorJson as + | { error?: { message?: unknown; details?: unknown }; message?: unknown } + | null + | undefined; + let errorMessage = String(obj?.error?.message || obj?.message || ""); + const details = obj?.error?.details; + if (Array.isArray(details)) { + for (const detail of details) { + const reason = (detail as { reason?: unknown } | null)?.reason; + if (reason) errorMessage += ` ${reason}`; + } + } + return errorMessage; +} + +function getChunkedOrFixedBody(bodyStr: string, stream: boolean): BodyInit { + if (stream) { + return new ReadableStream( + { + async start(controller) { + controller.enqueue(new TextEncoder().encode(bodyStr)); + controller.close(); + }, + }, + { highWaterMark: 16384 } + ); + } + return bodyStr; +} + +function cloneAntigravityRequestBody(body: unknown): unknown { + if (!body || typeof body !== "object") { + return body; + } + + try { + return structuredClone(body); + } catch { + return JSON.parse(JSON.stringify(body)); + } +} + +function serializeAntigravityRequest( + provider: string, + headers: Record, + body: unknown +): { headers: Record; bodyString: string } { + const serializedBody = cloneAntigravityRequestBody(body); + + if (!isCliCompatEnabled(provider)) { + return { headers, bodyString: JSON.stringify(serializedBody) }; + } + return applyFingerprint(provider, { ...headers }, serializedBody); +} + +function getRequestTargetModel(body: Record): string { + const target = body.model; + return typeof target === "string" && target.length > 0 ? target : "unknown"; +} + +function attachToolNameMap(payload: T, toolNameMap: Map | null): T { + if (!toolNameMap?.size || !payload || typeof payload !== "object") { + return payload; + } + + const copy = Array.isArray(payload) ? ([...payload] as T) : ({ ...(payload as object) } as T); + Object.defineProperty(copy, "_toolNameMap", { + value: toolNameMap, + enumerable: false, + configurable: true, + writable: true, + }); + return copy; +} + +/** Cloak the tool-name payload, then apply credits-first injection, for one attempt. */ +export function finalizeAntigravityRequestBody( + transformed: Record, + useCreditsFirst: boolean, + log: SafeAntigravityLog +): { + transformedBody: Record; + requestToolNameMap: Map | null; +} { + let transformedBody: Record = transformed; + let requestToolNameMap: Map | null = null; + + if (transformedBody && typeof transformedBody === "object") { + const cloaked = cloakAntigravityToolPayload(transformedBody); + transformedBody = cloaked.body; + requestToolNameMap = cloaked.toolNameMap; + } + + // Credits-first: inject GOOGLE_ONE_AI upfront so we never try the normal + // quota path. If credits are exhausted / disabled shouldUseCreditsFirst() + // returns false and we fall back to the legacy retry-on-429 flow. + if (useCreditsFirst) { + transformedBody = injectCreditsField(transformedBody); + log.debug("AG_CREDITS", "Credits-first enabled (ANTIGRAVITY_CREDITS=always)"); + } + + return { transformedBody, requestToolNameMap }; +} + +/** Debug-only dump of outgoing headers (mask Authorization) and envelope shape. */ +function dumpAntigravityRequestDebug( + finalHeaders: Record, + transformedBody: Record, + clientProfile: unknown, + log: SafeAntigravityLog +): void { + const safeHeaders = { ...finalHeaders }; + if (safeHeaders["Authorization"]) safeHeaders["Authorization"] = "Bearer ***"; + log.debug("AG_REQUEST_HEADERS", JSON.stringify(safeHeaders)); + + const envelope = transformedBody as Record; + const requestInner = envelope.request as Record | undefined; + log.debug( + "AG_REQUEST_ENVELOPE", + JSON.stringify({ + fieldOrder: Object.keys(envelope), + project: envelope.project, + requestId: envelope.requestId, + model: envelope.model, + userAgent: envelope.userAgent, + requestType: envelope.requestType, + enabledCreditTypes: envelope.enabledCreditTypes, + clientProfile, + sessionId: requestInner?.sessionId, + generationConfig: requestInner?.generationConfig, + }) + ); +} + +/** + * Send one Antigravity request attempt: serialize + apply the client-profile + * fingerprint, debug-dump the outgoing envelope, fetch with a readiness timeout, + * and transparently retry once without `x-goog-user-project` on a 403 (some + * projects reject that header). Returns the (possibly 403-retried) response + * plus the headers actually used for it. + */ +export async function sendAntigravityRequest( + provider: string, + url: string, + model: string, + headers: Record, + transformedBody: Record, + credentials: AntigravityCredentials, + stream: boolean, + signal: AbortSignal | null | undefined, + log: SafeAntigravityLog, + retryAttempt: number +): Promise<{ response: Response; finalHeaders: Record }> { + const serializedRequest = serializeAntigravityRequest(provider, headers, transformedBody); + let finalHeaders = serializedRequest.headers; + const clientProfile = applyAntigravityClientProfileHeaders( + finalHeaders, + credentials, + transformedBody + ); + + log.debug( + "TELEMETRY", + `[Antigravity] Execute - URL: ${url}, Model: ${model}, Target: ${getRequestTargetModel(transformedBody)}, RetryAttempt: ${retryAttempt}` + ); + + // Dump outgoing headers (mask Authorization) and envelope shape for debugging. + // Gated behind an explicit typeof check (not just calling log.debug() unconditionally) + // so the JSON.stringify work below is skipped entirely when debug logging is off. + if (typeof log.debug === "function") { + dumpAntigravityRequestDebug(finalHeaders, transformedBody, clientProfile, log); + } + + await prl.captureCurrentProviderBody(url, finalHeaders, serializedRequest.bodyString, log); + let response = await fetchAntigravityWithReadinessTimeout(url, { + method: "POST", + headers: finalHeaders, + body: getChunkedOrFixedBody(serializedRequest.bodyString, stream), + ...(stream ? { duplex: "half" } : {}), + signal, + }); + + if (response.status === HTTP_STATUS.FORBIDDEN && finalHeaders["x-goog-user-project"]) { + const retryHeaders = { ...finalHeaders }; + removeHeaderCaseInsensitive(retryHeaders, "x-goog-user-project"); + log.debug("RETRY", "403 with x-goog-user-project, retrying once without it"); + await prl.captureCurrentProviderBody(url, retryHeaders, serializedRequest.bodyString, log); + response = await fetchAntigravityWithReadinessTimeout(url, { + method: "POST", + headers: retryHeaders, + body: getChunkedOrFixedBody(serializedRequest.bodyString, stream), + ...(stream ? { duplex: "half" } : {}), + signal, + }); + finalHeaders = retryHeaders; + } + + if (!response.ok) { + log.warn( + "TELEMETRY", + `[Antigravity] Error Response - URL: ${url}, Status: ${response.status}, Model: ${model}` + ); + } + + return { response, finalHeaders }; +} + +/** + * Retry the SAME url with `enabledCreditTypes: ["GOOGLE_ONE_AI"]` injected, for a + * quota_exhausted 429 that hasn't already tried credits. Returns the result to hand + * back to the caller of execute() on success (or a non-429 status), or null if the + * credits retry also failed/429'd (caller falls through to the normal retry logic). + */ +export async function tryCreditsRetry( + provider: string, + url: string, + headers: Record, + transformedBody: Record, + requestToolNameMap: Map | null, + credentials: AntigravityCredentials, + stream: boolean, + signal: AbortSignal | null | undefined, + log: SafeAntigravityLog, + accountId: string, + onCreditsUpdate: OnAntigravityCreditsUpdate +): Promise { + log.info("AG_CREDITS", "Retrying with Google One AI credits"); + const creditsBody = injectCreditsField(transformedBody); + const serializedCreditsRequest = serializeAntigravityRequest(provider, headers, creditsBody); + const finalCreditsHeaders = serializedCreditsRequest.headers; + try { + await prl.captureCurrentProviderBody( + url, + finalCreditsHeaders, + serializedCreditsRequest.bodyString, + log + ); + const creditsResp = await fetchAntigravityWithReadinessTimeout(url, { + method: "POST", + headers: finalCreditsHeaders, + body: getChunkedOrFixedBody(serializedCreditsRequest.bodyString, stream), + ...(stream ? { duplex: "half" } : {}), + signal, + }); + if (creditsResp.ok || creditsResp.status !== HTTP_STATUS.RATE_LIMITED) { + log.info("AG_CREDITS", `Credits retry succeeded: ${creditsResp.status}`); + if (!stream && creditsResp.body) { + // Raw SSE pass-through + credits extraction (see + // streamingPassthrough.ts); 499s early if the client + // already disconnected instead of piping a cancelled body. + return buildSsePassthroughResult( + creditsResp.body, + creditsResp, + accountId, + onCreditsUpdate, + url, + finalCreditsHeaders, + attachToolNameMap(creditsBody, requestToolNameMap), + signal + ); + } + return { + response: creditsResp, + url, + headers: finalCreditsHeaders, + transformedBody: attachToolNameMap(creditsBody, requestToolNameMap), + }; + } + + // Credit retry also 429'd + handleCreditsFailure(credentials?.accessToken || ""); + log.warn("AG_CREDITS", "Credits retry also 429'd"); + + // Also mark in our legacy exhaustion map to avoid retrying other routes + markCreditsExhausted(accountId); + return null; + } catch (creditsErr) { + handleCreditsFailure(credentials?.accessToken || ""); + log.warn("AG_CREDITS", `Credits retry failed: ${creditsErr}`); + return null; + } +} + +/** + * If we have a 429 with a long retry time (> LONG_RETRY_THRESHOLD_MS), embed + * `retryAfterMs` in the response body so the caller (combo/account-fallback + * layer) can read it back out. Returns null (fall back to the original + * response handling) when the status/retryMs don't qualify, or on error. + */ +export async function tryEmbedLongRetryAfter( + response: Response, + retryMs: number | null, + url: string, + finalHeaders: Record, + transformedBody: Record, + requestToolNameMap: Map | null, + log: ExecutorLog | null | undefined +): Promise { + if ( + response.status !== HTTP_STATUS.RATE_LIMITED || + !retryMs || + retryMs <= LONG_RETRY_THRESHOLD_MS + ) { + return null; + } + try { + const respBody = await response.clone().text(); + let obj; + try { + obj = JSON.parse(respBody); + } catch { + obj = {}; + } + obj.retryAfterMs = retryMs; + const modifiedBody = JSON.stringify(obj); + const modifiedResponse = new Response(modifiedBody, { + status: response.status, + headers: response.headers, + }); + return { + response: modifiedResponse, + url, + headers: finalHeaders, + transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), + }; + } catch (err) { + log?.warn?.("RETRY", `Failed to embed retryAfterMs: ${err}`); + return null; + } +} + +/** Build the sanitized JSON error result shared by the non-streaming and streaming paths. */ +async function buildUpstreamErrorResult( + response: Response, + url: string, + finalHeaders: Record, + transformedBody: Record, + requestToolNameMap: Map | null +): Promise { + const rawBody = await response + .clone() + .text() + .catch(() => ""); + const errorBody = buildAntigravityUpstreamError(response.status, response.statusText, rawBody); + return { + response: new Response(JSON.stringify(errorBody), { + status: response.status, + headers: { "Content-Type": "application/json" }, + }), + url, + headers: finalHeaders, + transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), + }; +} + +/** + * For non-streaming clients, return the raw SSE stream with a + * credits-extraction TransformStream. chatCore's non-streaming path + * (readNonStreamingResponseBody + parseNonStreamingSSEPayload with + * Gemini format support) handles draining and conversion to JSON. + * This replaces the previous collectStreamToResponse() approach which + * had an artificial timeout (now the standard FETCH_BODY_TIMEOUT_MS + * of 10 min applies). + */ +async function buildNonStreamingExecuteOnceResult( + response: Response, + url: string, + finalHeaders: Record, + transformedBody: Record, + requestToolNameMap: Map | null, + accountId: string, + signal: AbortSignal | null | undefined, + onCreditsUpdate: OnAntigravityCreditsUpdate +): Promise { + // #3229: surface a real upstream error instead of masking a 4xx/5xx as an + // empty `chat.completion` envelope. + if (!response.ok) { + return buildUpstreamErrorResult(response, url, finalHeaders, transformedBody, requestToolNameMap); + } + + if (response.body) { + // Raw SSE pass-through + credits extraction (see + // streamingPassthrough.ts); 499s early if the client already + // disconnected instead of piping a cancelled body. + return buildSsePassthroughResult( + response.body, + response, + accountId, + onCreditsUpdate, + url, + finalHeaders, + attachToolNameMap(transformedBody, requestToolNameMap), + signal + ); + } + + // No body -- return as-is + return { + response, + url, + headers: finalHeaders, + transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), + }; +} + +/** + * Streaming path: wrap the response body in a pass-through TransformStream + * that extracts remainingCredits from the final SSE chunk(s) without + * consuming the stream. The client receives the unmodified SSE data. + * + * #2461: a non-ok upstream response (e.g. 403) must never be piped through the + * streaming pass-through below as if it were an SSE body. Google occasionally + * returns non-UTF8/binary error bodies (observed: gzip-magic-byte payloads) for + * 403s on this endpoint; reading/forwarding those raw bytes corrupts the + * client-visible error message. Mirror the non-streaming branch above and build + * a sanitized JSON error via buildAntigravityUpstreamError (hard rule #12) + * instead of streaming unknown bytes straight through. + */ +async function buildStreamingExecuteOnceResult( + response: Response, + url: string, + finalHeaders: Record, + transformedBody: Record, + requestToolNameMap: Map | null, + accountId: string, + signal: AbortSignal | null | undefined, + onCreditsUpdate: OnAntigravityCreditsUpdate +): Promise { + if (!response.ok) { + return buildUpstreamErrorResult(response, url, finalHeaders, transformedBody, requestToolNameMap); + } + + if (response.body) { + // If the downstream client aborts, cancel the upstream fetch body immediately + // to release the socket back to the Undici agent pool and prevent memory leaks. + if (signal) { + const abortHandler = () => { + try { + response.body?.cancel().catch(() => {}); + } catch (_) {} + }; + if (signal.aborted) { + abortHandler(); + } else { + signal.addEventListener("abort", abortHandler, { once: true }); + } + } + + const passThrough = createCreditsExtractionTransformImpl( + accountId, + onCreditsUpdate, + 16 * 1024 // 16KB sliding-window cap to prevent OOM + ); + const tappedBody = response.body.pipeThrough(passThrough); + const tappedResponse = new Response(tappedBody, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }); + return { + response: tappedResponse, + url, + headers: finalHeaders, + transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), + }; + } + + return { + response, + url, + headers: finalHeaders, + transformedBody: attachToolNameMap(transformedBody, requestToolNameMap), + }; +} + +/** Dispatch to the non-streaming or streaming final-result builder. */ +export async function buildFinalAntigravityResult( + stream: boolean, + response: Response, + url: string, + finalHeaders: Record, + transformedBody: Record, + requestToolNameMap: Map | null, + accountId: string, + signal: AbortSignal | null | undefined, + onCreditsUpdate: OnAntigravityCreditsUpdate +): Promise { + if (!stream) { + return buildNonStreamingExecuteOnceResult( + response, + url, + finalHeaders, + transformedBody, + requestToolNameMap, + accountId, + signal, + onCreditsUpdate + ); + } + return buildStreamingExecuteOnceResult( + response, + url, + finalHeaders, + transformedBody, + requestToolNameMap, + accountId, + signal, + onCreditsUpdate + ); +} diff --git a/open-sse/executors/antigravity/streamingPassthrough.ts b/open-sse/executors/antigravity/streamingPassthrough.ts new file mode 100644 index 0000000000..69282e030c --- /dev/null +++ b/open-sse/executors/antigravity/streamingPassthrough.ts @@ -0,0 +1,176 @@ +// Pure streaming pass-through helpers for the Antigravity executor (#7408): +// tap an upstream Gemini SSE Response through a credits-extraction +// TransformStream instead of buffering the whole body in the executor, so +// long-thinking models aren't killed by an artificial collection timeout. +// Extracted from antigravity.ts (no host state, no fetch/auth) -- the +// credit-balance cache itself stays in antigravity.ts; callers inject the +// update function below so the two modules don't import each other. + +/** Shape of one entry in a Gemini `remainingCredits` SSE payload array. */ +export type AntigravityCreditEntry = { + creditType?: string; + creditAmount?: string; +}; + +function asCreditRecord(value: unknown): Record | null { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as Record) + : null; +} + +/** + * Create a pass-through TransformStream that extracts `remainingCredits` + * from SSE data without consuming the stream. The downstream client + * receives the unmodified bytes. + * + * @param accountId Provider account ID for credit-balance persistence. + * @param onCreditsUpdate Invoked with the parsed GOOGLE_ONE_AI balance. + * Injected by the caller (antigravity.ts's + * updateAntigravityRemainingCredits) to avoid this module + * importing back the executor's credit-balance cache. + * @param bufferSize Optional sliding-window buffer cap in bytes. + * Pass 0 or omit for unlimited (non-streaming callers + * where the full body is already buffered upstream). + * The streaming path uses 16384 (16 KB) to prevent OOM + * on long-lived SSE connections. Credit-balance data + * appears near the end of the SSE stream (after + * content), so the sliding window captures it even at + * 16 KB -- only truly massive responses (>16 KB of + * consecutive non-newline content) would lose credits. + */ +export function createCreditsExtractionTransform( + accountId: string, + onCreditsUpdate: (accountId: string, balance: number) => void, + bufferSize = 0 +): TransformStream { + let buffer = ""; + const decoder = new TextDecoder(); + + return new TransformStream( + { + transform(chunk, controller) { + controller.enqueue(chunk); + try { + buffer += decoder.decode(chunk, { stream: true }); + // Sliding-window cap: truncate after the last complete newline + // in the discard region so SSE lines are never split mid-payload. + if (bufferSize > 0 && buffer.length > bufferSize) { + const lastNewline = buffer.lastIndexOf("\n", buffer.length - bufferSize); + if (lastNewline !== -1) { + buffer = buffer.slice(lastNewline + 1); + } else { + // No newline in the discard region -- incomplete line, discard entirely. + buffer = ""; + } + } + } catch { + /* decoding best-effort */ + } + }, + flush() { + try { + buffer += decoder.decode(); + } catch { + /* decoding best-effort */ + } + try { + const lines = buffer.split("\n"); + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data:")) continue; + const payload = trimmed.slice(5).trim(); + if (!payload || payload === "[DONE]") continue; + try { + const parsed = JSON.parse(payload); + if (Array.isArray(parsed?.remainingCredits)) { + const googleCredit = parsed.remainingCredits.find((c: unknown) => { + const credit = asCreditRecord(c); + return credit?.creditType === "GOOGLE_ONE_AI"; + }) as AntigravityCreditEntry | undefined; + if (googleCredit) { + const balance = parseInt(String(googleCredit.creditAmount ?? ""), 10); + if (!isNaN(balance)) onCreditsUpdate(accountId, balance); + } + } + } catch { + /* skip malformed lines */ + } + } + } catch { + /* credits extraction is best-effort */ + } + buffer = ""; + }, + }, + { highWaterMark: 16384 }, + { highWaterMark: 16384 } + ); +} + +/** Result shape returned to callers of AntigravityExecutor.execute(). */ +export type SsePassthroughResult = { + response: Response; + url: string; + headers: Record; + transformedBody: unknown; +}; + +/** Cancel `body` when `signal` aborts, releasing the upstream connection. */ +function cancelBodyOnAbort(body: ReadableStream, signal: AbortSignal): void { + signal.addEventListener( + "abort", + () => { + body.cancel().catch(() => {}); + }, + { once: true } + ); +} + +/** + * Build the non-streaming pass-through result: tap `body` through + * createCreditsExtractionTransform and wrap it in a same-status Response so + * chatCore's non-streaming path (readNonStreamingResponseBody + + * parseNonStreamingSSEPayload) can drain and parse the Gemini SSE without + * this executor buffering the whole stream itself. + * + * If the client already disconnected (`signal.aborted`), cancels the + * upstream body immediately and returns a bare 499 instead of piping a + * cancelled body through. + */ +export function buildSsePassthroughResult( + body: ReadableStream, + upstream: { status: number; statusText: string; headers: Headers }, + accountId: string, + onCreditsUpdate: (accountId: string, balance: number) => void, + url: string, + outHeaders: Record, + transformedBody: unknown, + signal: AbortSignal | null | undefined +): SsePassthroughResult { + // Client already disconnected — skip pipe + if (signal?.aborted) { + body.cancel().catch(() => {}); + return { + response: new Response(null, { status: 499 }), + url, + headers: outHeaders, + transformedBody: null, + }; + } + // Cancel upstream body on client disconnect + if (signal) cancelBodyOnAbort(body, signal); + + const tapped = body.pipeThrough( + createCreditsExtractionTransform(accountId, onCreditsUpdate, 16 * 1024) + ); + return { + response: new Response(tapped, { + status: upstream.status, + statusText: upstream.statusText, + headers: upstream.headers, + }), + url, + headers: outHeaders, + transformedBody, + }; +} diff --git a/open-sse/handlers/chatCore/nonStreamingSse.ts b/open-sse/handlers/chatCore/nonStreamingSse.ts index 5d7e642fc6..743ba81a91 100644 --- a/open-sse/handlers/chatCore/nonStreamingSse.ts +++ b/open-sse/handlers/chatCore/nonStreamingSse.ts @@ -4,6 +4,7 @@ import { parseSSEToClaudeResponse, parseSSEToOpenAIResponse, } from "../sseParser.ts"; +import { parseSSEToGeminiResponse } from "../sseParser/geminiResponse.ts"; import { getHeaderValueCaseInsensitive } from "./headers.ts"; export function parseNonStreamingSSEPayload( @@ -20,6 +21,7 @@ export function parseNonStreamingSSEPayload( }; queueFormat(preferredFormat); + queueFormat(FORMATS.GEMINI); queueFormat(FORMATS.OPENAI_RESPONSES); queueFormat(FORMATS.CLAUDE); queueFormat(FORMATS.OPENAI); @@ -30,7 +32,9 @@ export function parseNonStreamingSSEPayload( ? parseSSEToResponsesOutput(rawBody, fallbackModel) : format === FORMATS.CLAUDE ? parseSSEToClaudeResponse(rawBody, fallbackModel) - : parseSSEToOpenAIResponse(rawBody, fallbackModel); + : format === FORMATS.GEMINI || format === FORMATS.ANTIGRAVITY + ? parseSSEToGeminiResponse(rawBody, fallbackModel) + : parseSSEToOpenAIResponse(rawBody, fallbackModel); if (parsed && typeof parsed === "object") { return { body: parsed as Record, @@ -107,6 +111,21 @@ function hasClaudeTerminalMessageDelta(parsed: unknown, eventType: string): bool return typeof stopReason === "string" ? stopReason.length > 0 : stopReason != null; } +// Non-empty finishReason is terminal. Gemini SSE payloads from +// streamGenerateContent have candidates at the top level (no +// "response" wrapper). Any non-empty string signals stream end. +function hasGeminiTerminalFinishReason(parsed: unknown): boolean { + if (!parsed || typeof parsed !== "object") return false; + // Top-level candidates (streamGenerateContent?alt=sse) + const obj = parsed as Record; + const candidates = obj.candidates as unknown[] | undefined; + if (!Array.isArray(candidates) || candidates.length === 0) return false; + const candidate = candidates[0] as Record | undefined; + if (!candidate || typeof candidate !== "object") return false; + const finishReason = candidate.finishReason; + return typeof finishReason === "string" && finishReason.length > 0; +} + function processNonStreamingSseTerminalLine( state: NonStreamingSseTerminalState, rawLine: string @@ -130,12 +149,22 @@ function processNonStreamingSseTerminalLine( // Hot-path optimization: the terminal SSE events we look for (message_stop, // response.completed, …) all carry a top-level "type" field, OR are signalled by a - // preceding `event:` line (Claude). OpenAI chat.completion chunks carry neither and - // terminate with `[DONE]` (handled above), so parsing every one of them here is pure - // waste that compounds into the CPU-runaway on large buffered responses. Skip the - // JSON.parse unless the line could actually be a typed terminal. + // preceding `event:` line (Claude). Gemini signals completion via + // "finishReason" inside response.candidates[0]. OpenAI chat.completion chunks + // carry none of these and terminate with `[DONE]` (handled above), so parsing + // every one of them here is pure waste that compounds into the CPU-runaway on + // large buffered responses. Skip the JSON.parse unless the line could actually + // be a typed terminal. if ( !data.includes('"type"') && + // NOTE: "finishReason" is a superset match -- it triggers JSON.parse on + // every Gemini chunk that happens to contain the string (e.g. partial + // candidate payloads), not just the terminal one. This is intentional: + // the extra parses are cheap compared to the CPU-runaway we'd get from + // parsing ALL chunks unconditionally on large buffered responses, and + // the superset is safe (false positives just parse a non-terminal chunk + // and fall through to `return false`). + !data.includes('"finishReason"') && !(state.currentEvent === "message_delta" && data.includes("stop_reason")) ) { return isNonStreamingSseTerminalType(state.currentEvent); @@ -148,7 +177,9 @@ function processNonStreamingSseTerminalLine( ? parsed.type : state.currentEvent; return ( - isNonStreamingSseTerminalType(eventType) || hasClaudeTerminalMessageDelta(parsed, eventType) + isNonStreamingSseTerminalType(eventType) || + hasClaudeTerminalMessageDelta(parsed, eventType) || + hasGeminiTerminalFinishReason(parsed) ); } catch { // Keep reading malformed data so the parser can report a useful upstream error. diff --git a/open-sse/handlers/sseParser/geminiResponse.ts b/open-sse/handlers/sseParser/geminiResponse.ts new file mode 100644 index 0000000000..1657cf2030 --- /dev/null +++ b/open-sse/handlers/sseParser/geminiResponse.ts @@ -0,0 +1,214 @@ +// Gemini/Antigravity buffered-SSE -> chat.completion conversion (#7408). +// Extracted verbatim from sseParser.ts (file-size cap): pure parsing, no host +// state, following the handlers submodule pattern (chatCore/, responseSanitizer/). +import { normalizeOpenAICompatibleFinishReasonString } from "../../utils/finishReason.ts"; + +type AccumulatedToolCall = { + id: string; + index: number; + type: "function"; + function: { name: string; arguments: string }; +}; + +/** Mutable accumulator threaded through one SSE payload's worth of parsing. */ +type GeminiSSEAccumulator = { + textContent: string; + finishReason: string; + usage: Record | null; + sawContent: boolean; + toolCalls: AccumulatedToolCall[]; +}; + +function stripZeroWidth(value: unknown): unknown { + if (typeof value === "string") return value.replace(/[\u200B-\u200D\uFEFF]/g, ""); + return value; +} + +/** + * Detect the `[Tool call: name]\nArguments: {...}` textual convention some + * Gemini/Antigravity models emit instead of a native functionCall part. + */ +function tryParseTextualToolCall(text: string): { name: string; args: unknown } | null { + const normalized = text.replace(/[\u200B-\u200D\uFEFF]/g, ""); + const match = normalized.match( + /^[\s\S]*?\[Tool call:\s*([^\]\n]+)\]\s*\nArguments:\s*([\s\S]+?)\s*$/ + ); + if (!match) return null; + const name = match[1]?.trim(); + const rawArgs = match[2]?.trim(); + if (!name || !rawArgs) return null; + try { + return { name, args: stripZeroWidth(JSON.parse(rawArgs)) }; + } catch { + return null; + } +} + +/** Extract the markdown shortcut some Gemini variants send (top-level or nested). */ +function extractGeminiMarkdownShortcut(parsed: Record): string | null { + if (typeof parsed.markdown === "string") return parsed.markdown; + const response = parsed.response as Record | undefined; + return typeof response?.markdown === "string" ? response.markdown : null; +} + +/** Append one candidate content part (text or textual tool call) onto the accumulator. */ +function applyCandidatePart(part: Record, acc: GeminiSSEAccumulator): void { + if (typeof part.text !== "string" || part.thought || part.thoughtSignature) return; + + const textualToolCall = tryParseTextualToolCall(part.text); + if (textualToolCall) { + acc.toolCalls.push({ + id: `${textualToolCall.name}-${Date.now()}-${acc.toolCalls.length}`, + index: acc.toolCalls.length, + type: "function", + function: { + name: textualToolCall.name, + arguments: JSON.stringify(textualToolCall.args || {}), + }, + }); + } else { + acc.textContent += part.text; + } + acc.sawContent = true; +} + +/** Walk the first candidate's content parts, if present, mutating the accumulator. */ +function applyCandidateContentParts( + candidate: Record | undefined, + acc: GeminiSSEAccumulator +): void { + const content = candidate?.content as Record | undefined; + const parts = content?.parts; + if (!Array.isArray(parts)) return; + for (const part of parts) { + applyCandidatePart(part as Record, acc); + } +} + +/** Normalize and apply the candidate's finishReason, if present. */ +function applyFinishReason( + candidate: Record | undefined, + acc: GeminiSSEAccumulator +): void { + if (!candidate?.finishReason) return; + acc.finishReason = normalizeOpenAICompatibleFinishReasonString( + String(candidate.finishReason).toLowerCase() + ); +} + +/** Extract usageMetadata into the OpenAI-shaped usage object, if present. */ +function applyUsageMetadata(parsed: Record, acc: GeminiSSEAccumulator): void { + const response = parsed.response as Record | undefined; + const um = response?.usageMetadata as Record | undefined; + if (!um) return; + acc.usage = { + prompt_tokens: um.promptTokenCount || 0, + completion_tokens: um.candidatesTokenCount || 0, + total_tokens: um.totalTokenCount || 0, + }; +} + +/** Parse one `data:` line's JSON payload and fold it into the accumulator (best-effort). */ +function applyGeminiSSEDataLine(payload: string, acc: GeminiSSEAccumulator): void { + try { + const parsed = JSON.parse(payload) as Record; + + const markdown = extractGeminiMarkdownShortcut(parsed); + if (markdown) { + acc.textContent += markdown; + acc.sawContent = true; + } + + const response = parsed.response as Record | undefined; + const candidates = response?.candidates; + const candidate = Array.isArray(candidates) + ? (candidates[0] as Record | undefined) + : undefined; + + applyCandidateContentParts(candidate, acc); + applyFinishReason(candidate, acc); + applyUsageMetadata(parsed, acc); + } catch { + // Ignore malformed lines + } +} + +/** Assemble the final non-streaming chat.completion payload from the accumulator. */ +function buildChatCompletionFromAccumulator( + acc: GeminiSSEAccumulator, + fallbackModel: string +): Record { + const message: Record = { + role: "assistant", + content: acc.textContent || null, + }; + + let finishReason = acc.finishReason; + if (acc.toolCalls.length > 0) { + message.tool_calls = acc.toolCalls; + finishReason = "tool_calls"; + } + + const result: Record = { + id: `chatcmpl-${Date.now()}`, + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: fallbackModel || "unknown", + choices: [ + { + index: 0, + message, + finish_reason: finishReason, + }, + ], + }; + + if (acc.usage) { + result.usage = acc.usage; + } + + return result; +} + +/** + * Convert Gemini/Antigravity SSE chunks into a single non-streaming OpenAI + * chat.completion JSON response. Gemini SSE carries payloads like: + * + * data: {"markdown":"...chunk..."} + * data: {"response":{"candidates":[{"content":{"parts":[{"text":"..."}]},"finishReason":"STOP"}],"usageMetadata":{...}}} + * data: {"remainingCredits":[...]} + * + * Reuses the same parsing logic as processAntigravitySSEPayload() in sseCollect.ts + * so that format conversion is functionally equivalent to the previous + * collectStreamToResponse() approach. Intentional differences: + * - remainingCredits is NOT embedded into the result (handled separately + * by the credits-extraction TransformStream in antigravity.ts). + * - The synthetic `id` uses `chatcmpl-${Date.now()}` (no UUID suffix) + * because this path runs once per response, not per chunk. + */ +export function parseSSEToGeminiResponse( + rawSSE: string, + fallbackModel: string +): Record | null { + const lines = String(rawSSE || "").split("\n"); + const acc: GeminiSSEAccumulator = { + textContent: "", + finishReason: "stop", + usage: null, + sawContent: false, + toolCalls: [], + }; + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data:")) continue; + const payload = trimmed.slice(5).trim(); + if (!payload || payload === "[DONE]") continue; + + applyGeminiSSEDataLine(payload, acc); + } + + if (!acc.sawContent && acc.toolCalls.length === 0) return null; + + return buildChatCompletionFromAccumulator(acc, fallbackModel); +} diff --git a/tests/unit/antigravity-streaming-passthrough.test.ts b/tests/unit/antigravity-streaming-passthrough.test.ts new file mode 100644 index 0000000000..23715cbae3 --- /dev/null +++ b/tests/unit/antigravity-streaming-passthrough.test.ts @@ -0,0 +1,202 @@ +// Antigravity streaming-passthrough behavior (#7408): the non-streaming drain +// path (raw SSE + chatCore-side parse) and the credits-extraction pass-through +// TransformStream. Moved verbatim from tests/unit/executor-antigravity.test.ts +// (frozen file-size cap) — same tests, same asserts. +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import { + AntigravityExecutor, + createCreditsExtractionTransform, +} from "../../open-sse/executors/antigravity.ts"; +import { parseSSEToGeminiResponse } from "../../open-sse/handlers/sseParser/geminiResponse.ts"; +import { + clearAntigravityVersionCache, + seedAntigravityVersionCache, +} from "../../open-sse/services/antigravityVersion.ts"; + +type ChatCompletionPayload = { + object?: string; + choices: Array<{ + message: { content: string }; + finish_reason: string; + }>; + usage?: { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + }; +}; + +test.afterEach(() => { + clearAntigravityVersionCache(); +}); + +test("AntigravityExecutor.execute auto-retries short 429 responses and collects SSE for non-stream clients", async () => { + const executor = new AntigravityExecutor(); + const originalFetch = globalThis.fetch; + const originalSetTimeout = globalThis.setTimeout; + const calls = []; + seedAntigravityVersionCache("2026.04.17-test"); + + globalThis.fetch = async (url) => { + calls.push(String(url)); + + if (calls.length === 1) { + return new Response(JSON.stringify({ error: { message: "rate limited" } }), { + status: 429, + headers: { "Content-Type": "application/json" }, + }); + } + + return new Response( + [ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hello "}]},"finishReason":"STOP"}]}}\n\n', + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"again"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3,"totalTokenCount":5}}}\n\n', + ].join(""), + { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + } + ); + }; + globalThis.setTimeout = ((callback) => { + (callback as () => void)(); + return 0; + }) as typeof setTimeout; + + try { + const result = await executor.execute({ + model: "antigravity/gemini-2.5-flash", + body: { request: { contents: [] } }, + stream: false, + credentials: { accessToken: "token", projectId: "project-1" }, + log: { debug() {}, warn() {} }, + }); + // Non-streaming now returns raw SSE; parse it the way chatCore would. + const rawSSE = await result.response.text(); + const parsed = parseSSEToGeminiResponse(rawSSE, "antigravity/gemini-2.5-flash"); + assert.ok(parsed, "parseSSEToGeminiResponse should parse the SSE"); + const payload = parsed as ChatCompletionPayload; + + assert.equal(calls.length, 2); + assert.equal(result.response.status, 200); + assert.equal(payload.choices[0].message.content, "Hello again"); + assert.deepEqual(payload.usage, { + prompt_tokens: 2, + completion_tokens: 3, + total_tokens: 5, + }); + } finally { + globalThis.fetch = originalFetch; + globalThis.setTimeout = originalSetTimeout; + } +}); + +// --------------------------------------------------------------------------- +// createCreditsExtractionTransform -- credits extraction with buffer cap +// --------------------------------------------------------------------------- + +test("createCreditsExtractionTransform extracts remainingCredits from SSE data", async () => { + const encoder = new TextEncoder(); + const sseData = [ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}]}}\n\n', + 'data: {"remainingCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"42"}]}\n\n', + ].join(""); + + const transform = createCreditsExtractionTransform("test-account"); + const readable = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sseData)); + controller.close(); + }, + }); + + // Consume the stream through the transform + const output = readable.pipeThrough(transform); + const reader = output.getReader(); + const chunks: Uint8Array[] = []; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + } + + // Data must pass through unmodified + const collected = new TextDecoder().decode( + new Uint8Array( + chunks.reduce((acc, c) => acc + c.length, 0) > 0 ? Buffer.concat(chunks) : new Uint8Array(0) + ) + ); + assert.ok(collected.includes("hello")); + assert.ok(collected.includes("remainingCredits")); +}); + +test("createCreditsExtractionTransform with buffer cap truncates large buffers", async () => { + const encoder = new TextEncoder(); + // Build a payload larger than 1KB + const largeText = "x".repeat(2000); + const ssePayload = JSON.stringify({ + response: { + candidates: [{ content: { parts: [{ text: largeText }] } }], + }, + }); + const sseLine = `data: ${ssePayload}\n\n`; + // Append a credits line at the end + const creditsLine = + 'data: {"remainingCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"99"}]}\n\n'; + const fullData = sseLine + creditsLine; + + // Use a 512-byte buffer cap -- the large text line should be discarded + const transform = createCreditsExtractionTransform("test-account", 512); + const readable = new ReadableStream({ + start(controller) { + // Send in small chunks to exercise the sliding-window logic + const encoded = encoder.encode(fullData); + const chunkSize = 256; + for (let i = 0; i < encoded.length; i += chunkSize) { + controller.enqueue(encoded.slice(i, i + chunkSize)); + } + controller.close(); + }, + }); + + const output = readable.pipeThrough(transform); + const reader = output.getReader(); + while (true) { + const { done } = await reader.read(); + if (done) break; + } + + // The transform should not throw -- buffer cap just limits what the + // flush handler can see. If the credits line was within the last 512 + // bytes it will be found; otherwise it's a graceful no-op. + // Either way, no crash or OOM. + assert.ok(true); +}); + +test("createCreditsExtractionTransform handles malformed SSE gracefully", async () => { + const encoder = new TextEncoder(); + const badData = "not valid sse\ndata: {broken json\n\ndata: [DONE]\n\n"; + + const transform = createCreditsExtractionTransform("test-account"); + const readable = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(badData)); + controller.close(); + }, + }); + + const output = readable.pipeThrough(transform); + const reader = output.getReader(); + const chunks: Uint8Array[] = []; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + } + + // Data passes through unmodified, no crash on malformed input + const collected = new TextDecoder().decode(Buffer.concat(chunks)); + assert.ok(collected.includes("not valid sse")); +}); diff --git a/tests/unit/executor-antigravity.test.ts b/tests/unit/executor-antigravity.test.ts index 465af7a718..1cdc87e6b6 100644 --- a/tests/unit/executor-antigravity.test.ts +++ b/tests/unit/executor-antigravity.test.ts @@ -631,62 +631,9 @@ test("AntigravityExecutor.refreshCredentials refreshes Google OAuth tokens", asy } }); -test("AntigravityExecutor.execute auto-retries short 429 responses and collects SSE for non-stream clients", async () => { - const executor = new AntigravityExecutor(); - const originalFetch = globalThis.fetch; - const originalSetTimeout = globalThis.setTimeout; - const calls = []; - seedAntigravityVersionCache("2026.04.17-test"); - - globalThis.fetch = async (url) => { - calls.push(String(url)); - - if (calls.length === 1) { - return new Response(JSON.stringify({ error: { message: "rate limited" } }), { - status: 429, - headers: { "Content-Type": "application/json" }, - }); - } - - return new Response( - [ - 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hello "}]},"finishReason":"STOP"}]}}\n\n', - 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"again"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3,"totalTokenCount":5}}}\n\n', - ].join(""), - { - status: 200, - headers: { "Content-Type": "text/event-stream" }, - } - ); - }; - globalThis.setTimeout = ((callback) => { - (callback as () => void)(); - return 0; - }) as typeof setTimeout; - - try { - const result = await executor.execute({ - model: "antigravity/gemini-2.5-flash", - body: { request: { contents: [] } }, - stream: false, - credentials: { accessToken: "token", projectId: "project-1" }, - log: { debug() {}, warn() {} }, - }); - const payload = (await result.response.json()) as ChatCompletionPayload; - - assert.equal(calls.length, 2); - assert.equal(result.response.status, 200); - assert.equal(payload.choices[0].message.content, "Hello again"); - assert.deepEqual(payload.usage, { - prompt_tokens: 2, - completion_tokens: 3, - total_tokens: 5, - }); - } finally { - globalThis.fetch = originalFetch; - globalThis.setTimeout = originalSetTimeout; - } -}); +// The non-streaming passthrough drain test ("auto-retries short 429 responses and +// collects SSE for non-stream clients") lives in +// tests/unit/antigravity-streaming-passthrough.test.ts with the other passthrough tests. test("AntigravityExecutor.execute embeds retryAfterMs when the upstream asks for a long wait", async () => { const executor = new AntigravityExecutor(); diff --git a/tests/unit/sse-parser.test.ts b/tests/unit/sse-parser.test.ts index 7511834df0..36fcd7deda 100644 --- a/tests/unit/sse-parser.test.ts +++ b/tests/unit/sse-parser.test.ts @@ -3,6 +3,8 @@ import assert from "node:assert/strict"; const { parseSSEToOpenAIResponse, parseSSEToClaudeResponse, parseSSEToResponsesOutput } = await import("../../open-sse/handlers/sseParser.ts"); +const { parseSSEToGeminiResponse } = + await import("../../open-sse/handlers/sseParser/geminiResponse.ts"); test("parseSSEToOpenAIResponse parses a single SSE event with a done marker", () => { const rawSSE = [ @@ -103,12 +105,12 @@ test("parseSSEToClaudeResponse parses text, thinking, tool_use, and usage events assert.equal(parsed.id, "msg_1"); assert.equal(parsed.model, "claude-3-5-sonnet"); - assert.equal((parsed.content[0] as any).type, "thinking"); - assert.equal((parsed as any).content[0].thinking, "step 1"); - assert.equal((parsed as any).content[0].signature, "sig-1"); - (assert as any).equal((parsed.content[1] as any).text, "Hello"); - assert.equal((parsed.content[2] as any).type, "tool_use"); - (assert as any).deepEqual((parsed.content[2] as any).input, { q: "docs" }); + assert.equal((parsed.content[0] as { type: string }).type, "thinking"); + assert.equal((parsed.content[0] as { thinking: string }).thinking, "step 1"); + assert.equal((parsed.content[0] as { signature: string }).signature, "sig-1"); + assert.equal((parsed.content[1] as { text: string }).text, "Hello"); + assert.equal((parsed.content[2] as { type: string }).type, "tool_use"); + assert.deepEqual((parsed.content[2] as { input: unknown }).input, { q: "docs" }); assert.equal(parsed.stop_reason, "tool_use"); assert.equal(parsed.stop_sequence, "END"); assert.deepEqual(parsed.usage, { input_tokens: 10, output_tokens: 4 }); @@ -130,7 +132,7 @@ test("parseSSEToClaudeResponse tolerates event-only types and missing blank sepa assert.equal(parsed.id, "msg_event_fallback"); assert.equal(parsed.model, "claude-sonnet-4-6"); - assert.equal((parsed.content[0] as any).text, "event fallback ok"); + assert.equal((parsed.content[0] as { text: string }).text, "event fallback ok"); assert.deepEqual(parsed.usage, { input_tokens: 3, output_tokens: 2 }); }); @@ -152,9 +154,9 @@ test("parseSSEToClaudeResponse merges signature_delta into an existing thinking const parsed = parseSSEToClaudeResponse(rawSSE, "fallback-model"); - assert.equal((parsed.content[0] as any).type, "thinking"); - assert.equal((parsed.content[0] as any).thinking, "first second"); - assert.equal((parsed.content[0] as any).signature, "sig-1"); + assert.equal((parsed.content[0] as { type: string }).type, "thinking"); + assert.equal((parsed.content[0] as { thinking: string }).thinking, "first second"); + assert.equal((parsed.content[0] as { signature: string }).signature, "sig-1"); }); test("parseSSEToClaudeResponse preserves signature_delta when it arrives before thinking_delta", () => { @@ -172,9 +174,9 @@ test("parseSSEToClaudeResponse preserves signature_delta when it arrives before const parsed = parseSSEToClaudeResponse(rawSSE, "fallback-model"); - assert.equal((parsed.content[0] as any).type, "thinking"); - assert.equal((parsed.content[0] as any).thinking, "later thinking"); - assert.equal((parsed.content[0] as any).signature, "sig-before"); + assert.equal((parsed.content[0] as { type: string }).type, "thinking"); + assert.equal((parsed.content[0] as { thinking: string }).thinking, "later thinking"); + assert.equal((parsed.content[0] as { signature: string }).signature, "sig-before"); }); test("parseSSEToClaudeResponse ignores malformed payloads and returns null when nothing valid remains", () => { @@ -334,3 +336,98 @@ test("parseSSEToOpenAIResponse deduplicates repeated tool call snapshots", () => assert.equal(toolCall.function.arguments, args); assert.equal(JSON.parse(toolCall.function.arguments).command, "find /tmp -name test.txt"); }); + +// --------------------------------------------------------------------------- +// parseSSEToGeminiResponse +// --------------------------------------------------------------------------- + +test("parseSSEToGeminiResponse extracts text content from candidate parts", () => { + const rawSSE = [ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hello "}]}}]}}', + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"world"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3,"totalTokenCount":8}}}', + ].join("\n"); + + const parsed = parseSSEToGeminiResponse(rawSSE, "gemini-2.5-flash"); + + assert.ok(parsed); + assert.equal(parsed.object, "chat.completion"); + assert.equal(parsed.choices[0].message.content, "Hello world"); + assert.equal(parsed.choices[0].finish_reason, "stop"); + assert.deepEqual(parsed.usage, { + prompt_tokens: 5, + completion_tokens: 3, + total_tokens: 8, + }); +}); + +test("parseSSEToGeminiResponse handles markdown shortcut format", () => { + const rawSSE = [ + 'data: {"markdown":"Hello "}', + 'data: {"markdown":"world"}', + 'data: {"response":{"candidates":[{"finishReason":"STOP"}]}}', + ].join("\n"); + + const parsed = parseSSEToGeminiResponse(rawSSE, "gemini-2.5-flash"); + + assert.ok(parsed); + assert.equal(parsed.choices[0].message.content, "Hello world"); + assert.equal(parsed.choices[0].finish_reason, "stop"); +}); + +test("parseSSEToGeminiResponse extracts tool calls from textual format", () => { + const rawSSE = [ + `data: ${JSON.stringify({ + response: { + candidates: [ + { + content: { + parts: [ + { + text: '[Tool call: search_files]\nArguments: {"path":"/tmp"}', + }, + ], + }, + finishReason: "STOP", + }, + ], + }, + })}`, + ].join("\n"); + + const parsed = parseSSEToGeminiResponse(rawSSE, "gemini-3.5-flash-low"); + + assert.ok(parsed); + assert.equal(parsed.choices[0].finish_reason, "tool_calls"); + const toolCalls = parsed.choices[0].message.tool_calls; + assert.equal(toolCalls.length, 1); + assert.equal(toolCalls[0].function.name, "search_files"); + assert.deepEqual(JSON.parse(toolCalls[0].function.arguments), { path: "/tmp" }); +}); + +test("parseSSEToGeminiResponse returns null for empty or non-content SSE", () => { + assert.equal(parseSSEToGeminiResponse("", "model"), null); + assert.equal(parseSSEToGeminiResponse("data: [DONE]\n", "model"), null); + assert.equal(parseSSEToGeminiResponse("event: ping\n", "model"), null); +}); + +test("parseSSEToGeminiResponse ignores thought/thoughtSignature parts", () => { + const rawSSE = [ + `data: ${JSON.stringify({ + response: { + candidates: [ + { + content: { + parts: [{ text: "internal reasoning", thought: true }, { text: "visible answer" }], + }, + finishReason: "STOP", + }, + ], + }, + })}`, + ].join("\n"); + + const parsed = parseSSEToGeminiResponse(rawSSE, "model"); + + assert.ok(parsed); + assert.equal(parsed.choices[0].message.content, "visible answer"); +});