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({