fix(google): preserve and replay thought signatures across stream chunks

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<string, unknown>;
@@ -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;
 }
 
 /**

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<AdapterEvent, "continue" | "content" | "terminate"> {
         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" });

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" },

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