From d541ffbb657e4f8b87eb4986143c75613062af00 Mon Sep 17 00:00:00 2001 From: agentHits Date: Thu, 20 Aug 2026 15:35:35 +0300 Subject: [PATCH 1/4] fix(google): support unwrapped freeform tool arguments in antigravity replay cache - Freeform/custom tools (such as default_api:exec) are emitted to clients as custom_tool_call with { input: '{"cmd":...}' } while upstream observation records parsed function arguments { cmd: ... }. - applyAntigravityReplay now unwraps { input: string } to match against observed JSON argument signatures when exact matching misses, fixing 400 errors during early turns and replayed freeform execution (#2125). --- src/adapters/google-antigravity-replay.ts | 29 ++++++++++++-- src/adapters/google.ts | 3 +- tests/google-antigravity-replay.test.ts | 12 ++++++ ...google-signature-history-roundtrip.test.ts | 40 ++++++++++++++++++- 4 files changed, 79 insertions(+), 5 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 745ac44f9b..c2743be937 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -702,11 +702,34 @@ 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; + 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") { + try { + const parsedInput = JSON.parse(argsObj.input); + 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 && ck) { part.thoughtSignature = call.signature; - entry.byCall.delete(ck); - entry.byCall.set(ck, { ...call, touchedAtMs: now }); + if (matchedKey) { + 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..7e5709c744 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -205,6 +205,18 @@ 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("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({ From be656941d484eefd5f40c70b72282935a50a6459 Mon Sep 17 00:00:00 2001 From: agentHits Date: Thu, 20 Aug 2026 15:48:07 +0300 Subject: [PATCH 2/4] fix(google): restore signature on matched alternate key and bound input length pre-parse - Restore call.signature when call && matchedKey (not gated on ck), ensuring whitespace-padded wrapped arguments whose ck overflows 64 KiB still restore their signature if the parsed inner JSON is within bounds. - Bound argsObj.input.trim() to REPLAY_MAX_CANONICAL_ARGS_BYTES before calling JSON.parse in the custom tool replay unwrap, preventing oversized remote payloads from triggering unbounded allocations. - Add regression tests for oversized whitespace-wrapped input restoring correctly and oversized JSON payloads being rejected before parse (addressing review feedback from @Ingwannu on #2198). --- src/adapters/google-antigravity-replay.ts | 13 ++++++------ tests/google-antigravity-replay.test.ts | 25 +++++++++++++++++++++++ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index c2743be937..96240edb7f 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -709,7 +709,10 @@ export function applyAntigravityReplay(model: string, sessionId: string, content // 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") { + if ( + typeof argsObj.input === "string" + && utf8.encode(argsObj.input.trim()).byteLength <= REPLAY_MAX_CANONICAL_ARGS_BYTES + ) { try { const parsedInput = JSON.parse(argsObj.input); if (parsedInput && typeof parsedInput === "object") { @@ -724,12 +727,10 @@ export function applyAntigravityReplay(model: string, sessionId: string, content } } } - if (call && ck) { + if (call && matchedKey) { part.thoughtSignature = call.signature; - if (matchedKey) { - entry.byCall.delete(matchedKey); - entry.byCall.set(matchedKey, { ...call, touchedAtMs: now }); - } + entry.byCall.delete(matchedKey); + entry.byCall.set(matchedKey, { ...call, touchedAtMs: now }); touched = true; } } diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 7e5709c744..27d2da4b0d 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -217,6 +217,31 @@ describe("antigravity reasoning-replay cache", () => { 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")]); + // 65 KiB of leading whitespace makes the wrapped object exceed 64 KiB, so functionCallKey(fc.name, fc.args) is null/undefined. + // However, the parsed inner object is small and produces a valid alternate key. + const bigWhitespaceJson = " ".repeat(66 * 1024) + JSON.stringify({ cmd: "x" }); + const contents = [{ + role: "model", + parts: [{ functionCall: { name: "default_api:exec", args: { input: bigWhitespaceJson } } }], + }]; + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe("sig-whitespace-overflow"); + }); + + 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 skipped before parsing. + const oversizedPayload = JSON.stringify({ data: "a".repeat(100 * 1024) }); + const contents = [{ + role: "model", + parts: [{ functionCall: { name: "default_api:exec", args: { input: oversizedPayload } } }], + }]; + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + }); + test("claude models do not use the replay cache", () => { expect(antigravityUsesReplayCache("claude-opus-4.6")).toBe(false); expect(antigravityUsesReplayCache("gemini-3-pro")).toBe(true); From 99412cc85a65db31eca92947440d7ef7a03ee97c Mon Sep 17 00:00:00 2001 From: agentHits Date: Thu, 20 Aug 2026 16:00:46 +0300 Subject: [PATCH 3/4] fix(google): parse validated trimmed input directly and test parse boundaries - Compute trimmedInput and pass that exact validated string to JSON.parse, preventing large whitespace-padded prefixes/suffixes from reaching JSON.parse. - Add parse seam assertions proving that trimmed small payloads are parsed directly and that oversized valid payloads never reach JSON.parse (addressing feedback from @Ingwannu on #2198). --- src/adapters/google-antigravity-replay.ts | 16 ++++----- tests/google-antigravity-replay.test.ts | 43 ++++++++++++++++++----- 2 files changed, 43 insertions(+), 16 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 96240edb7f..5f12a4253e 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -709,12 +709,11 @@ export function applyAntigravityReplay(model: string, sessionId: string, content // 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" - && utf8.encode(argsObj.input.trim()).byteLength <= REPLAY_MAX_CANONICAL_ARGS_BYTES - ) { - try { - const parsedInput = JSON.parse(argsObj.input); + if (typeof argsObj.input === "string") { + const trimmedInput = argsObj.input.trim(); + if (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)) { @@ -722,8 +721,9 @@ export function applyAntigravityReplay(model: string, sessionId: string, content matchedKey = altKey; } } - } catch { - // not JSON, keep default + } catch { + // not JSON, keep default + } } } } diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 27d2da4b0d..f507b62e10 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -219,27 +219,54 @@ describe("antigravity reasoning-replay cache", () => { 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")]); - // 65 KiB of leading whitespace makes the wrapped object exceed 64 KiB, so functionCallKey(fc.name, fc.args) is null/undefined. - // However, the parsed inner object is small and produces a valid alternate key. - const bigWhitespaceJson = " ".repeat(66 * 1024) + JSON.stringify({ cmd: "x" }); + // 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 } } }], }]; - applyAntigravityReplay(MODEL, SESSION, contents); - expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe("sig-whitespace-overflow"); + 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 skipped before parsing. + // 100 KiB of valid JSON payload exceeds REPLAY_MAX_CANONICAL_ARGS_BYTES and must never reach JSON.parse. + let parseCalled = false; + const originalParse = JSON.parse; + JSON.parse = (text, reviver) => { + if (typeof text === "string" && text.includes('"data":')) { + parseCalled = true; + } + return originalParse(text, reviver); + }; const oversizedPayload = JSON.stringify({ data: "a".repeat(100 * 1024) }); const contents = [{ role: "model", parts: [{ functionCall: { name: "default_api:exec", args: { input: oversizedPayload } } }], }]; - applyAntigravityReplay(MODEL, SESSION, contents); - expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + try { + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + expect(parseCalled).toBe(false); + } finally { + JSON.parse = originalParse; + } }); test("claude models do not use the replay cache", () => { From 4e89471a7b9390fb8b1575fe528fe26282fa1c9b Mon Sep 17 00:00:00 2001 From: agentHits Date: Thu, 20 Aug 2026 16:29:59 +0300 Subject: [PATCH 4/4] fix(google): check string length before utf8.encode in replay unwrap - Guard trimmedInput.length <= REPLAY_MAX_CANONICAL_ARGS_BYTES before calling utf8.encode(trimmedInput), preventing large ASCII/serialized input from allocating proportional Uint8Array buffers on the request path. - Add regression test asserting that oversized inputs are rejected before TextEncoder.encode and JSON.parse (addressing review feedback from @Ingwannu and CodeRabbit on #2198). --- src/adapters/google-antigravity-replay.ts | 5 ++++- tests/google-antigravity-replay.test.ts | 12 +++++++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 5f12a4253e..3d574b54d9 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -711,7 +711,10 @@ export function applyAntigravityReplay(model: string, sessionId: string, content const argsObj = fc.args as Record; if (typeof argsObj.input === "string") { const trimmedInput = argsObj.input.trim(); - if (utf8.encode(trimmedInput).byteLength <= REPLAY_MAX_CANONICAL_ARGS_BYTES) { + 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") { diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index f507b62e10..8fab813cca 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -246,15 +246,23 @@ describe("antigravity reasoning-replay cache", () => { 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 never reach JSON.parse. + // 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", @@ -264,8 +272,10 @@ describe("antigravity reasoning-replay cache", () => { 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; } });