smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): fail over stalled Antigravity streams

Browse Source
This commit is contained in:
usr-bin-roygbiv
2026-08-10 01:58:11 +00:00
parent 45e12e5bb7
commit c5d1a0aa9c
4 changed files with 145 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/CHANGELOG.md
+4
View File
@@ -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
packages/ai/src/providers/google-gemini-cli.ts
+31 -11
View File
@@ -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
+18 -14
View File
@@ -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
+92
View File
@@ -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." }]);
});