diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 745ac44f9b..3d574b54d9 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/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; + 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; } } diff --git a/src/adapters/google.ts b/src/adapters/google.ts index f152d91741..5ce985dd4b 100644 --- a/src/adapters/google.ts +++ b/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 ?? {}) }; diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 0ef9524533..8fab813cca 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/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); diff --git a/tests/google-signature-history-roundtrip.test.ts b/tests/google-signature-history-roundtrip.test.ts index b40c21f074..c914825bee 100644 --- a/tests/google-signature-history-roundtrip.test.ts +++ b/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({