From 73db117583b7f962571a0020f4299f24b750276e Mon Sep 17 00:00:00 2001 From: chilung Date: Mon, 17 Aug 2026 00:08:47 +0800 Subject: [PATCH 1/3] fix(google): retry transient 429/5xx for AI Studio direct requests AI Studio direct (generativelanguage.googleapis.com) requests went through the default server fetch path, which retries connection resets but never HTTP error statuses. During capacity spikes the upstream returns 503 UNAVAILABLE ("This model is currently experiencing high demand") and every affected turn failed immediately (sendCount=1 in the request log), while Vertex and Antigravity already had Kiro-style bounded retry. Route direct AI Studio through the shared Google retry wrapper with the existing surface preserved: raw Provider error : text (no classification) and single-shot 400 semantics (no request-shape compatibility replay). Transient 500/502/503/504 and plain rate-limit 429s are now retried up to 3 attempts with Retry-After honoring and jittered backoff. Covered by focused tests: a transient 503 retries into success, a final 400 keeps its raw body and is not replayed, and rate-limit 429s stay bounded at 3 attempts with the raw body returned on exhaustion. --- src/adapters/google-http.ts | 38 +++++++++++++++++++++++----- src/adapters/google.ts | 14 ++++++++--- tests/google-vertex-http.test.ts | 43 +++++++++++++++++++++++++++++--- 3 files changed, 82 insertions(+), 13 deletions(-) diff --git a/src/adapters/google-http.ts b/src/adapters/google-http.ts index de849cde3c..a77e3c1a97 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,12 +27,19 @@ 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; let lastError: unknown; let activeRequest = request; @@ -40,7 +52,7 @@ export async function fetchGoogleWithRetry(label: string, request: AdapterReques headers: activeRequest.headers, body: activeRequest.body, }, timeoutMs, ctx.abortSignal, ctx.stream); - if (res.status === 400 && !compatibilityReplayUsed) { + if (res.status === 400 && repairInvalid400 && !compatibilityReplayUsed) { let payloadText = ""; try { payloadText = await readDisplaySafeErrorPayloadText(res.clone(), ctx.abortSignal); @@ -89,6 +101,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..6ab94b8427 100644 --- a/src/adapters/google.ts +++ b/src/adapters/google.ts @@ -17,7 +17,7 @@ import type { import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolAllowedByChoice } from "../types"; import { contentPartsToText, parseDataUrl } from "./image"; import { getVertexAccessToken } from "../lib/gcp-adc"; -import { fetchAntigravityWithRetry, fetchVertexWithRetry } from "./google-http"; +import { fetchAntigravityWithRetry, fetchDirectGeminiWithRetry, fetchVertexWithRetry } from "./google-http"; import { safeAntigravityHttpErrorMessage, safeVertexHttpErrorMessage } from "./google-errors"; import { isVertexTruncatedTurn, vertexTruncationErrorMessage } from "./google-truncation"; import { ANTIGRAVITY_REQUEST_UA, antigravitySessionId, isLikelyRealThoughtSignature, sanitizeAntigravityClaudeSignatures } from "./google-antigravity-wire"; @@ -377,8 +377,11 @@ 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. + // AI-Studio direct keeps the default raw "Provider error : " surface and its + // single-shot 400 semantics, but gains the same bounded transient 429/5xx retry so the + // frequent "model is currently experiencing high demand" 503s are retried instead of + // failing the turn immediately. ...(provider.googleMode === "vertex" || provider.googleMode === "cloud-code-assist" ? { fetchResponse: (request: AdapterRequest, ctx?: AdapterFetchContext): Promise => @@ -386,7 +389,10 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte formatErrorBody: (status: number, _headers: Headers, payloadText: string): string => (provider.googleMode === "cloud-code-assist" ? safeAntigravityHttpErrorMessage : safeVertexHttpErrorMessage)(status, payloadText), } - : {}), + : { + fetchResponse: (request: AdapterRequest, ctx?: AdapterFetchContext): Promise => + fetchDirectGeminiWithRetry(request, ctx), + }), async buildRequest(parsed: OcxParsedRequest) { const routedModelId = provider.googleMode === "cloud-code-assist" diff --git a/tests/google-vertex-http.test.ts b/tests/google-vertex-http.test.ts index e640f6fabe..25c5fcbd8f 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; @@ -212,6 +212,43 @@ 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 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); + }); }); describe("safeVertexHttpErrorMessage classification + redaction", () => { @@ -249,13 +286,13 @@ describe("safeVertexHttpErrorMessage classification + redaction", () => { }); describe("adapter fetchResponse wiring", () => { - test("vertex adapter exposes fetchResponse; ai-studio does not", async () => { + test("vertex and ai-studio adapters expose fetchResponse; ai-studio keeps default error formatting", 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); expect(typeof vertex.fetchResponse).toBe("function"); expect(typeof vertex.formatErrorBody).toBe("function"); - expect(aistudio.fetchResponse).toBeUndefined(); + expect(typeof aistudio.fetchResponse).toBe("function"); expect(aistudio.formatErrorBody).toBeUndefined(); }); From 97d8e4bbe89d57806995a0b01da2188398652138 Mon Sep 17 00:00:00 2001 From: chilung Date: Mon, 17 Aug 2026 02:23:24 +0800 Subject: [PATCH 2/3] fix(google): keep quota-exhausted 429 responses single-shot in raw mode --- src/adapters/google-http.ts | 7 ++++--- tests/google-vertex-http.test.ts | 25 +++++++++++++++++++------ 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/src/adapters/google-http.ts b/src/adapters/google-http.ts index a77e3c1a97..d23c42985c 100644 --- a/src/adapters/google-http.ts +++ b/src/adapters/google-http.ts @@ -73,10 +73,11 @@ export async function fetchGoogleWithRetry( } // 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), }); diff --git a/tests/google-vertex-http.test.ts b/tests/google-vertex-http.test.ts index 25c5fcbd8f..55a8140aca 100644 --- a/tests/google-vertex-http.test.ts +++ b/tests/google-vertex-http.test.ts @@ -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); }); @@ -237,6 +237,19 @@ describe("vertex retry fetch", () => { 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([ From 844b313bc7acfc3f35ed30f9139292c5b7e89cfc Mon Sep 17 00:00:00 2001 From: chilung Date: Mon, 17 Aug 2026 15:03:41 +0800 Subject: [PATCH 3/3] fix(google): route direct Gemini retries through canonical transport and preserve key pool failover --- src/adapters/base.ts | 2 + src/adapters/google-http.ts | 3 +- src/adapters/google.ts | 14 +++--- src/images/loop.ts | 1 + src/lib/upstream-retry.ts | 3 +- src/server/responses/core.ts | 22 ++++++++-- tests/google-vertex-http.test.ts | 20 ++++++++- tests/request-pacing.test.ts | 20 +++++++++ tests/server-key-failover-e2e.test.ts | 61 +++++++++++++++++++++++++++ tests/upstream-http-version.test.ts | 18 ++++++++ 10 files changed, 148 insertions(+), 16 deletions(-) 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 d23c42985c..f7b90de87e 100644 --- a/src/adapters/google-http.ts +++ b/src/adapters/google-http.ts @@ -41,6 +41,7 @@ export async function fetchGoogleWithRetry( ): 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; @@ -51,7 +52,7 @@ export async function fetchGoogleWithRetry( method: activeRequest.method, headers: activeRequest.headers, body: activeRequest.body, - }, timeoutMs, ctx.abortSignal, ctx.stream); + }, timeoutMs, ctx.abortSignal, ctx.stream, executor); if (res.status === 400 && repairInvalid400 && !compatibilityReplayUsed) { let payloadText = ""; try { diff --git a/src/adapters/google.ts b/src/adapters/google.ts index 6ab94b8427..8b8d577dee 100644 --- a/src/adapters/google.ts +++ b/src/adapters/google.ts @@ -17,7 +17,7 @@ import type { import { isAllowedToolChoice, namespacedToolName, resolveToolChoiceWireName, toolAllowedByChoice } from "../types"; import { contentPartsToText, parseDataUrl } from "./image"; import { getVertexAccessToken } from "../lib/gcp-adc"; -import { fetchAntigravityWithRetry, fetchDirectGeminiWithRetry, fetchVertexWithRetry } from "./google-http"; +import { fetchAntigravityWithRetry, fetchVertexWithRetry } from "./google-http"; import { safeAntigravityHttpErrorMessage, safeVertexHttpErrorMessage } from "./google-errors"; import { isVertexTruncatedTurn, vertexTruncationErrorMessage } from "./google-truncation"; import { ANTIGRAVITY_REQUEST_UA, antigravitySessionId, isLikelyRealThoughtSignature, sanitizeAntigravityClaudeSignatures } from "./google-antigravity-wire"; @@ -378,10 +378,9 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte name: "google", // Vertex + Antigravity get Kiro-style retry/timeout + classified, redacted errors. - // AI-Studio direct keeps the default raw "Provider error : " surface and its - // single-shot 400 semantics, but gains the same bounded transient 429/5xx retry so the - // frequent "model is currently experiencing high demand" 503s are retried instead of - // failing the turn immediately. + // 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 => @@ -389,10 +388,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte formatErrorBody: (status: number, _headers: Headers, payloadText: string): string => (provider.googleMode === "cloud-code-assist" ? safeAntigravityHttpErrorMessage : safeVertexHttpErrorMessage)(status, payloadText), } - : { - fetchResponse: (request: AdapterRequest, ctx?: AdapterFetchContext): Promise => - fetchDirectGeminiWithRetry(request, ctx), - }), + : {}), async buildRequest(parsed: OcxParsedRequest) { const routedModelId = provider.googleMode === "cloud-code-assist" 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 55a8140aca..496fb56a1a 100644 --- a/tests/google-vertex-http.test.ts +++ b/tests/google-vertex-http.test.ts @@ -262,6 +262,19 @@ describe("vertex retry fetch", () => { 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", () => { @@ -299,13 +312,16 @@ describe("safeVertexHttpErrorMessage classification + redaction", () => { }); describe("adapter fetchResponse wiring", () => { - test("vertex and ai-studio adapters expose fetchResponse; ai-studio keeps default error formatting", 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 aistudio.fetchResponse).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"); + }); });