Merge PR #8130: fix(ai): fail over stalled Antigravity streams (@usr-bin-roygbiv)

packages/ai/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Fixed
 
 - Fixed the Amazon Bedrock and Cursor transports ignoring `StreamOptions.headers`; both built their request headers from scratch, so caller-supplied tracing or attribution headers were silently dropped while working on every other provider ([#8107](https://github.com/can1357/oh-my-pi/pull/8107) by [@svperfecta](https://github.com/svperfecta)).
+- 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
 

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<CloudCodeAssistResponseChunk>(
-					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<CloudCodeAssistResponseChunk>(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}`;

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<TApi extends Api>(
 	(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<TApi extends Api>(
 					(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).

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<FetchImpl>[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<void>();
+	const dailyBodyStall = Promise.withResolvers<void>();
+	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." }]);
+});