From bbabde20d0aa5a14cf8fdb9b75f28753432c1e1c Mon Sep 17 00:00:00 2001 From: Ingwannu Date: Sat, 1 Aug 2026 06:52:20 +0000 Subject: [PATCH] bound Antigravity replay cache --- src/adapters/google-antigravity-replay.ts | 169 +++++++++++++++++++--- structure/04_transports-and-sidecars.md | 15 ++ tests/google-antigravity-replay.test.ts | 122 +++++++++++++++- 3 files changed, 282 insertions(+), 24 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 2fced61b0..9ad608067 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -1,3 +1,5 @@ +import { createHash } from "node:crypto"; + /** * Antigravity (Cloud Code Assist) thoughtSignature reasoning-replay cache. * @@ -10,21 +12,50 @@ * Claude-on-Antigravity uses inline signature sanitization instead (see google-antigravity-wire). */ +interface ReplayCall { + signature: string; + sizeBytes: number; +} + interface ReplayEntry { - /** thoughtSignature keyed by functionCall identity (name + canonical args). */ - byCall: Map; + /** thoughtSignature keyed by a fixed-size digest of functionCall identity. */ + byCall: Map; + sizeBytes: number; expiresAtMs: number; } const MIN_SIGNATURE_LEN = 16; const REPLAY_TTL_MS = 60 * 60 * 1000; // 1h const REPLAY_MAX_ENTRIES = 10_240; -const REPLAY_EVICT_BATCH = 128; +const REPLAY_MAX_CALLS_PER_SESSION = 256; +const REPLAY_MAX_BYTES_PER_SESSION = 2 * 1024 * 1024; +const REPLAY_MAX_TOTAL_BYTES = 32 * 1024 * 1024; +const REPLAY_MAX_SIGNATURE_BYTES = 64 * 1024; +const REPLAY_SWEEP_INTERVAL_MS = 60 * 1000; + +interface ReplayLimits { + maxEntries: number; + maxCallsPerSession: number; + maxBytesPerSession: number; + maxTotalBytes: number; + maxSignatureBytes: number; +} + +const DEFAULT_LIMITS: ReplayLimits = { + maxEntries: REPLAY_MAX_ENTRIES, + maxCallsPerSession: REPLAY_MAX_CALLS_PER_SESSION, + maxBytesPerSession: REPLAY_MAX_BYTES_PER_SESSION, + maxTotalBytes: REPLAY_MAX_TOTAL_BYTES, + maxSignatureBytes: REPLAY_MAX_SIGNATURE_BYTES, +}; const replayCache = new Map(); +let replayBytes = 0; +let lastSweepAtMs = 0; +let limits = { ...DEFAULT_LIMITS }; function replayKey(model: string, sessionId: string): string { - return `${model}::session:${sessionId}`; + return createHash("sha256").update(model).update("\0").update(sessionId).digest("hex"); } /** Recursively canonicalize a JSON value: object keys sorted, arrays preserved. */ @@ -36,7 +67,7 @@ function canonicalJson(value: unknown): string { return `{${entries.join(",")}}`; } -/** Stable identity for a functionCall part: name + recursively canonicalized args. */ +/** Stable fixed-size identity for a functionCall part: name + recursively canonicalized args. */ function functionCallKey(name: unknown, args: unknown): string | undefined { if (typeof name !== "string" || name.length === 0) return undefined; let argsKey = ""; @@ -45,7 +76,7 @@ function functionCallKey(name: unknown, args: unknown): string | undefined { } catch { argsKey = ""; } - return `${name}::${argsKey}`; + return createHash("sha256").update(name).update("\0").update(argsKey).digest("hex"); } function extractSignature(part: Record): string | undefined { @@ -57,12 +88,40 @@ function extractSignature(part: Record): string | undefined { return undefined; } -function evictIfNeeded(): void { - if (replayCache.size <= REPLAY_MAX_ENTRIES) return; - const oldest = [...replayCache.entries()] - .sort((a, b) => a[1].expiresAtMs - b[1].expiresAtMs) - .slice(0, REPLAY_EVICT_BATCH); - for (const [key] of oldest) replayCache.delete(key); +function deleteSession(key: string): void { + const entry = replayCache.get(key); + if (!entry) return; + replayBytes -= entry.sizeBytes; + if (replayBytes < 0) replayBytes = 0; + replayCache.delete(key); +} + +function deleteCall(entry: ReplayEntry, key: string): void { + const call = entry.byCall.get(key); + if (!call) return; + entry.byCall.delete(key); + entry.sizeBytes -= call.sizeBytes; + replayBytes -= call.sizeBytes; + if (entry.sizeBytes < 0) entry.sizeBytes = 0; + if (replayBytes < 0) replayBytes = 0; +} + +function sweepExpired(now: number, force = false): void { + if (!force && now - lastSweepAtMs < REPLAY_SWEEP_INTERVAL_MS) return; + lastSweepAtMs = now; + for (const [key, entry] of replayCache) { + if (entry.expiresAtMs <= now) deleteSession(key); + } +} + +function enforceLimits(now: number): void { + const underPressure = replayCache.size > limits.maxEntries || replayBytes > limits.maxTotalBytes; + sweepExpired(now, underPressure); + while (replayCache.size > limits.maxEntries || replayBytes > limits.maxTotalBytes) { + const oldest = replayCache.keys().next().value; + if (oldest === undefined) break; + deleteSession(oldest); + } } /** Gemini/Flash/Agent use the replay cache; Claude does not (inline sanitization instead). */ @@ -78,8 +137,15 @@ export function antigravityUsesReplayCache(model: string): boolean { */ export function observeAntigravityReplay(model: string, sessionId: string, parts: unknown[]): void { if (!antigravityUsesReplayCache(model) || !Array.isArray(parts) || parts.length === 0) return; + const now = Date.now(); const key = replayKey(model, sessionId); - const entry = replayCache.get(key) ?? { byCall: new Map(), expiresAtMs: 0 }; + let entry = replayCache.get(key); + if (entry?.expiresAtMs !== undefined && entry.expiresAtMs <= now) { + deleteSession(key); + entry = undefined; + } + sweepExpired(now); + entry ??= { byCall: new Map(), sizeBytes: 0, expiresAtMs: 0 }; let changed = false; for (const raw of parts) { if (!raw || typeof raw !== "object") continue; @@ -89,12 +155,35 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts const fc = part.functionCall as { name?: unknown; args?: unknown } | undefined; const ck = fc ? functionCallKey(fc.name, fc.args) : undefined; if (!ck) continue; // only function-call signatures are replayable by identity - if (entry.byCall.get(ck) !== sig) { entry.byCall.set(ck, sig); changed = true; } + const signatureBytes = Buffer.byteLength(sig, "utf8"); + const sizeBytes = Buffer.byteLength(ck, "utf8") + signatureBytes; + if (signatureBytes > limits.maxSignatureBytes || sizeBytes > limits.maxBytesPerSession) continue; + const existing = entry.byCall.get(ck); + if (existing?.signature === sig) continue; + if (existing) deleteCall(entry, ck); + entry.byCall.set(ck, { signature: sig, sizeBytes }); + entry.sizeBytes += sizeBytes; + replayBytes += sizeBytes; + changed = true; } if (!changed && replayCache.has(key)) return; - entry.expiresAtMs = Date.now() + REPLAY_TTL_MS; + if (!changed) return; + while ( + entry.byCall.size > limits.maxCallsPerSession + || entry.sizeBytes > limits.maxBytesPerSession + ) { + const oldest = entry.byCall.keys().next().value; + if (oldest === undefined) break; + deleteCall(entry, oldest); + } + if (entry.byCall.size === 0) { + deleteSession(key); + return; + } + entry.expiresAtMs = now + REPLAY_TTL_MS; + replayCache.delete(key); replayCache.set(key, entry); - evictIfNeeded(); + enforceLimits(now); } /** @@ -104,11 +193,14 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts */ export function applyAntigravityReplay(model: string, sessionId: string, contents: unknown[]): unknown[] { if (!antigravityUsesReplayCache(model) || !Array.isArray(contents)) return contents; - const entry = replayCache.get(replayKey(model, sessionId)); - if (!entry || entry.expiresAtMs <= Date.now()) { - if (entry) replayCache.delete(replayKey(model, sessionId)); + const now = Date.now(); + const key = replayKey(model, sessionId); + const entry = replayCache.get(key); + if (!entry || entry.expiresAtMs <= now) { + if (entry) deleteSession(key); return contents; } + let touched = false; for (const c of contents as { role?: string; parts?: unknown[] }[]) { if (!c || typeof c !== "object" || c.role !== "model" || !Array.isArray(c.parts)) continue; for (const raw of c.parts) { @@ -118,19 +210,52 @@ 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 sig = ck ? entry.byCall.get(ck) : undefined; - if (sig) part.thoughtSignature = sig; + const call = ck ? entry.byCall.get(ck) : undefined; + if (call && ck) { + part.thoughtSignature = call.signature; + entry.byCall.delete(ck); + entry.byCall.set(ck, call); + touched = true; + } } } + if (touched) { + replayCache.delete(key); + replayCache.set(key, entry); + } return contents; } /** Drop the cache entry when upstream rejects a signature (clear-on-invalid). */ export function clearAntigravityReplay(model: string, sessionId: string): void { - replayCache.delete(replayKey(model, sessionId)); + deleteSession(replayKey(model, sessionId)); +} + +/** Test-only observability for count/byte-bound regressions. */ +export function antigravityReplayMetricsForTests(): { + sessions: number; + calls: number; + totalBytes: number; + largestSessionBytes: number; +} { + let calls = 0; + let largestSessionBytes = 0; + for (const entry of replayCache.values()) { + calls += entry.byCall.size; + largestSessionBytes = Math.max(largestSessionBytes, entry.sizeBytes); + } + return { sessions: replayCache.size, calls, totalBytes: replayBytes, largestSessionBytes }; +} + +/** Test seam: lower limits without allocating production-sized fixtures. */ +export function setAntigravityReplayLimitsForTests(overrides: Partial | null): void { + __resetAntigravityReplayCache(); + limits = overrides ? { ...DEFAULT_LIMITS, ...overrides } : { ...DEFAULT_LIMITS }; } /** Test seam. */ export function __resetAntigravityReplayCache(): void { replayCache.clear(); + replayBytes = 0; + lastSweepAtMs = 0; } diff --git a/structure/04_transports-and-sidecars.md b/structure/04_transports-and-sidecars.md index 5a39c59e2..b490f01f3 100644 --- a/structure/04_transports-and-sidecars.md +++ b/structure/04_transports-and-sidecars.md @@ -439,6 +439,21 @@ combo whose remaining eligible targets use other providers. ## Transport inventory +Antigravity Gemini replay retains upstream `thoughtSignature` values only for later matching tool +calls. Session and function-call identities are SHA-256 digests rather than raw model/session/argument +strings. Each session retains at most 256 calls and 2 MiB, one signature is limited to 64 KiB, and +the cache retains at most 32 MiB globally. Both calls and sessions use least-recently-used eviction; +expired sessions are removed on exact access and by a sweep no more than once per minute unless +budget pressure requires it. + +[Decision Log] +- 목적과 의도: Preserve Gemini interleaved-thinking compatibility without retaining unbounded tool arguments or signatures. +- 기존 구현 및 제약 조건: The cache had only a 10,240-session count limit; raw canonical arguments were Map keys and every call inside a live session was unbounded. +- 검토한 주요 대안: Clear the session after every turn; retain raw identities behind a count-only LRU; hash identities and enforce per-call, per-session, and global byte budgets. +- 선택한 방식: Use fixed-size SHA-256 identities, exact UTF-8 signature accounting, nested LRU limits, and bounded TTL sweeping. +- 다른 대안 대신 이 방식을 선택한 이유: Clearing per turn breaks multi-step tool replay, while count-only limits still permit a few very large calls to dominate heap. +- 장점, 단점 및 영향: Ordinary long tool loops retain their newest 256 signatures; older calls in an extreme loop may miss replay and receive the upstream signature error instead of growing memory without bound. + The sections above cover the transports with load-bearing invariants. The rest of the transport surface is listed here so a maintainer can find the owner without grepping: diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 67da01182..e3de13959 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -1,14 +1,15 @@ import { afterEach, describe, expect, test } from "bun:test"; import { + antigravityReplayMetricsForTests, antigravityUsesReplayCache, applyAntigravityReplay, clearAntigravityReplay, observeAntigravityReplay, - __resetAntigravityReplayCache, + setAntigravityReplayLimitsForTests, } from "../src/adapters/google-antigravity-replay"; import { sanitizeAntigravityClaudeSignatures } from "../src/adapters/google-antigravity-wire"; -afterEach(() => __resetAntigravityReplayCache()); +afterEach(() => setAntigravityReplayLimitsForTests(null)); const SIG = "sig-1234567890abcdef"; // >= 16 chars const MODEL = "gemini-3-pro"; @@ -107,6 +108,123 @@ describe("antigravity reasoning-replay cache", () => { applyAntigravityReplay("claude-opus-4.6", SESSION, contents); expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); }); + + test("retains long sequential loops up to the 256-call production boundary", () => { + for (let i = 0; i < 260; i++) { + observeAntigravityReplay(MODEL, SESSION, [ + fcPart(`call-${i}`, { i }, `sig-${String(i).padStart(3, "0")}-${"x".repeat(16)}`), + ]); + } + expect(antigravityReplayMetricsForTests()).toMatchObject({ sessions: 1, calls: 256 }); + + const oldest = [{ role: "model", parts: [fcPart("call-0", { i: 0 })] }]; + const newest = [{ role: "model", parts: [fcPart("call-259", { i: 259 })] }]; + applyAntigravityReplay(MODEL, SESSION, oldest); + applyAntigravityReplay(MODEL, SESSION, newest); + expect((oldest[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + expect((newest[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toContain("sig-259"); + }); + + test("session byte pressure evicts the least-recently-used call", () => { + setAntigravityReplayLimitsForTests({ maxBytesPerSession: 180 }); + for (const name of ["one", "two", "three"]) { + observeAntigravityReplay(MODEL, SESSION, [fcPart(name, {}, `sig-${name}-${"x".repeat(16)}`)]); + } + + const metrics = antigravityReplayMetricsForTests(); + expect(metrics.totalBytes).toBeLessThanOrEqual(180); + expect(metrics.calls).toBe(2); + const contents = ["one", "two", "three"].map(name => ({ role: "model", parts: [fcPart(name, {})] })); + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + expect((contents[2].parts[0] as { thoughtSignature?: string }).thoughtSignature).toContain("sig-three"); + }); + + test("apply refreshes call and session LRU without extending TTL", () => { + setAntigravityReplayLimitsForTests({ maxEntries: 2, maxCallsPerSession: 2 }); + observeAntigravityReplay(MODEL, "one", [fcPart("old", {}, "sig-old-1234567890")]); + observeAntigravityReplay(MODEL, "one", [fcPart("new", {}, "sig-new-1234567890")]); + applyAntigravityReplay(MODEL, "one", [{ role: "model", parts: [fcPart("old", {})] }]); + observeAntigravityReplay(MODEL, "one", [fcPart("third", {}, "sig-third-12345678")]); + + const calls = ["old", "new", "third"].map(name => ({ role: "model", parts: [fcPart(name, {})] })); + applyAntigravityReplay(MODEL, "one", calls); + expect((calls[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toContain("sig-old"); + expect((calls[1].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + + observeAntigravityReplay(MODEL, "two", [fcPart("tool", {}, "sig-two-1234567890")]); + applyAntigravityReplay(MODEL, "one", [{ role: "model", parts: [fcPart("old", {})] }]); + observeAntigravityReplay(MODEL, "three", [fcPart("tool", {}, "sig-three-12345678")]); + expect(antigravityReplayMetricsForTests().sessions).toBe(2); + const evicted = [{ role: "model", parts: [fcPart("tool", {})] }]; + applyAntigravityReplay(MODEL, "two", evicted); + expect((evicted[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + }); + + test("global byte pressure evicts the least-recently-used session", () => { + setAntigravityReplayLimitsForTests({ maxTotalBytes: 210 }); + for (const session of ["session-one", "session-two", "session-three"]) { + observeAntigravityReplay(MODEL, session, [fcPart("tool", { session }, `sig-${session}-${"x".repeat(16)}`)]); + } + + expect(antigravityReplayMetricsForTests()).toMatchObject({ sessions: 2, calls: 2 }); + const first = [{ role: "model", parts: [fcPart("tool", { session: "session-one" })] }]; + const last = [{ role: "model", parts: [fcPart("tool", { session: "session-three" })] }]; + applyAntigravityReplay(MODEL, "session-one", first); + applyAntigravityReplay(MODEL, "session-three", last); + expect((first[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + expect((last[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toContain("session-three"); + }); + + test("oversized replacement preserves the prior valid signature", () => { + setAntigravityReplayLimitsForTests({ maxSignatureBytes: 32 }); + observeAntigravityReplay(MODEL, SESSION, [fcPart("tool", { a: 1 }, "sig-valid-1234567890")]); + const before = antigravityReplayMetricsForTests(); + observeAntigravityReplay(MODEL, SESSION, [fcPart("tool", { a: 1 }, "x".repeat(33))]); + + const contents = [{ role: "model", parts: [fcPart("tool", { a: 1 })] }]; + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature) + .toBe("sig-valid-1234567890"); + expect(antigravityReplayMetricsForTests().totalBytes).toBe(before.totalBytes); + }); + + test("duplicate replacement keeps exact byte accounting", () => { + observeAntigravityReplay(MODEL, SESSION, [fcPart("tool", { a: 1 }, "sig-first-123456789")]); + const firstBytes = antigravityReplayMetricsForTests().totalBytes; + observeAntigravityReplay(MODEL, SESSION, [fcPart("tool", { a: 1 }, "sig-other-123456789")]); + + const metrics = antigravityReplayMetricsForTests(); + expect(metrics).toMatchObject({ sessions: 1, calls: 1 }); + expect(metrics.totalBytes).toBe(firstBytes); + }); + + test("large canonical arguments retain only a fixed-size digest key", () => { + observeAntigravityReplay(MODEL, SESSION, [ + fcPart("huge", { text: "x".repeat(512 * 1024) }, "sig-fixed-key-123456"), + ]); + + const metrics = antigravityReplayMetricsForTests(); + expect(metrics).toMatchObject({ sessions: 1, calls: 1 }); + expect(metrics.totalBytes).toBeLessThan(256); + }); + + test("expired exact sessions release calls and byte accounting", () => { + const realNow = Date.now; + let current = 1_000; + Date.now = () => current; + try { + observeAntigravityReplay(MODEL, SESSION, [fcPart("tool", {}, SIG)]); + current += 60 * 60 * 1_000 + 1; + const contents = [{ role: "model", parts: [fcPart("tool", {})] }]; + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + expect(antigravityReplayMetricsForTests()) + .toEqual({ sessions: 0, calls: 0, totalBytes: 0, largestSessionBytes: 0 }); + } finally { + Date.now = realNow; + } + }); }); describe("claude-on-antigravity inline signature sanitization", () => {