diff --git a/src/adapters/base.ts b/src/adapters/base.ts index 8789a03463..cb067fe659 100644 --- a/src/adapters/base.ts +++ b/src/adapters/base.ts @@ -82,4 +82,6 @@ export interface AdapterFetchContext { returnRawErrors?: boolean; /** Whether the upstream response will be consumed as a stream; adapters may select low-latency transport settings. */ stream?: boolean; + /** Custom fetch executor to use for physical upstream network requests (defaults to globalThis.fetch). */ + executor?: typeof globalThis.fetch; } diff --git a/src/adapters/google-http.ts b/src/adapters/google-http.ts index de849cde3c..f7b90de87e 100644 --- a/src/adapters/google-http.ts +++ b/src/adapters/google-http.ts @@ -14,6 +14,11 @@ const GOOGLE_RETRY_ATTEMPTS = 3; const GOOGLE_RETRY_BASE_MS = 250; const GOOGLE_RETRY_MAX_MS = 2_000; +export interface GoogleRetryOptions { + /** Repair-and-replay structurally invalid 400 bodies (Vertex/Antigravity behavior). */ + repairInvalid400?: boolean; +} + async function normalizeFinalGoogleError(label: string, res: Response, signal?: AbortSignal): Promise { return normalizeUpstreamHttpErrorResponse(res, { signal, @@ -22,13 +27,21 @@ async function normalizeFinalGoogleError(label: string, res: Response, signal?: } /** - * Fetch a Google-family upstream (Vertex / Antigravity) with Kiro-style hardening: per-attempt - * timeout (`AbortSignal.any([parent, timeout])`), bounded retry on transient status / network - * errors, `Retry-After` honoring, jittered exponential backoff, and a classified + redacted final - * error body. `label` is the provider-facing prefix used in error messages. + * Fetch a Google-family upstream with Kiro-style hardening: per-attempt timeout + * (`AbortSignal.any([parent, timeout])`), bounded retry on transient status / network errors, + * `Retry-After` honoring, jittered exponential backoff, and (unless raw mode is used) a + * classified + redacted final error body. `label` is the provider-facing prefix used in error + * messages. */ -export async function fetchGoogleWithRetry(label: string, request: AdapterRequest, ctx: AdapterFetchContext = {}): Promise { +export async function fetchGoogleWithRetry( + label: string, + request: AdapterRequest, + ctx: AdapterFetchContext = {}, + opts: GoogleRetryOptions = {}, +): Promise { + const repairInvalid400 = opts.repairInvalid400 ?? true; const timeoutMs = ctx.timeoutMs ?? 200_000; + const executor = ctx.executor ?? globalThis.fetch; let lastError: unknown; let activeRequest = request; let compatibilityReplayUsed = false; @@ -39,8 +52,8 @@ export async function fetchGoogleWithRetry(label: string, request: AdapterReques method: activeRequest.method, headers: activeRequest.headers, body: activeRequest.body, - }, timeoutMs, ctx.abortSignal, ctx.stream); - if (res.status === 400 && !compatibilityReplayUsed) { + }, timeoutMs, ctx.abortSignal, ctx.stream, executor); + if (res.status === 400 && repairInvalid400 && !compatibilityReplayUsed) { let payloadText = ""; try { payloadText = await readDisplaySafeErrorPayloadText(res.clone(), ctx.abortSignal); @@ -61,10 +74,11 @@ export async function fetchGoogleWithRetry(label: string, request: AdapterReques } // A 429 may be a transient rate limit (retry) or hard quota exhaustion (do NOT retry — // it won't recover for hours and burns retries). Peek the body to tell them apart. - if (res.status === 429 && !ctx.returnRawErrors) { - const peek = await readDisplaySafeErrorPayloadText(res, ctx.abortSignal); + if (res.status === 429) { + const peekTarget = ctx.returnRawErrors ? res.clone() : res; + const peek = await readDisplaySafeErrorPayloadText(peekTarget, ctx.abortSignal); if (isQuotaExhaustedBody(peek)) { - return normalizeUpstreamHttpErrorResponse(res, { + return ctx.returnRawErrors ? res : normalizeUpstreamHttpErrorResponse(res, { signal: ctx.abortSignal, formatMessage: payloadText => safeGoogleHttpErrorMessage(label, res.status, payloadText || peek), }); @@ -89,6 +103,20 @@ export async function fetchGoogleWithRetry(label: string, request: AdapterReques throw lastError ?? new Error(`${label} fetch failed`); } +/** + * AI Studio direct (`generativelanguage.googleapis.com`) retry wrapper. + * + * Direct requests keep the default server error surface — the raw `Provider error : + * ` text the shared Responses path formats — and keep single-shot 400 semantics (no + * request-shape compatibility replay). The wrapper exists for the failure mode observed in + * production: AI Studio's transient `503 UNAVAILABLE` "model is currently experiencing high + * demand" spikes, plus plain rate-limit 429s, both of which previously failed immediately + * because the default server fetch path only retries connection resets. + */ +export function fetchDirectGeminiWithRetry(request: AdapterRequest, ctx: AdapterFetchContext = {}): Promise { + return fetchGoogleWithRetry("Gemini", request, { ...ctx, returnRawErrors: true }, { repairInvalid400: false }); +} + /** Vertex AI retry wrapper. */ export function fetchVertexWithRetry(request: AdapterRequest, ctx: AdapterFetchContext = {}): Promise { return fetchGoogleWithRetry("Vertex AI", request, ctx); diff --git a/src/adapters/google.ts b/src/adapters/google.ts index 72914d9559..8b8d577dee 100644 --- a/src/adapters/google.ts +++ b/src/adapters/google.ts @@ -377,8 +377,10 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte return { name: "google", - // Vertex + Antigravity get Kiro-style retry/timeout + classified, redacted errors. AI-Studio - // Gemini keeps the default server fetch path (fetchResponse stays undefined so server.ts falls back). + // Vertex + Antigravity get Kiro-style retry/timeout + classified, redacted errors. + // Direct AI-Studio uses the canonical server transport (fetchWithTransientRetry), which + // retries transient 5xx responses through providerFetch while preserving multi-key pool + // 429 rotation and raw error formatting. ...(provider.googleMode === "vertex" || provider.googleMode === "cloud-code-assist" ? { fetchResponse: (request: AdapterRequest, ctx?: AdapterFetchContext): Promise => diff --git a/src/images/loop.ts b/src/images/loop.ts index 65a71276c6..1699906d94 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -508,6 +508,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { const attemptTimeout = clearableDeadline(timeoutMs, abortSignal); const headers = new Headers(init.headers); @@ -222,7 +223,7 @@ export async function fetchWithAttemptDeadline( headers.set("accept-encoding", "identity"); } try { - return await fetch(url, { + return await executor(url, { ...init, headers, signal: attemptTimeout.signal, diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index a5ad0191fc..55b5797327 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -3583,9 +3583,13 @@ async function handleResponsesInner( abortSignal: upstream.signal, timeoutMs: connectMs, stream: parsed.stream, + executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, + modelId: route.modelId, + }), }); } else { - upstreamResponse = await fetchWithResetRetry( + upstreamResponse = await fetchWithTransientRetry( recovery => { noteAttemptSend(logCtx.activeAttempt, inputTokenEstimate, recovery); return fetchWithHeaderTimeout(builtInitialRequest.url, applyUpstreamRecoveryInit({ @@ -3673,7 +3677,15 @@ async function handleResponsesInner( try { if (activeAdapter.fetchResponse) { await waitForProviderRequestSlot(route.providerName, route.provider, route.modelId, upstream.signal); - return await activeAdapter.fetchResponse(retryRequest, { abortSignal: upstream.signal, timeoutMs: connectMs, stream: parsed.stream }); + return await activeAdapter.fetchResponse(retryRequest, { + abortSignal: upstream.signal, + timeoutMs: connectMs, + stream: parsed.stream, + executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, + modelId: route.modelId, + }), + }); } return await fetchWithHeaderTimeout(retryRequest.url, { method: retryRequest.method, headers: retryRequest.headers, body: retryRequest.body, @@ -3999,9 +4011,13 @@ async function handleResponsesInner( abortSignal: upstream.signal, timeoutMs: connectMs, stream: nextParsed.stream, + executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, + modelId: nextParsed.modelId, + }), }); } - return await fetchWithResetRetry( + return await fetchWithTransientRetry( recovery => { noteAttemptSend(logCtx.activeAttempt, continuationEstimate, recovery ?? replayKind); return fetchWithHeaderTimeout( diff --git a/tests/google-vertex-http.test.ts b/tests/google-vertex-http.test.ts index e640f6fabe..496fb56a1a 100644 --- a/tests/google-vertex-http.test.ts +++ b/tests/google-vertex-http.test.ts @@ -1,6 +1,6 @@ import { afterEach, describe, expect, test } from "bun:test"; import type { AdapterRequest } from "../src/adapters/base"; -import { fetchAntigravityWithRetry, fetchVertexWithRetry } from "../src/adapters/google-http"; +import { fetchAntigravityWithRetry, fetchDirectGeminiWithRetry, fetchVertexWithRetry } from "../src/adapters/google-http"; import { safeVertexHttpErrorMessage, retryableGoogleStatus } from "../src/adapters/google-errors"; const realFetch = globalThis.fetch; @@ -171,16 +171,16 @@ describe("vertex retry fetch", () => { expect(mock.calls).toHaveLength(1); }); - test("raw quota errors keep bounded retry counts without body peeking", async () => { + test("raw mode does NOT retry a quota-exhausted 429 and preserves raw response", async () => { const raw = vertexError(429, "RESOURCE_EXHAUSTED", "Quota exceeded for billing"); const mock = mockFetch([ - new Response(raw, { status: 429, headers: { "Retry-After": "0" } }), - new Response(raw, { status: 429, headers: { "Retry-After": "0" } }), - new Response(raw, { status: 429, headers: { "Retry-After": "0", "x-final": "yes" } }), + new Response(raw, { status: 429, headers: { "Retry-After": "0", "x-raw-quota": "1" } }), + new Response("ok", { status: 200 }), ]); const res = await fetchVertexWithRetry(request, { timeoutMs: 5_000, returnRawErrors: true }); - expect(mock.calls).toHaveLength(3); - expect(res.headers.get("x-final")).toBe("yes"); + expect(mock.calls).toHaveLength(1); + expect(res.status).toBe(429); + expect(res.headers.get("x-raw-quota")).toBe("1"); expect(await res.text()).toBe(raw); }); @@ -212,6 +212,69 @@ describe("vertex retry fetch", () => { controller.abort(); await expect(p).rejects.toBeDefined(); }); + + test("direct AI Studio retries a transient 503 then returns the successful response", async () => { + const mock = mockFetch([ + new Response(vertexError(503, "UNAVAILABLE", "This model is currently experiencing high demand."), { status: 503, headers: { "Retry-After": "0" } }), + new Response("ok", { status: 200 }), + ]); + const res = await fetchDirectGeminiWithRetry(request, { timeoutMs: 5_000 }); + expect(res.status).toBe(200); + expect(await res.text()).toBe("ok"); + expect(mock.calls).toHaveLength(2); + }); + + test("direct AI Studio keeps the raw final error body and does not replay repaired 400s", async () => { + const raw = vertexError(400, "INVALID_ARGUMENT", "tools.0.custom.input_schema: JSON schema is invalid"); + const mock = mockFetch([ + new Response(raw, { status: 400, headers: { "x-provider-error": "raw" } }), + new Response("ok", { status: 200 }), + ]); + const res = await fetchDirectGeminiWithRetry(request, { timeoutMs: 5_000 }); + expect(res.status).toBe(400); + expect(res.headers.get("x-provider-error")).toBe("raw"); + expect(await res.text()).toBe(raw); + expect(mock.calls).toHaveLength(1); + }); + + test("direct AI Studio does NOT retry a quota-exhausted 429 (single attempt, raw body returned)", async () => { + const raw = vertexError(429, "RESOURCE_EXHAUSTED", "Quota exceeded for quota metric 'Generate Content API requests'"); + const mock = mockFetch([ + new Response(raw, { status: 429, headers: { "Retry-After": "0", "x-direct-raw": "quota" } }), + new Response("ok", { status: 200 }), + ]); + const res = await fetchDirectGeminiWithRetry(request, { timeoutMs: 5_000 }); + expect(mock.calls).toHaveLength(1); + expect(res.status).toBe(429); + expect(res.headers.get("x-direct-raw")).toBe("quota"); + expect(await res.text()).toBe(raw); + }); + + test("direct AI Studio retries transient rate-limit 429s (bounded, raw body on exhaustion)", async () => { + const raw = vertexError(429, "RESOURCE_EXHAUSTED", "rate limit, try again"); + const mock = mockFetch([ + new Response(raw, { status: 429, headers: { "Retry-After": "0" } }), + new Response(raw, { status: 429, headers: { "Retry-After": "0" } }), + new Response(raw, { status: 429, headers: { "Retry-After": "0", "x-final": "yes" } }), + ]); + const res = await fetchDirectGeminiWithRetry(request, { timeoutMs: 5_000 }); + expect(mock.calls).toHaveLength(3); + expect(res.headers.get("x-final")).toBe("yes"); + expect(await res.text()).toBe(raw); + }); + + test("fetchGoogleWithRetry routes physical attempts through ctx.executor when provided", async () => { + const executorCalls: RequestInit[] = []; + const customExecutor: typeof fetch = (async (_url: string | URL | Request, init?: RequestInit) => { + executorCalls.push(init ?? {}); + return new Response("executor-ok", { status: 200 }); + }) as typeof fetch; + + const res = await fetchVertexWithRetry(request, { timeoutMs: 5_000, executor: customExecutor }); + expect(res.status).toBe(200); + expect(await res.text()).toBe("executor-ok"); + expect(executorCalls).toHaveLength(1); + }); }); describe("safeVertexHttpErrorMessage classification + redaction", () => { @@ -249,12 +312,15 @@ describe("safeVertexHttpErrorMessage classification + redaction", () => { }); describe("adapter fetchResponse wiring", () => { - test("vertex adapter exposes fetchResponse; ai-studio does not", async () => { + test("vertex and antigravity adapters expose fetchResponse; ai-studio direct delegates to canonical server transport", async () => { const { createGoogleAdapter } = await import("../src/adapters/google"); const vertex = createGoogleAdapter({ adapter: "google", baseUrl: "https://aiplatform.googleapis.com", googleMode: "vertex" } as never); const aistudio = createGoogleAdapter({ adapter: "google", baseUrl: "https://generativelanguage.googleapis.com", apiKey: "k" } as never); + const antigravity = createGoogleAdapter({ adapter: "google", baseUrl: "https://daily-cloudcode-pa.googleapis.com", googleMode: "cloud-code-assist" } as never); expect(typeof vertex.fetchResponse).toBe("function"); expect(typeof vertex.formatErrorBody).toBe("function"); + expect(typeof antigravity.fetchResponse).toBe("function"); + expect(typeof antigravity.formatErrorBody).toBe("function"); expect(aistudio.fetchResponse).toBeUndefined(); expect(aistudio.formatErrorBody).toBeUndefined(); }); diff --git a/tests/request-pacing.test.ts b/tests/request-pacing.test.ts index 4a2dedc60a..7df4c59c3a 100644 --- a/tests/request-pacing.test.ts +++ b/tests/request-pacing.test.ts @@ -259,4 +259,24 @@ describe("provider request pacing queue", () => { const second = await fetchWithHeaderTimeout("https://example.test/v1/chat/completions", {}, new AbortController().signal, 50, false, executor); expect(second.status).toBe(200); }); + + test("Google AI Studio providerFetch paces each attempt through waitForPacing", async () => { + let pacingWaited = 0; + const configured: OcxProviderConfig = { + adapter: "google", + baseUrl: "https://generativelanguage.googleapis.com", + apiKey: "key", + requestPacing: { enabled: true, minIntervalMs: 50 }, + fetch: (async () => new Response("ok")) as typeof fetch, + }; + const executor = providerFetch(configured, undefined, { providerName: "google-direct", modelId: "gemini-2.5-flash" }); + const originalWaitForPacing = executor.waitForPacing; + executor.waitForPacing = async (signal) => { + pacingWaited++; + await originalWaitForPacing?.(signal); + }; + const res = await fetchWithHeaderTimeout("https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent", {}, new AbortController().signal, 500, false, executor); + expect(res.status).toBe(200); + expect(pacingWaited).toBe(1); + }); }); diff --git a/tests/server-key-failover-e2e.test.ts b/tests/server-key-failover-e2e.test.ts index be0d56df7d..030ad9bbb5 100644 --- a/tests/server-key-failover-e2e.test.ts +++ b/tests/server-key-failover-e2e.test.ts @@ -465,4 +465,65 @@ describe("server 429 key failover (end-to-end)", () => { await server.stop(true); } }); + + test("Google AI Studio apiKeyPool rotates on 429 without redundant single-key retries (A sent once, B sent once)", async () => { + const originalFetch = globalThis.fetch; + const seenKeys: string[] = []; + globalThis.fetch = (async (input, init) => { + const url = input instanceof Request ? input.url : String(input); + if (url.includes("generativelanguage.googleapis.com")) { + const headers = new Headers(init?.headers); + const key = headers.get("x-goog-api-key") ?? ""; + seenKeys.push(key); + if (seenKeys.length === 1) { + return new Response(JSON.stringify({ + error: { code: 429, message: "Rate limit exceeded", status: "RESOURCE_EXHAUSTED" }, + }), { + status: 429, + headers: { "retry-after": "30", "content-type": "application/json" }, + }); + } + return new Response(JSON.stringify({ + candidates: [{ + content: { role: "model", parts: [{ text: "response from key B" }] }, + finishReason: "STOP", + }], + usageMetadata: { promptTokenCount: 3, candidatesTokenCount: 4, totalTokenCount: 7 }, + }), { headers: { "content-type": "application/json" } }); + } + return originalFetch(input, init); + }) as typeof fetch; + + const config: OcxConfig = { + port: 0, hostname: "127.0.0.1", defaultProvider: "google-direct", + providers: { + "google-direct": { + adapter: "google", + baseUrl: "https://generativelanguage.googleapis.com", + authMode: "key", + apiKey: "key-alpha-111", + apiKeyPool: [ + { id: "k1", key: "key-alpha-111", addedAt: 1 }, + { id: "k2", key: "key-beta-222", addedAt: 2 }, + ], + }, + }, + } as OcxConfig; + saveConfig(config); + const server = startServer(0); + try { + const res = await fetch(new URL("/v1/responses", server.url), { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "google-direct/gemini-2.5-flash", input: "hi", stream: false }), + }); + expect(res.status).toBe(200); + const json = await res.json() as { output?: Array<{ content?: Array<{ text?: string }> }> }; + expect(json.output?.[0]?.content?.[0]?.text).toBe("response from key B"); + expect(seenKeys).toEqual(["key-alpha-111", "key-beta-222"]); + } finally { + await server.stop(true); + globalThis.fetch = originalFetch; + } + }); }); diff --git a/tests/upstream-http-version.test.ts b/tests/upstream-http-version.test.ts index d23bce3f50..34d555fa0b 100644 --- a/tests/upstream-http-version.test.ts +++ b/tests/upstream-http-version.test.ts @@ -116,4 +116,22 @@ describe("providerFetch upstreamHttpVersion propagation", () => { await fetcher(HTTPS_URL); expect((seen.init as RequestInit & { protocol?: string })?.protocol).toBe("http1.1"); }); + + test("Google AI Studio providerFetch attaches pinned protocol to all attempts", async () => { + const seen: RequestInit[] = []; + const fetcher = providerFetch({ + adapter: "google", + baseUrl: "https://generativelanguage.googleapis.com", + apiKey: "ai-key", + upstreamHttpVersion: "http1.1", + fetch: (async (_url, init) => { + seen.push(init ?? {}); + return new Response("ok"); + }) as typeof fetch, + }); + + await fetcher("https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"); + expect(seen).toHaveLength(1); + expect((seen[0] as RequestInit & { protocol?: string }).protocol).toBe("http1.1"); + }); });