From c5d1a0aa9c4afca37a6261ce9f8df7c1ee1d32ee Mon Sep 17 00:00:00 2001 From: usr-bin-roygbiv Date: Mon, 10 Aug 2026 01:58:11 +0000 Subject: [PATCH] fix(ai): fail over stalled Antigravity streams --- packages/ai/CHANGELOG.md | 4 + .../ai/src/providers/google-gemini-cli.ts | 42 ++++++--- .../ai/src/providers/register-builtins.ts | 32 ++++--- ...gle-gemini-cli-first-event-timeout.test.ts | 92 +++++++++++++++++++ 4 files changed, 145 insertions(+), 25 deletions(-) create mode 100644 packages/ai/test/google-gemini-cli-first-event-timeout.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 0e26ffbf0..a0ee0964d 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed Antigravity Flash turns hanging after successful response headers when the endpoint never emitted an SSE event; the provider now cancels the stalled body and fails over after 60 seconds while retaining the longer allowance for Pro reasoning starts. + ## [17.2.12] - 2026-08-08 ### Fixed diff --git a/packages/ai/src/providers/google-gemini-cli.ts b/packages/ai/src/providers/google-gemini-cli.ts index 8631e9900..7335db688 100644 --- a/packages/ai/src/providers/google-gemini-cli.ts +++ b/packages/ai/src/providers/google-gemini-cli.ts @@ -31,7 +31,7 @@ import { normalizeSystemPrompts } from "../utils"; import { AssistantMessageEventStream } from "../utils/event-stream"; import { extractGoogleValidationUrl, formatGoogleValidationRequiredMessage } from "../utils/google-validation"; import type { RawHttpRequestDump } from "../utils/http-inspector"; -import { armPreResponseTimeout, getStreamFirstEventTimeoutMs } from "../utils/idle-iterator"; +import { armPreResponseTimeout, getStreamFirstEventTimeoutMs, iterateWithIdleTimeout } from "../utils/idle-iterator"; // Refresh is the sole responsibility of AuthStorage (broker-aware, single-flighted); // the stream provider trusts the access token threaded through `options.apiKey`. import { normalizeSchemaForCCA } from "../utils/schema"; @@ -325,6 +325,9 @@ export { // Retry configuration const MAX_RETRIES = 3; const BASE_DELAY_MS = 1000; +const FLASH_FIRST_EVENT_TIMEOUT_MS = 60_000; +const DEFAULT_FIRST_EVENT_TIMEOUT_MS = 300_000; +const FIRST_EVENT_TIMEOUT_ERROR = "Cloud Code Assist stream timed out while waiting for the first event"; const RATE_LIMIT_BUDGET_MS = 5 * 60 * 1000; const CLAUDE_THINKING_BETA_HEADER = "interleaved-thinking-2025-05-14"; const GOOGLE_GEMINI_REFRESH_SKEW_MS = 60_000; @@ -616,12 +619,16 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = ( headers: requestHeaders, }; - // Direct callers that skip `register-builtins` (which installs the - // iterator-level watchdog) need a pre-response timer alongside - // `timeout: false`; otherwise a stalled Cloud Code Assist proxy - // would hang forever. Floor matches the lazy wrapper's 5min default. + // The provider owns the first-event watchdog so a silent successful + // response can fail over to the alternate Antigravity endpoint before + // anything user-visible has streamed. Flash should not inherit the + // five-minute allowance reserved for cold Pro reasoning starts. const firstEventTimeoutMs = - options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(undefined, 300_000); + options?.streamFirstEventTimeoutMs ?? + getStreamFirstEventTimeoutMs( + undefined, + model.id.includes("flash") ? FLASH_FIRST_EVENT_TIMEOUT_MS : DEFAULT_FIRST_EVENT_TIMEOUT_MS, + ); const callerSignal = options?.signal; const toolNames = new Set(context.tools?.map(t => t.name) ?? []); const isFlashLeakModel = model.id.includes("flash"); @@ -755,11 +762,24 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = ( } }; - for await (const chunk of readSseJson( - activeResponse.body!, - options?.signal, - event => options?.onSseEvent?.({ event: event.event, data: event.data, raw: [...event.raw] }, model), - )) { + const responseAbortController = new AbortController(); + const responseSignal = options?.signal + ? AbortSignal.any([options.signal, responseAbortController.signal]) + : responseAbortController.signal; + const chunks = iterateWithIdleTimeout( + readSseJson(activeResponse.body, responseSignal, event => + options?.onSseEvent?.({ event: event.event, data: event.data, raw: [...event.raw] }, model), + ), + { + firstItemTimeoutMs: firstEventTimeoutMs, + errorMessage: FIRST_EVENT_TIMEOUT_ERROR, + firstItemErrorMessage: FIRST_EVENT_TIMEOUT_ERROR, + onFirstItemTimeout: () => + responseAbortController.abort(new AIError.StreamTimeoutError(FIRST_EVENT_TIMEOUT_ERROR)), + abortSignal: options?.signal, + }, + ); + for await (const chunk of chunks) { if (chunk.error) { const detail = chunk.error.message || chunk.error.status || "unknown error"; const message = `Cloud Code Assist stream error: ${detail}`; diff --git a/packages/ai/src/providers/register-builtins.ts b/packages/ai/src/providers/register-builtins.ts index 1b7a59e20..d2dbfd48c 100644 --- a/packages/ai/src/providers/register-builtins.ts +++ b/packages/ai/src/providers/register-builtins.ts @@ -202,6 +202,11 @@ interface LazyStreamLimits { * stream timeouts. Keep the lazy loader from racing it with generic errors. */ providerHandlesStreamTimeouts?: boolean; + /** + * The provider retries or fails over when no first event arrives, while the + * lazy wrapper continues to own steady-state idle detection. + */ + providerHandlesFirstEventTimeouts?: boolean; /** * Apply OpenAI-family idle timeout precedence in the lazy wrapper. Used by * local backends whose users historically tune slow prompt-processing gaps @@ -210,16 +215,13 @@ interface LazyStreamLimits { openAIIdleEnvFloorsFirstEvent?: boolean; } /** - * Cloud Code Assist (google-gemini-cli / google-antigravity) routinely takes - * longer than the global 100s default to emit its first SSE event when serving - * the heavier Gemini 3.x Pro tiers at high thinking levels. Bump the first-event - * floor to five minutes so callers stop seeing spurious "stream timed out while - * waiting for the first event" aborts on legitimate cold reasoning starts. - * The steady-state idle watchdog stays on the global default since the upstream - * emits thinking tokens frequently once it gets going. + * Cloud Code Assist owns first-event detection because Antigravity can return + * successful headers and then never emit an SSE event. Keeping the watchdog in + * the provider lets it fail over before surfacing an error; the lazy wrapper + * still catches post-first-event stalls. */ const GOOGLE_GEMINI_CLI_LAZY_STREAM_LIMITS: LazyStreamLimits = { - defaultFirstEventTimeoutMs: 300_000, + providerHandlesFirstEventTimeouts: true, }; const PROVIDER_HANDLED_STREAM_TIMEOUTS: LazyStreamLimits = { @@ -241,6 +243,7 @@ function forwardStream( (async () => { try { const providerHandlesStreamTimeouts = limits?.providerHandlesStreamTimeouts === true; + const providerHandlesFirstEventTimeouts = limits?.providerHandlesFirstEventTimeouts === true; // Per-model catalog compat can widen the fallback watchdog for hosts // with no keepalive events (e.g. Bedrock reasoning models that go // quiet for minutes mid-thinking, issue #4758). Caller options and @@ -258,12 +261,13 @@ function forwardStream( (limits?.openAIIdleEnvFloorsFirstEvent ? getOpenAIStreamIdleTimeoutMs(idleTimeoutFallbackMs) : getStreamIdleTimeoutMs(idleTimeoutFallbackMs))); - const firstItemTimeoutMs = providerHandlesStreamTimeouts - ? 0 - : (options.streamFirstEventTimeoutMs ?? - (limits?.openAIIdleEnvFloorsFirstEvent - ? getOpenAIStreamFirstEventTimeoutMs(idleTimeoutMs, limits.defaultFirstEventTimeoutMs) - : getStreamFirstEventTimeoutMs(idleTimeoutMs, limits?.defaultFirstEventTimeoutMs))); + const firstItemTimeoutMs = + providerHandlesStreamTimeouts || providerHandlesFirstEventTimeouts + ? 0 + : (options.streamFirstEventTimeoutMs ?? + (limits?.openAIIdleEnvFloorsFirstEvent + ? getOpenAIStreamFirstEventTimeoutMs(idleTimeoutMs, limits.defaultFirstEventTimeoutMs) + : getStreamFirstEventTimeoutMs(idleTimeoutMs, limits?.defaultFirstEventTimeoutMs))); // Providers with a server-driven local tool bridge (e.g. the Cursor // exec channel) mark their stream busy while a local tool runs; the // watchdog must not read that silence as a provider stall (#4593). diff --git a/packages/ai/test/google-gemini-cli-first-event-timeout.test.ts b/packages/ai/test/google-gemini-cli-first-event-timeout.test.ts new file mode 100644 index 000000000..0b06d00d9 --- /dev/null +++ b/packages/ai/test/google-gemini-cli-first-event-timeout.test.ts @@ -0,0 +1,92 @@ +import { afterEach, expect, test, vi } from "bun:test"; +import { streamGoogleGeminiCli } from "@oh-my-pi/pi-ai/providers/google-gemini-cli"; +import type { Context, FetchImpl, Model } from "@oh-my-pi/pi-ai/types"; +import { buildModel } from "@oh-my-pi/pi-catalog/build"; + +const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com"; +const ANTIGRAVITY_SANDBOX_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"; +const FLASH_FIRST_EVENT_TIMEOUT_MS = 60_000; +const context: Context = { messages: [{ role: "user", content: "hi", timestamp: 1 }] }; +const antigravityModel: Model<"google-gemini-cli"> = buildModel({ + id: "gemini-3-flash", + name: "Gemini 3 Flash (Antigravity)", + api: "google-gemini-cli", + provider: "google-antigravity", + baseUrl: ANTIGRAVITY_DAILY_ENDPOINT, + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 32_000, +}); + +afterEach(() => { + vi.useRealTimers(); +}); + +function endpointFromInput(input: Parameters[0]): string { + const url = input instanceof Request ? input.url : input.toString(); + return url.startsWith(ANTIGRAVITY_SANDBOX_ENDPOINT) ? ANTIGRAVITY_SANDBOX_ENDPOINT : ANTIGRAVITY_DAILY_ENDPOINT; +} + +function responseWithUrl(response: Response, endpoint: string): Response { + Object.defineProperty(response, "url", { value: `${endpoint}/v1internal:streamGenerateContent?alt=sse` }); + return response; +} + +test("Antigravity Flash fails over when headers arrive without a first SSE event", async () => { + const requestedEndpoints: string[] = []; + let dailyBodyCancelled = false; + const dailyBodyReadStarted = Promise.withResolvers(); + const dailyBodyStall = Promise.withResolvers(); + vi.useFakeTimers(); + + const fetchMock: FetchImpl = async input => { + const endpoint = endpointFromInput(input); + requestedEndpoints.push(endpoint); + if (endpoint === ANTIGRAVITY_SANDBOX_ENDPOINT) { + const body = `data: ${JSON.stringify({ + response: { + candidates: [{ content: { parts: [{ text: "Recovered after stall." }] }, finishReason: "STOP" }], + }, + })}\n\n`; + return responseWithUrl( + new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }), + endpoint, + ); + } + + return responseWithUrl( + new Response( + new ReadableStream({ + pull() { + dailyBodyReadStarted.resolve(); + return dailyBodyStall.promise; + }, + cancel() { + dailyBodyStall.resolve(); + dailyBodyCancelled = true; + }, + }), + { status: 200, headers: { "content-type": "text/event-stream" } }, + ), + endpoint, + ); + }; + + const stream = streamGoogleGeminiCli(antigravityModel, context, { + apiKey: JSON.stringify({ token: "token", projectId: "proj-123" }), + antigravityEndpointMode: "auto", + fetch: fetchMock, + }); + const resultPromise = stream.result(); + await dailyBodyReadStarted.promise; + expect(vi.getTimerCount()).toBeGreaterThan(0); + vi.advanceTimersByTime(FLASH_FIRST_EVENT_TIMEOUT_MS * 2); + const result = await resultPromise; + + expect(requestedEndpoints).toEqual([ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_SANDBOX_ENDPOINT]); + expect(dailyBodyCancelled).toBe(true); + expect(result.stopReason).toBe("stop"); + expect(result.content).toEqual([{ type: "text", text: "Recovered after stall." }]); +});