From 1adcfde0cfff10589eefdaa33aae077ffe404bcc Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Thu, 20 Aug 2026 05:50:07 +0900 Subject: [PATCH] fix(google): preserve and replay thought signatures across stream chunks Antigravity thinking models 400 on multi-turn tool loops: the historical tool calls are replayed without their thought_signature, so exec/wait/agent tools break once history is long enough to matter. Two holes, both in the #897 "pair the signature to the next call and drop the rest" contract: pendingThoughtSig was cleared by the first functionCall, so a parallel tool turn left every call after the first unsigned. It now survives the turn. A signature and its function call can arrive in different SSE chunks, so the pending value was lost at the chunk boundary. It is now carried across. The parser also reads the snake_case thought_signature alongside the camelCase form, which is what the upstream actually sends on some turns. Carries @agentHits's #2127 unchanged. Closes #2125 --- src/adapters/google-antigravity-replay.ts | 24 ++++++--- src/adapters/google.ts | 25 +++++++-- tests/google-antigravity-replay.test.ts | 27 ++++++++++ ...google-signature-history-roundtrip.test.ts | 53 +++++++++++++++++++ 4 files changed, 117 insertions(+), 12 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index d3ffa480c8..745ac44f9b 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -599,12 +599,20 @@ export function antigravityUsesReplayCache(model: string): boolean { * Observe a parsed CCA chunk's `candidates[0].content.parts` and record thought signatures keyed by * the functionCall identity (name + args). Accumulates across the whole session so a sequential * multi-step tool loop keeps EVERY prior call's signature, not just the latest part-index slot. - * A signature on a standalone thought part is paired with the next functionCall in the same - * array (#897); a call's own signature takes precedence and an unpaired one is dropped. + * A signature on a standalone thought part applies to the functionCall parts that follow it in + * the same array AND to later arrays of the same turn: streaming splits a thought part and its + * calls across SSE chunks, so `carriedThoughtSig` threads the still-unpaired signature from the + * previous chunk and the return value hands the remainder to the next one (#897, #2125). A call's + * own signature always takes precedence over a carried one. * `parts` is the already-unwrapped `response.candidates[0].content.parts`. */ -export function observeAntigravityReplay(model: string, sessionId: string, parts: unknown[]): void { - if (!antigravityUsesReplayCache(model) || !Array.isArray(parts) || parts.length === 0) return; +export function observeAntigravityReplay( + model: string, + sessionId: string, + parts: unknown[], + carriedThoughtSig?: string, +): string | undefined { + if (!antigravityUsesReplayCache(model) || !Array.isArray(parts) || parts.length === 0) return carriedThoughtSig; ensureReplaySnapshotLoaded(); const now = Date.now(); deleteExpiredReplaySessionsThrottled(now); @@ -618,7 +626,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts lastActiveAtMs: 0, }; let inserted = false; - let pendingThoughtSig: string | undefined; + let pendingThoughtSig: string | undefined = carriedThoughtSig; for (const raw of parts) { if (!raw || typeof raw !== "object") continue; const part = raw as Record; @@ -632,7 +640,6 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts continue; } const callSig = sig ?? pendingThoughtSig; // a signature on the call part itself wins - pendingThoughtSig = undefined; if (!callSig) continue; const ck = functionCallKey(fc.name, fc.args); if (!ck) continue; // only function-call signatures are replayable by identity @@ -645,7 +652,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts replayBytes += sizeBytes; inserted = true; } - if (!inserted) return; + if (!inserted) return pendingThoughtSig; // Charge the fixed outer key only when the session is actually stored. if (!existing) replayBytes += REPLAY_SESSION_KEY_BYTES; evictInnerCalls(entry); @@ -659,7 +666,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts } else { replayBytes -= REPLAY_SESSION_KEY_BYTES; } - return; + return pendingThoughtSig; } entry.expiresAtMs = now + REPLAY_TTL_MS; entry.lastActiveAtMs = now; @@ -668,6 +675,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts evictIfNeeded(); enforceAppOwnedMemoryBudget(); markReplayDirty(); + return pendingThoughtSig; } /** diff --git a/src/adapters/google.ts b/src/adapters/google.ts index a3be021a76..f152d91741 100644 --- a/src/adapters/google.ts +++ b/src/adapters/google.ts @@ -342,6 +342,7 @@ interface GoogleResponsePart { text?: string; thought?: boolean; thoughtSignature?: string; + thought_signature?: string; functionCall?: { name: string; args: unknown }; } @@ -352,8 +353,9 @@ interface GoogleResponsePart { */ function googleToolCallMetadataFromPart( part: GoogleResponsePart, + fallbackSignature?: string, ): { providerMetadata: OcxProviderOpaqueToolCallMetadata } | undefined { - const signature = part.thoughtSignature; + const signature = part.thoughtSignature ?? part.thought_signature ?? fallbackSignature; if (!isLikelyRealThoughtSignature(signature)) return undefined; return { providerMetadata: { google: { thoughtSignature: signature } } }; } @@ -601,6 +603,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte let lastFinishReason: string | undefined; let sawAnyFrame = false; let sawTerminalSignal = false; + let pendingStreamThoughtSig: string | undefined; const handleDataLine = async function* (line: string): AsyncGenerator { const payload = line.slice(5).trim(); @@ -697,10 +700,19 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte const replaySession = provider.googleMode === "cloud-code-assist" ? antigravitySession : vertexReplaySession; if ((provider.googleMode === "cloud-code-assist" || provider.googleMode === "vertex") && parts && replayModel && replaySession) { - observeAntigravityReplay(replayModel, replaySession, parts as unknown[]); + pendingStreamThoughtSig = observeAntigravityReplay( + replayModel, + replaySession, + parts as unknown[], + pendingStreamThoughtSig, + ); } if (parts) { for (const part of parts) { + const sig = part.thoughtSignature ?? part.thought_signature; + if (part.thought === true && sig && isLikelyRealThoughtSignature(sig)) { + pendingStreamThoughtSig = sig; + } const textEvent = googlePartTextEvent(part); if (textEvent) { emittedContentEvent = true; @@ -729,7 +741,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte type: "tool_call_start", id, name: restoreGoogleToolName(part.functionCall.name), - ...googleToolCallMetadataFromPart(part), + ...googleToolCallMetadataFromPart(part, pendingStreamThoughtSig), }; yield { type: "tool_call_delta", arguments: JSON.stringify(part.functionCall.args ?? {}) }; yield { type: "tool_call_end" }; @@ -926,7 +938,12 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte && replayModel && replaySession) { observeAntigravityReplay(replayModel, replaySession, candidates[0].content.parts as unknown[]); } + let pendingThoughtSig: string | undefined; for (const part of candidates[0].content.parts) { + const sig = part.thoughtSignature ?? part.thought_signature; + if (part.thought === true && sig && isLikelyRealThoughtSignature(sig)) { + pendingThoughtSig = sig; + } const textEvent = googlePartTextEvent(part); if (textEvent) events.push(textEvent); const inline = (part as { inlineData?: { mimeType?: string; data?: string } }).inlineData; @@ -950,7 +967,7 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte type: "tool_call_start", id, name: restoreGoogleToolName(part.functionCall.name), - ...googleToolCallMetadataFromPart(part), + ...googleToolCallMetadataFromPart(part, pendingThoughtSig), }); events.push({ type: "tool_call_delta", arguments: JSON.stringify(part.functionCall.args ?? {}) }); events.push({ type: "tool_call_end" }); diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index eb02fd6006..0ef9524533 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -97,6 +97,33 @@ describe("antigravity reasoning-replay cache", () => { expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe(SIG); }); + test("a signature on a standalone thought part replays onto all subsequent functionCalls in the same turn", () => { + observeAntigravityReplay(MODEL, SESSION, [ + { thought: true, thoughtSignature: SIG }, + fcPart("get_x", { a: 1 }), + fcPart("get_y", { b: 2 }), + ]); + const contents = [{ + role: "model", + parts: [ + { functionCall: { name: "get_x", args: { a: 1 } } }, + { functionCall: { name: "get_y", args: { b: 2 } } }, + ], + }]; + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe(SIG); + expect((contents[0].parts[1] as { thoughtSignature?: string }).thoughtSignature).toBe(SIG); + }); + + test("carried thought signature across separate stream chunks pairs with subsequent functionCalls", () => { + const carried = observeAntigravityReplay(MODEL, SESSION, [{ thought: true, thought_signature: SIG }]); + expect(carried).toBe(SIG); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 })], carried); + const contents = [{ role: "model", parts: [{ functionCall: { name: "get_x", args: { a: 1 } } }] }]; + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe(SIG); + }); + test("a call's own signature wins over a preceding standalone thought signature", () => { observeAntigravityReplay(MODEL, SESSION, [ { thought: true, thoughtSignature: "sig-standalone-aaaa" }, diff --git a/tests/google-signature-history-roundtrip.test.ts b/tests/google-signature-history-roundtrip.test.ts index 499c7c8c3f..b40c21f074 100644 --- a/tests/google-signature-history-roundtrip.test.ts +++ b/tests/google-signature-history-roundtrip.test.ts @@ -132,6 +132,59 @@ describe("#1735 thought signature survives history replay", () => { expect(signatures).toEqual([SIGNATURE, SIGNATURE_B]); }); + test("a standalone thought part's signature attaches to all subsequent functionCalls in non-streaming parse", async () => { + const adapter = createGoogleAdapter(provider); + await adapter.buildRequest(firstTurn()); + const events = await adapter.parseResponse!(new Response(JSON.stringify(googleBody([ + { text: "thinking...", thought: true, thoughtSignature: SIGNATURE }, + { functionCall: { name: "shell_command", args: { command: "pwd" } } }, + { functionCall: { name: "shell_command", args: { command: "ls" } } }, + ])))); + const starts = events.filter((e: AdapterEvent) => e.type === "tool_call_start"); + expect(starts.length).toBe(2); + expect("providerMetadata" in starts[0] ? starts[0].providerMetadata?.google?.thoughtSignature : undefined) + .toBe(SIGNATURE); + expect("providerMetadata" in starts[1] ? starts[1].providerMetadata?.google?.thoughtSignature : undefined) + .toBe(SIGNATURE); + }); + + test("streaming SSE chunks carry thought signature across chunk boundaries to function calls", async () => { + const adapter = createGoogleAdapter(provider); + await adapter.buildRequest(firstTurn()); + // Each SSE frame arrives as its own transport chunk so the signature has to + // survive the chunk boundary between the thought part and the function calls. + const frames = [ + `data: ${JSON.stringify(googleBody([{ text: "thinking...", thought: true, thought_signature: SIGNATURE }]))}\n\n`, + `data: ${JSON.stringify(googleBody([{ functionCall: { name: "shell_command", args: { command: "pwd" } } }]))}\n\n`, + `data: ${JSON.stringify(googleBody([{ functionCall: { name: "shell_command", args: { command: "ls" } } }]))}\n\n`, + // usageMetadata is the terminal signal; no [DONE] sentinel needed. + `data: ${JSON.stringify({ usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 2 } })}\n\n`, + ]; + + const stream = new ReadableStream({ + start(controller) { + const encoder = new TextEncoder(); + for (const frame of frames) controller.enqueue(encoder.encode(frame)); + controller.close(); + }, + }); + + const events: AdapterEvent[] = []; + for await (const event of adapter.parseStream(new Response(stream))) { + events.push(event); + } + + // The turn completes cleanly: no error events, terminal done last. + expect(events.some((e: AdapterEvent) => e.type === "error")).toBe(false); + expect(events[events.length - 1]?.type).toBe("done"); + const starts = events.filter((e: AdapterEvent) => e.type === "tool_call_start"); + expect(starts.length).toBe(2); + expect("providerMetadata" in starts[0] ? starts[0].providerMetadata?.google?.thoughtSignature : undefined) + .toBe(SIGNATURE); + expect("providerMetadata" in starts[1] ? starts[1].providerMetadata?.google?.thoughtSignature : undefined) + .toBe(SIGNATURE); + }); + test("a signature replayed through Responses history reaches the rebuilt Google part", async () => { // No cache is warmed here: this is a cold process replaying client-supplied history. const parsed = parseRequestScoped({