fix(google): match unwrapped freeform tool arguments in antigravity replay cache

src/adapters/google-antigravity-replay.ts

@@ -702,11 +702,38 @@ export function applyAntigravityReplay(model: string, sessionId: string, content
       if (!fc) continue;
       if (part.thoughtSignature !== undefined || part.thought_signature !== undefined) continue;
       const ck = functionCallKey(fc.name, fc.args);
-      const call = ck ? entry.byCall.get(ck) : undefined;
-      if (call && ck) {
+      let call = ck ? entry.byCall.get(ck) : undefined;
+      let matchedKey = ck;
+      if (!call && typeof fc.name === "string" && typeof fc.args === "object" && fc.args !== null) {
+        // Freeform / custom tool replay unwrap:
+        // The client replays custom_tool_call with arguments: { input: "..." }.
+        // Upstream was invoked with args: { input: "..." } or raw string or parsed JSON.
+        const argsObj = fc.args as Record<string, unknown>;
+        if (typeof argsObj.input === "string") {
+          const trimmedInput = argsObj.input.trim();
+          if (
+            trimmedInput.length <= REPLAY_MAX_CANONICAL_ARGS_BYTES
+            && utf8.encode(trimmedInput).byteLength <= REPLAY_MAX_CANONICAL_ARGS_BYTES
+          ) {
+            try {
+              const parsedInput = JSON.parse(trimmedInput);
+            if (parsedInput && typeof parsedInput === "object") {
+              const altKey = functionCallKey(fc.name, parsedInput);
+              if (altKey && entry.byCall.has(altKey)) {
+                call = entry.byCall.get(altKey);
+                matchedKey = altKey;
+              }
+            }
+            } catch {
+              // not JSON, keep default
+            }
+          }
+        }
+      }
+      if (call && matchedKey) {
         part.thoughtSignature = call.signature;
-        entry.byCall.delete(ck);
-        entry.byCall.set(ck, { ...call, touchedAtMs: now });
+        entry.byCall.delete(matchedKey);
+        entry.byCall.set(matchedKey, { ...call, touchedAtMs: now });
         touched = true;
       }
     }

src/adapters/google.ts

@@ -737,10 +737,11 @@ export function createGoogleAdapter(provider: OcxProviderConfig): ProviderAdapte
               const id = `call_${crypto.randomUUID().slice(0, 8)}`;
               toolCallsStarted++;
               emittedContentEvent = true;
+              const restoredName = restoreGoogleToolName(part.functionCall.name);
               yield {
                 type: "tool_call_start",
                 id,
-                name: restoreGoogleToolName(part.functionCall.name),
+                name: restoredName,
                 ...googleToolCallMetadataFromPart(part, pendingStreamThoughtSig),
               };
               yield { type: "tool_call_delta", arguments: JSON.stringify(part.functionCall.args ?? {}) };

tests/google-antigravity-replay.test.ts

@@ -205,6 +205,80 @@ describe("antigravity reasoning-replay cache", () => {
     expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe("sig-orderindep00000");
   });
 
+  test("matches freeform/custom_tool_call {input: string} against observed parsed JSON args", () => {
+    // Upstream saw functionCall with parsed args { cmd: "ls -la" }
+    observeAntigravityReplay(MODEL, SESSION, [fcPart("default_api:exec", { cmd: "ls -la" }, "sig-freeform-exec-1111")]);
+    // Client replays custom_tool_call with serialized input { input: '{"cmd":"ls -la"}' }
+    const contents = [{
+      role: "model",
+      parts: [{ functionCall: { name: "default_api:exec", args: { input: JSON.stringify({ cmd: "ls -la" }) } } }],
+    }];
+    applyAntigravityReplay(MODEL, SESSION, contents);
+    expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe("sig-freeform-exec-1111");
+  });
+
+  test("matches alternate key even when wrapped representation exceeds canonical key limit (ck undefined)", () => {
+    observeAntigravityReplay(MODEL, SESSION, [fcPart("default_api:exec", { cmd: "x" }, "sig-whitespace-overflow")]);
+    // 2 MiB of leading/trailing whitespace makes the raw string large, but trimmed payload is small (under 64 KiB).
+    // It must parse the trimmed slice directly rather than passing the 2 MiB string to JSON.parse.
+    const smallJson = JSON.stringify({ cmd: "x" });
+    const bigWhitespaceJson = " ".repeat(2 * 1024 * 1024) + smallJson + " ".repeat(1024);
+    let parsedString = "";
+    const originalParse = JSON.parse;
+    JSON.parse = (text, reviver) => {
+      if (typeof text === "string" && text.includes('"cmd":"x"')) {
+        parsedString = text;
+      }
+      return originalParse(text, reviver);
+    };
+    const contents = [{
+      role: "model",
+      parts: [{ functionCall: { name: "default_api:exec", args: { input: bigWhitespaceJson } } }],
+    }];
+    try {
+      applyAntigravityReplay(MODEL, SESSION, contents);
+      expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe("sig-whitespace-overflow");
+      expect(parsedString).toBe(smallJson);
+    } finally {
+      JSON.parse = originalParse;
+    }
+  });
+
+  test("rejects oversized input before JSON.parse without matching or allocating", () => {
+    observeAntigravityReplay(MODEL, SESSION, [fcPart("default_api:exec", { data: "huge" }, "sig-should-not-match")]);
+    // 100 KiB of valid JSON payload exceeds REPLAY_MAX_CANONICAL_ARGS_BYTES and must be rejected before TextEncoder.encode / JSON.parse.
+    let parseCalled = false;
+    let encodeCalledOnPayload = false;
+    const originalParse = JSON.parse;
+    const originalEncode = TextEncoder.prototype.encode;
+    JSON.parse = (text, reviver) => {
+      if (typeof text === "string" && text.includes('"data":')) {
+        parseCalled = true;
+      }
+      return originalParse(text, reviver);
+    };
+    TextEncoder.prototype.encode = function (input) {
+      if (typeof input === "string" && input.includes('"data":')) {
+        encodeCalledOnPayload = true;
+      }
+      return originalEncode.call(this, input);
+    };
+    const oversizedPayload = JSON.stringify({ data: "a".repeat(100 * 1024) });
+    const contents = [{
+      role: "model",
+      parts: [{ functionCall: { name: "default_api:exec", args: { input: oversizedPayload } } }],
+    }];
+    try {
+      applyAntigravityReplay(MODEL, SESSION, contents);
+      expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined();
+      expect(parseCalled).toBe(false);
+      expect(encodeCalledOnPayload).toBe(false);
+    } finally {
+      JSON.parse = originalParse;
+      TextEncoder.prototype.encode = originalEncode;
+    }
+  });
+
   test("claude models do not use the replay cache", () => {
     expect(antigravityUsesReplayCache("claude-opus-4.6")).toBe(false);
     expect(antigravityUsesReplayCache("gemini-3-pro")).toBe(true);

tests/google-signature-history-roundtrip.test.ts

@@ -7,7 +7,7 @@ import { mkdtempSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { createGoogleAdapter as createGoogleAdapterProduction } from "../src/adapters/google";
-import { __resetAntigravityReplayCache } from "../src/adapters/google-antigravity-replay";
+import { __resetAntigravityReplayCache, observeAntigravityReplay } from "../src/adapters/google-antigravity-replay";
 import { parseRequest } from "../src/responses/parser";
 import {
   flushThoughtSignatureReplayForTests,
@@ -257,6 +257,44 @@ describe("#1735 thought signature survives history replay", () => {
     expect(part?.thoughtSignature).toBe(SIGNATURE_B);
   });
 
+  test("a custom_tool_call without call_id store entry falls back to in-memory replay cache by unwrapped args", async () => {
+    const adapter = createGoogleAdapter({
+      ...provider,
+      googleMode: "cloud-code-assist",
+      baseUrl: "https://daily-cloudcode-pa.googleapis.com",
+      project: "test-proj",
+      apiKey: "test-token",
+    });
+    const parsedDummy = parseRequestScoped({
+      model: MODEL,
+      input: [{ type: "message", role: "user", content: [{ type: "input_text", text: "test-session-freeform" }] }],
+      tools: [{ type: "function", name: "default_api:exec", description: "run", parameters: { type: "object" } }],
+    }, undefined);
+    const dummyReq = await adapter.buildRequest(parsedDummy);
+    const wireModel = JSON.parse(dummyReq.body as string).model;
+    const wireSession = JSON.parse(dummyReq.body as string).request.sessionId;
+    const wireToolName = JSON.parse(dummyReq.body as string).request.tools[0].functionDeclarations[0].name;
+
+    // Warm up the Antigravity replay cache with parsed function args:
+    observeAntigravityReplay(wireModel, wireSession, [
+      { functionCall: { name: wireToolName, args: { cmd: "whoami" } }, thoughtSignature: SIGNATURE },
+    ]);
+    const parsed = parseRequestScoped({
+      model: MODEL,
+      input: [
+        { type: "message", role: "user", content: [{ type: "input_text", text: "test-session-freeform" }] },
+        { type: "custom_tool_call", call_id: "call_custom_unscoped", name: "default_api:exec", input: JSON.stringify({ cmd: "whoami" }) },
+        { type: "custom_tool_call_output", call_id: "call_custom_unscoped", output: "agent" },
+      ],
+      tools: [{ type: "function", name: "default_api:exec", description: "run", parameters: { type: "object" } }],
+    }, undefined); // unscoped so durable store cannot hit
+    const request = await adapter.buildRequest(parsed);
+    const reqObj = JSON.parse(request.body as string);
+    const contents = reqObj.request.contents;
+    const modelTurn = contents.find((c: { role: string }) => c.role === "model");
+    expect(modelTurn.parts[0].thoughtSignature).toBe(SIGNATURE);
+  });
+
   test("a tool_search_call replay is re-signed from the proxy-side store", async () => {
     rememberThoughtSignatureForReplay("call_ts_1", SIGNATURE, scopeFor());
     const parsed = parseRequestScoped({