From 9b6d16c7b9063f07cf85beb09a23b95becf86743 Mon Sep 17 00:00:00 2001 From: agentHits Date: Wed, 19 Aug 2026 18:13:29 +0300 Subject: [PATCH 1/2] fix(google): preserve and replay thought signatures across stream chunks and multi-tool turns (#2125) --- src/adapters/google-antigravity-replay.ts | 21 +++++--- src/adapters/google.ts | 25 ++++++++-- tests/google-antigravity-replay.test.ts | 27 ++++++++++ ...google-signature-history-roundtrip.test.ts | 50 +++++++++++++++++++ 4 files changed, 111 insertions(+), 12 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index d3ffa480c8..111a502041 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -601,10 +601,15 @@ export function antigravityUsesReplayCache(model: string): boolean { * 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. - * `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; +* `parts` is the already-unwrapped `response.candidates[0].content.parts`. +*/ +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 +623,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 +637,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 +649,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 +663,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 +672,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..8707d4904d 100644 --- a/tests/google-signature-history-roundtrip.test.ts +++ b/tests/google-signature-history-roundtrip.test.ts @@ -132,6 +132,56 @@ 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()); + const ssePayload = [ + `data: ${JSON.stringify(googleBody([{ text: "thinking...", thought: true, thought_signature: SIGNATURE }]))}`, + "", + `data: ${JSON.stringify(googleBody([{ functionCall: { name: "shell_command", args: { command: "pwd" } } }]))}`, + "", + `data: ${JSON.stringify(googleBody([{ functionCall: { name: "shell_command", args: { command: "ls" } } }]))}`, + "", + "data: [DONE]", + "", + ].join("\n"); + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(ssePayload)); + controller.close(); + }, + }); + + const events: AdapterEvent[] = []; + for await (const event of adapter.parseStream(new Response(stream))) { + events.push(event); + } + + 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({ From 3b076504eed0b62016d2f2466f740520740bd5ee Mon Sep 17 00:00:00 2001 From: agentHits Date: Wed, 19 Aug 2026 18:59:52 +0300 Subject: [PATCH 2/2] fix(google): align replay JSDoc with the same-turn contract and harden the streaming test - JSDoc now documents that a standalone thought signature applies to all subsequent functionCall parts of the turn and is carried across SSE chunks (carriedThoughtSig in / returned out), replacing the stale one-call pairing wording from #897; star alignment restored. - Streaming roundtrip test enqueues each SSE frame as a separate transport chunk, drops the [DONE] sentinel in favour of the usageMetadata terminal signal, and asserts error-free completion with a terminal done event (CodeRabbit finding). --- src/adapters/google-antigravity-replay.ts | 11 +++++--- ...google-signature-history-roundtrip.test.ts | 25 +++++++++++-------- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 111a502041..745ac44f9b 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -599,10 +599,13 @@ 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. -* `parts` is the already-unwrapped `response.candidates[0].content.parts`. -*/ + * 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, diff --git a/tests/google-signature-history-roundtrip.test.ts b/tests/google-signature-history-roundtrip.test.ts index 8707d4904d..b40c21f074 100644 --- a/tests/google-signature-history-roundtrip.test.ts +++ b/tests/google-signature-history-roundtrip.test.ts @@ -151,20 +151,20 @@ describe("#1735 thought signature survives history replay", () => { test("streaming SSE chunks carry thought signature across chunk boundaries to function calls", async () => { const adapter = createGoogleAdapter(provider); await adapter.buildRequest(firstTurn()); - const ssePayload = [ - `data: ${JSON.stringify(googleBody([{ text: "thinking...", thought: true, thought_signature: SIGNATURE }]))}`, - "", - `data: ${JSON.stringify(googleBody([{ functionCall: { name: "shell_command", args: { command: "pwd" } } }]))}`, - "", - `data: ${JSON.stringify(googleBody([{ functionCall: { name: "shell_command", args: { command: "ls" } } }]))}`, - "", - "data: [DONE]", - "", - ].join("\n"); + // 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) { - controller.enqueue(new TextEncoder().encode(ssePayload)); + const encoder = new TextEncoder(); + for (const frame of frames) controller.enqueue(encoder.encode(frame)); controller.close(); }, }); @@ -174,6 +174,9 @@ describe("#1735 thought signature survives history replay", () => { 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)