From b9ec99edcb3db30a38c49963523054a117904023 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:34:31 +0800 Subject: [PATCH 01/10] fix(google-antigravity): persist thought-signature replay across restarts The antigravity replay cache was process-local, so a tool-call continuation on /v1/chat/completions 400'd with 'missing a thought_signature' once the cache was lost (proxy restart, TTL expiry, or eviction pressure). The OpenAI-format surface has no field for a caller to carry the signature back, so the conversation was stuck. Snapshot the cache to antigravity-replay.json in the config directory: debounced atomic writes on mutation, lazy bounded load on first access, drop expired sessions at load, and a shutdown flush wired into drainAndShutdown. The CCA session id is deterministic (anchored on the first user message), so a restarted proxy re-derives the same key and re-injects the signatures. Claude-on-Antigravity is unaffected (it uses inline sanitization, not the replay cache). --- .../src/content/docs/ja/reference/adapters.md | 1 + .../src/content/docs/ko/reference/adapters.md | 1 + .../src/content/docs/reference/adapters.md | 2 + .../src/content/docs/ru/reference/adapters.md | 3 +- .../content/docs/zh-cn/reference/adapters.md | 3 +- src/adapters/google-antigravity-replay.ts | 173 +++++++++++++++++- src/server/lifecycle.ts | 13 +- tests/google-antigravity-replay.test.ts | 142 +++++++++++++- 8 files changed, 318 insertions(+), 20 deletions(-) diff --git a/docs-site/src/content/docs/ja/reference/adapters.md b/docs-site/src/content/docs/ja/reference/adapters.md index 1bf9f9bcf..6fc781dc3 100644 --- a/docs-site/src/content/docs/ja/reference/adapters.md +++ b/docs-site/src/content/docs/ja/reference/adapters.md @@ -75,6 +75,7 @@ HTTP リトライ ループの対象外です。 - システムプロンプト → `systemInstruction`；メッセージ → `contents[]`（assistant → `model`）；ツール → `functionDeclarations`。data URL 画像 → `inline_data`。 - Gemini が tool-call id を省略すると合成します。Vertex と Antigravity では実際の `thoughtSignature` 値を保存・再利用し、tool-result の継続ターンでも reasoning の連続性を保ちます。 + 署名キャッシュは設定ディレクトリにスナップショットされるため、プロキシ再起動後も継続ターンを維持できます。 ## `kiro` diff --git a/docs-site/src/content/docs/ko/reference/adapters.md b/docs-site/src/content/docs/ko/reference/adapters.md index 0ed83cc09..5995053e6 100644 --- a/docs-site/src/content/docs/ko/reference/adapters.md +++ b/docs-site/src/content/docs/ko/reference/adapters.md @@ -88,6 +88,7 @@ interface ProviderAdapter { `functionDeclarations`. data URL 이미지 → `inline_data`. - Gemini가 tool-call id를 생략하면 합성합니다. Vertex와 Antigravity에서는 실제 `thoughtSignature` 값을 보존하고 재사용해 tool-result 후속 턴에서도 reasoning 연속성을 유지합니다. + 서명 캐시는 설정 디렉터리에 스냅샷되므로 프록시 재시작 후에도 후속 턴이 유지됩니다. ## `kiro` diff --git a/docs-site/src/content/docs/reference/adapters.md b/docs-site/src/content/docs/reference/adapters.md index 7119ba7d7..4bfa661d7 100644 --- a/docs-site/src/content/docs/reference/adapters.md +++ b/docs-site/src/content/docs/reference/adapters.md @@ -92,6 +92,8 @@ of the HTTP retry loop. `functionDeclarations`. Data-URL images → `inline_data`. - Tool-call ids are synthesized when Gemini omits them. Vertex and Antigravity preserve and replay real `thoughtSignature` values so tool-result continuations retain Gemini reasoning continuity. + The signature cache is snapshotted to the config directory, so continuations also survive proxy + restarts. - **Inline image output:** when the model is one of the explicit image-capable chat IDs (`gemini-3.1-flash-image`, `gemini-2.0-flash-preview-image-generation`, or `gemini-3-pro-image-preview`), the adapter sends `responseModalities: ["TEXT", "IMAGE"]`. diff --git a/docs-site/src/content/docs/ru/reference/adapters.md b/docs-site/src/content/docs/ru/reference/adapters.md index a1f3950fa..b9269cedb 100644 --- a/docs-site/src/content/docs/ru/reference/adapters.md +++ b/docs-site/src/content/docs/ru/reference/adapters.md @@ -93,7 +93,8 @@ interface ProviderAdapter { инструменты → `functionDeclarations`. Изображения из data-URL → `inline_data`. - Идентификаторы вызовов инструментов синтезируются, когда Gemini их опускает. Vertex и Antigravity сохраняют и повторно передают настоящие значения `thoughtSignature`, чтобы непрерывность - рассуждений сохранялась после возврата результата инструмента. + рассуждений сохранялась после возврата результата инструмента. Кэш подписей сохраняется в каталог + конфигурации, поэтому последующие ходы переживают и перезапуск прокси. ## `kiro` diff --git a/docs-site/src/content/docs/zh-cn/reference/adapters.md b/docs-site/src/content/docs/zh-cn/reference/adapters.md index 9bde3ab29..c5511524a 100644 --- a/docs-site/src/content/docs/zh-cn/reference/adapters.md +++ b/docs-site/src/content/docs/zh-cn/reference/adapters.md @@ -81,7 +81,8 @@ interface ProviderAdapter { - 系统提示词 → `systemInstruction`；消息 → `contents[]`（assistant → `model`）；工具 → `functionDeclarations`；data URL 图像 → `inline_data`。 - Gemini 省略 tool-call id 时会合成 id。Vertex 与 Antigravity 会保留并重放真实 - `thoughtSignature`，使 tool-result 后续 turn 保持 reasoning continuity。 + `thoughtSignature`，使 tool-result 后续 turn 保持 reasoning continuity。签名缓存会快照到配置 + 目录，因此代理重启后后续 turn 仍可继续。 ## `kiro` diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 0334b930a..7212a9e54 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -1,4 +1,7 @@ import { createHash } from "node:crypto"; +import { chmodSync, existsSync, mkdirSync, readFileSync, statSync } from "node:fs"; +import { dirname, join } from "node:path"; +import { atomicWriteFileAsync, getConfigDir } from "../config"; import { enforceAppOwnedMemoryBudget } from "../lib/app-owned-memory"; /** @@ -38,6 +41,13 @@ export const ANTIGRAVITY_REPLAY_MAX_TOTAL_BYTES = 64 * 1024 * 1024; const REPLAY_MAX_SIGNATURE_BYTES = 64 * 1024; /** Fixed 64-hex outer key length, counted once per session entry. */ const REPLAY_SESSION_KEY_BYTES = 64; +const REPLAY_SNAPSHOT_FILE = "antigravity-replay.json"; +const REPLAY_SNAPSHOT_VERSION = 1; +const REPLAY_SNAPSHOT_DEBOUNCE_MS = 2_000; +/** Write bound for the durable snapshot (mirrors the responses-state cap). */ +const REPLAY_SNAPSHOT_MAX_BYTES = 24 * 1024 * 1024; +/** Refuse-to-parse ceiling for an existing snapshot file. */ +const REPLAY_SNAPSHOT_REFUSE_BYTES = 32 * 1024 * 1024; interface ReplayLimits { maxCallsPerSession: number; @@ -57,6 +67,138 @@ let replayLimits = { ...DEFAULT_REPLAY_LIMITS }; let replayBytes = 0; let replayOldestSessionKey: string | undefined; let replayOldestAt: number | null = null; +let replaySnapshotLoaded = false; +let replaySnapshotPersistTimer: ReturnType | null = null; +let replaySnapshotPersistGate: Promise = Promise.resolve(); + +function replaySnapshotPath(): string { + return join(getConfigDir(), REPLAY_SNAPSHOT_FILE); +} + +function loadReplaySnapshotEntry(key: string, value: unknown): void { + if (!/^[0-9a-f]{64}$/.test(key)) return; + if (!value || typeof value !== "object" || Array.isArray(value)) return; + const rec = value as { byCall?: unknown; expiresAtMs?: unknown }; + if (typeof rec.expiresAtMs !== "number" || !Number.isFinite(rec.expiresAtMs)) return; + // Expired sessions are dropped at load; a stale snapshot is self-healing. + if (rec.expiresAtMs <= Date.now()) return; + if (!Array.isArray(rec.byCall)) return; + const byCall = new Map(); + let bytes = REPLAY_SESSION_KEY_BYTES; + for (const pair of rec.byCall) { + if (!Array.isArray(pair) || pair.length !== 2 || typeof pair[0] !== "string") continue; + if (!/^[0-9a-f]{64}$/.test(pair[0])) continue; + const call = pair[1] as { signature?: unknown; sizeBytes?: unknown; touchedAtMs?: unknown } | null; + if (!call || typeof call !== "object" || Array.isArray(call)) continue; + if (typeof call.signature !== "string" || call.signature.length < MIN_SIGNATURE_LEN) continue; + if (typeof call.sizeBytes !== "number" || !Number.isSafeInteger(call.sizeBytes) || call.sizeBytes < 0) continue; + if (typeof call.touchedAtMs !== "number" || !Number.isFinite(call.touchedAtMs)) continue; + if (call.sizeBytes > replayLimits.maxSignatureBytes) continue; + byCall.set(pair[0], { + signature: call.signature, + sizeBytes: call.sizeBytes, + touchedAtMs: call.touchedAtMs, + }); + bytes += call.sizeBytes; + } + if (byCall.size === 0) return; + const entry: ReplayEntry = { byCall, bytes, expiresAtMs: rec.expiresAtMs, oldestAtMs: null }; + // Account BEFORE trimming: evictInnerCalls decrements the global byte count + // through deleteReplayCall, so the entry must already be on the books. + replayCache.set(key, entry); + replayBytes += entry.bytes; + // Same per-session caps as live writes, in case limits changed across versions. + evictInnerCalls(entry); + if (entry.byCall.size === 0) { + deleteReplaySession(key); + return; + } + refreshReplaySessionCandidate(key, entry); +} + +/** + * Lazy load of the durable snapshot on first cache access, so signatures observed + * before a proxy restart are available again once the session id re-derives to the + * same key (the session id is anchored on the first user message text). Load is + * best-effort: missing, corrupt, or oversized files start the cache empty. + */ +function ensureReplaySnapshotLoaded(): void { + if (replaySnapshotLoaded) return; + replaySnapshotLoaded = true; + try { + const path = replaySnapshotPath(); + if (!existsSync(path)) return; + const stat = statSync(path); + // Bound the read BEFORE parse: the 24 MiB write cap constrains snapshots this + // process wrote, not a pre-existing oversized file. + if (!stat.isFile() || stat.size > REPLAY_SNAPSHOT_REFUSE_BYTES) return; + const raw = JSON.parse(readFileSync(path, "utf8")) as { version?: unknown; sessions?: unknown }; + if (raw.version !== REPLAY_SNAPSHOT_VERSION || !Array.isArray(raw.sessions)) return; + for (const entry of raw.sessions) { + if (!Array.isArray(entry) || entry.length !== 2 || typeof entry[0] !== "string") continue; + loadReplaySnapshotEntry(entry[0], entry[1]); + } + } catch { + // Missing/corrupt snapshot: start empty. + } + // Load-time admission must respect the same global caps as live writes. + evictIfNeeded(); + enforceAppOwnedMemoryBudget(); +} + +function scheduleReplaySnapshotPersist(): void { + if (replaySnapshotPersistTimer) return; + replaySnapshotPersistTimer = setTimeout(() => { + void persistReplaySnapshotNow(); + }, REPLAY_SNAPSHOT_DEBOUNCE_MS); + (replaySnapshotPersistTimer as { unref?: () => void }).unref?.(); +} + +async function persistReplaySnapshotNow(): Promise { + if (replaySnapshotPersistTimer) { + clearTimeout(replaySnapshotPersistTimer); + replaySnapshotPersistTimer = null; + } + // Serialize writers so a flush and a debounced write cannot race on temps/ACL. + const previous = replaySnapshotPersistGate; + let release!: () => void; + replaySnapshotPersistGate = new Promise(resolve => { release = resolve; }); + await previous; + try { + const sessions: Array<[string, unknown]> = []; + let total = 0; + // Newest-first so the most recent sessions survive the snapshot byte cap. + for (const [key, entry] of [...replayCache].reverse()) { + const byCall = [...entry.byCall].map(([callKey, call]) => [ + callKey, + { signature: call.signature, sizeBytes: call.sizeBytes, touchedAtMs: call.touchedAtMs }, + ]); + const persistEntry: [string, unknown] = [key, { byCall, expiresAtMs: entry.expiresAtMs }]; + const size = Buffer.byteLength(JSON.stringify(persistEntry), "utf8"); + if (total + size > REPLAY_SNAPSHOT_MAX_BYTES) break; + total += size; + sessions.push(persistEntry); + } + sessions.reverse(); + mkdirSync(dirname(replaySnapshotPath()), { recursive: true, mode: 0o700 }); + try { chmodSync(dirname(replaySnapshotPath()), 0o700); } catch { /* best-effort (e.g. Windows) */ } + await atomicWriteFileAsync(replaySnapshotPath(), JSON.stringify({ version: REPLAY_SNAPSHOT_VERSION, sessions })); + } catch { + // The snapshot is a cache, not a source of truth: every failure is swallowed. + } finally { + release(); + } +} + +/** Flush any pending debounced snapshot write (graceful shutdown / deterministic tests). */ +export async function flushAntigravityReplay(): Promise { + if (replaySnapshotPersistTimer) { + await persistReplaySnapshotNow(); + return; + } + // No pending timer: still await any in-flight write so shutdown does not race it. + await replaySnapshotPersistGate; +} /** * Fixed-size identity for a (model, sessionId) pair: SHA-256 over @@ -372,6 +514,7 @@ 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; + ensureReplaySnapshotLoaded(); const now = Date.now(); deleteExpiredReplaySessionsThrottled(now); const key = replayKey(model, sessionId); @@ -418,8 +561,12 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts // The fixed session overhead can exceed the per-session cap on its own // (test-sized limits): an entry holding zero calls is unusable — drop it // instead of retaining an unevictable shell. - if (existing) deleteReplaySession(key); - else replayBytes -= REPLAY_SESSION_KEY_BYTES; + if (existing) { + deleteReplaySession(key); + scheduleReplaySnapshotPersist(); + } else { + replayBytes -= REPLAY_SESSION_KEY_BYTES; + } return; } entry.expiresAtMs = now + REPLAY_TTL_MS; @@ -427,6 +574,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts refreshReplaySessionCandidate(key, entry); evictIfNeeded(); enforceAppOwnedMemoryBudget(); + scheduleReplaySnapshotPersist(); } /** @@ -436,6 +584,7 @@ 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; + ensureReplaySnapshotLoaded(); const now = Date.now(); deleteExpiredReplaySessionsThrottled(now); const entry = replayCache.get(replayKey(model, sessionId)); @@ -461,13 +610,17 @@ export function applyAntigravityReplay(model: string, sessionId: string, content } } } - if (touched) refreshReplaySessionCandidate(replayKey(model, sessionId), entry); + if (touched) { + refreshReplaySessionCandidate(replayKey(model, sessionId), entry); + scheduleReplaySnapshotPersist(); + } return contents; } /** Drop the cache entry when upstream rejects a signature (clear-on-invalid). */ export function clearAntigravityReplay(model: string, sessionId: string): void { - deleteReplaySession(replayKey(model, sessionId)); + ensureReplaySnapshotLoaded(); + if (deleteReplaySession(replayKey(model, sessionId)) > 0) scheduleReplaySnapshotPersist(); } export function antigravityReplayMetrics(): { @@ -476,6 +629,7 @@ export function antigravityReplayMetrics(): { totalBytes: number; largestSessionBytes: number; } { + ensureReplaySnapshotLoaded(); let calls = 0; let largestSessionBytes = 0; for (const entry of replayCache.values()) { @@ -492,6 +646,7 @@ export function antigravityReplayRetainedStoreSnapshot(): { pinnedBytes: number; oldestAt: number | null; } { + ensureReplaySnapshotLoaded(); return { count: replayCache.size, bytes: replayBytes, @@ -502,7 +657,10 @@ export function antigravityReplayRetainedStoreSnapshot(): { } export function evictOldestAntigravityReplayForBudget(): number { - return replayOldestSessionKey === undefined ? 0 : deleteReplaySession(replayOldestSessionKey); + ensureReplaySnapshotLoaded(); + const removed = replayOldestSessionKey === undefined ? 0 : deleteReplaySession(replayOldestSessionKey); + if (removed > 0) scheduleReplaySnapshotPersist(); + return removed; } export function setAntigravityReplayLimitsForTests(limits?: Partial): void { @@ -512,6 +670,11 @@ export function setAntigravityReplayLimitsForTests(limits?: Partial setAntigravityReplayLimitsForTests()); +// Sandbox OPENCODEX_HOME: the replay cache now snapshots to disk, and these tests must +// never touch the real ~/.opencodex. +let replayTestHome: string; +const priorOpenCodexHome = process.env["OPENCODEX_HOME"]; + +beforeEach(() => { + replayTestHome = mkdtempSync(join(tmpdir(), "ocx-antigravity-replay-test-")); + process.env["OPENCODEX_HOME"] = replayTestHome; +}); + +afterEach(() => { + rmSync(replayTestHome, { recursive: true, force: true }); + if (priorOpenCodexHome === undefined) delete process.env["OPENCODEX_HOME"]; + else process.env["OPENCODEX_HOME"] = priorOpenCodexHome; +}); + const SIG = "sig-1234567890abcdef"; // >= 16 chars const MODEL = "gemini-3-pro"; const SESSION = "-12345"; -describe("antigravity reasoning-replay cache", () => { - // Signatures are keyed by functionCall identity (name + args), so observe/apply use functionCall parts. - const fcPart = (name: string, args: unknown, sig?: string, nested = false) => { - const part: Record = { functionCall: { name, args } }; - if (sig && nested) part.extra_content = { google: { thought_signature: sig } }; - else if (sig) part.thoughtSignature = sig; - return part; - }; +// Signatures are keyed by functionCall identity (name + args), so observe/apply use functionCall parts. +const fcPart = (name: string, args: unknown, sig?: string, nested = false) => { + const part: Record = { functionCall: { name, args } }; + if (sig && nested) part.extra_content = { google: { thought_signature: sig } }; + else if (sig) part.thoughtSignature = sig; + return part; +}; +describe("antigravity reasoning-replay cache", () => { test("observe then apply re-injects the signature onto the matching functionCall part", () => { observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); const contents = [ @@ -521,3 +542,106 @@ describe("antigravity replay fixed-size key identities", () => { expect(antigravityReplayMetrics().totalBytes).toBe(0); }); }); + +describe("durable antigravity replay snapshot", () => { + const snapshotPath = () => join(getConfigDir(), "antigravity-replay.json"); + + test("observed signatures survive a cache reset via the disk snapshot", async () => { + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + await flushAntigravityReplay(); + expect(existsSync(snapshotPath())).toBe(true); + // Proxy-restart simulation: wipe the in-process cache, then reload lazily. + setAntigravityReplayLimitsForTests(); + 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("flush writes the pending snapshot without waiting for the debounce", async () => { + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + await flushAntigravityReplay(); + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { version?: unknown; sessions?: unknown[] }; + expect(raw.version).toBe(1); + expect(raw.sessions).toHaveLength(1); + expect((raw.sessions![0] as unknown[])[0]).toBe(antigravityReplayKeyForTests(MODEL, SESSION)); + }); + + test("expired sessions are dropped when the snapshot is loaded", () => { + const now = Date.now(); + const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); + const staleCallKey = antigravityFunctionCallKeyForTests("get_stale", {}); + writeFileSync(snapshotPath(), JSON.stringify({ + version: 1, + sessions: [ + [antigravityReplayKeyForTests(MODEL, "-stale"), { + expiresAtMs: now - 1_000, + byCall: [[staleCallKey, { signature: SIG, sizeBytes: SIG.length, touchedAtMs: now }]], + }], + [antigravityReplayKeyForTests(MODEL, SESSION), { + expiresAtMs: now + 60_000, + byCall: [[callKey, { signature: SIG, sizeBytes: SIG.length, touchedAtMs: now }]], + }], + ], + })); + const staleContents = [{ role: "model", parts: [{ functionCall: { name: "get_stale", args: {} } }] }]; + applyAntigravityReplay(MODEL, "-stale", staleContents); + expect((staleContents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + 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("loaded sessions are trimmed to the current per-session caps", () => { + setAntigravityReplayLimitsForTests({ maxBytesPerSession: 300 }); + const now = Date.now(); + const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); + const trimKey = antigravityFunctionCallKeyForTests("get_y", {}); + writeFileSync(snapshotPath(), JSON.stringify({ + version: 1, + sessions: [ + [antigravityReplayKeyForTests(MODEL, SESSION), { + expiresAtMs: now + 60_000, + // Oldest (get_y) is evicted first on load: only get_x survives the cap. + byCall: [ + [trimKey, { signature: SIG, sizeBytes: 200, touchedAtMs: now }], + [callKey, { signature: SIG, sizeBytes: 200, touchedAtMs: now }], + ], + }], + ], + })); + 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); + const metrics = antigravityReplayMetrics(); + expect(metrics.sessions).toBe(1); + expect(metrics.calls).toBe(1); + // 300-byte cap: one 200-byte call plus the 64-byte fixed session key. + expect(metrics.totalBytes).toBe(64 + 200); + }); + + test("corrupt or unknown-version snapshots start the cache empty", () => { + writeFileSync(snapshotPath(), "{not valid json"); + 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).toBeUndefined(); + setAntigravityReplayLimitsForTests(); + writeFileSync(snapshotPath(), JSON.stringify({ version: 99, sessions: [] })); + applyAntigravityReplay(MODEL, SESSION, contents); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); + }); + + test("clear-on-invalid removes the session from the next snapshot", async () => { + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + clearAntigravityReplay(MODEL, SESSION); + await flushAntigravityReplay(); + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; + expect(raw.sessions).toHaveLength(0); + }); + + test("oversized snapshot files are refused before parsing", () => { + writeFileSync(snapshotPath(), "x".repeat(33 * 1024 * 1024)); + 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).toBeUndefined(); + }); +}); From 863aaf8d786a988a923e6bed03e9499af6d56293 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:14:41 +0800 Subject: [PATCH 02/10] fix(google-antigravity): recompute snapshot sizes, flush to latest mutation, surface persist failures --- src/adapters/google-antigravity-replay.ts | 77 +++++++++++----- tests/google-antigravity-replay.test.ts | 103 ++++++++++++++++++++-- 2 files changed, 153 insertions(+), 27 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 7212a9e54..9b9902b5d 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -1,7 +1,7 @@ import { createHash } from "node:crypto"; import { chmodSync, existsSync, mkdirSync, readFileSync, statSync } from "node:fs"; import { dirname, join } from "node:path"; -import { atomicWriteFileAsync, getConfigDir } from "../config"; +import { atomicWriteFileAsync, getConfigDir, type AtomicWriteAsyncTestSeam } from "../config"; import { enforceAppOwnedMemoryBudget } from "../lib/app-owned-memory"; /** @@ -70,6 +70,11 @@ let replayOldestAt: number | null = null; let replaySnapshotLoaded = false; let replaySnapshotPersistTimer: ReturnType | null = null; let replaySnapshotPersistGate: Promise = Promise.resolve(); +/** Mutation generation: bumped on every cache change that needs persisting. */ +let replayMutationGeneration = 0; +/** Generation of the data the last successful snapshot write actually captured. */ +let replayWrittenGeneration = 0; +let replaySnapshotWriteSeam: AtomicWriteAsyncTestSeam | undefined; function replaySnapshotPath(): string { return join(getConfigDir(), REPLAY_SNAPSHOT_FILE); @@ -88,18 +93,23 @@ function loadReplaySnapshotEntry(key: string, value: unknown): void { for (const pair of rec.byCall) { if (!Array.isArray(pair) || pair.length !== 2 || typeof pair[0] !== "string") continue; if (!/^[0-9a-f]{64}$/.test(pair[0])) continue; - const call = pair[1] as { signature?: unknown; sizeBytes?: unknown; touchedAtMs?: unknown } | null; + const call = pair[1] as { signature?: unknown; touchedAtMs?: unknown } | null; if (!call || typeof call !== "object" || Array.isArray(call)) continue; if (typeof call.signature !== "string" || call.signature.length < MIN_SIGNATURE_LEN) continue; - if (typeof call.sizeBytes !== "number" || !Number.isSafeInteger(call.sizeBytes) || call.sizeBytes < 0) continue; if (typeof call.touchedAtMs !== "number" || !Number.isFinite(call.touchedAtMs)) continue; - if (call.sizeBytes > replayLimits.maxSignatureBytes) continue; + // Never trust a serialized sizeBytes: a forged snapshot could claim a tiny + // size for a huge signature and bypass every byte cap. Recompute from the + // signature itself, exactly like the live write path. + const signatureBytes = utf8.encode(call.signature).byteLength; + if (signatureBytes > replayLimits.maxSignatureBytes) continue; + const callBytes = utf8.encode(pair[0]).byteLength + signatureBytes; + if (callBytes > replayLimits.maxBytesPerSession) continue; byCall.set(pair[0], { signature: call.signature, - sizeBytes: call.sizeBytes, + sizeBytes: callBytes, touchedAtMs: call.touchedAtMs, }); - bytes += call.sizeBytes; + bytes += callBytes; } if (byCall.size === 0) return; const entry: ReplayEntry = { byCall, bytes, expiresAtMs: rec.expiresAtMs, oldestAtMs: null }; @@ -146,10 +156,15 @@ function ensureReplaySnapshotLoaded(): void { enforceAppOwnedMemoryBudget(); } -function scheduleReplaySnapshotPersist(): void { +function markReplayDirty(): void { + replayMutationGeneration += 1; if (replaySnapshotPersistTimer) return; replaySnapshotPersistTimer = setTimeout(() => { - void persistReplaySnapshotNow(); + void persistReplaySnapshotNow().catch(() => { + // Redacted static message only: never echo the underlying error, which + // can carry paths or other environment details. + console.warn("[antigravity] replay snapshot persist failed; cached signatures will not survive a restart"); + }); }, REPLAY_SNAPSHOT_DEBOUNCE_MS); (replaySnapshotPersistTimer as { unref?: () => void }).unref?.(); } @@ -164,6 +179,7 @@ async function persistReplaySnapshotNow(): Promise { let release!: () => void; replaySnapshotPersistGate = new Promise(resolve => { release = resolve; }); await previous; + const writeGeneration = replayMutationGeneration; try { const sessions: Array<[string, unknown]> = []; let total = 0; @@ -171,7 +187,7 @@ async function persistReplaySnapshotNow(): Promise { for (const [key, entry] of [...replayCache].reverse()) { const byCall = [...entry.byCall].map(([callKey, call]) => [ callKey, - { signature: call.signature, sizeBytes: call.sizeBytes, touchedAtMs: call.touchedAtMs }, + { signature: call.signature, touchedAtMs: call.touchedAtMs }, ]); const persistEntry: [string, unknown] = [key, { byCall, expiresAtMs: entry.expiresAtMs }]; const size = Buffer.byteLength(JSON.stringify(persistEntry), "utf8"); @@ -182,9 +198,13 @@ async function persistReplaySnapshotNow(): Promise { sessions.reverse(); mkdirSync(dirname(replaySnapshotPath()), { recursive: true, mode: 0o700 }); try { chmodSync(dirname(replaySnapshotPath()), 0o700); } catch { /* best-effort (e.g. Windows) */ } - await atomicWriteFileAsync(replaySnapshotPath(), JSON.stringify({ version: REPLAY_SNAPSHOT_VERSION, sessions })); - } catch { - // The snapshot is a cache, not a source of truth: every failure is swallowed. + await atomicWriteFileAsync( + replaySnapshotPath(), + JSON.stringify({ version: REPLAY_SNAPSHOT_VERSION, sessions }), + undefined, + replaySnapshotWriteSeam, + ); + replayWrittenGeneration = writeGeneration; } finally { release(); } @@ -192,12 +212,17 @@ async function persistReplaySnapshotNow(): Promise { /** Flush any pending debounced snapshot write (graceful shutdown / deterministic tests). */ export async function flushAntigravityReplay(): Promise { - if (replaySnapshotPersistTimer) { - await persistReplaySnapshotNow(); - return; + // Persist until the durable generation catches up with the latest mutation: + // a change that lands while a writer is in flight must not be lost when + // shutdown exits right after the first write completes. + for (;;) { + if (replaySnapshotPersistTimer || replayMutationGeneration > replayWrittenGeneration) { + await persistReplaySnapshotNow(); + } + // No pending timer: still await any in-flight write so shutdown does not race it. + await replaySnapshotPersistGate; + if (replayMutationGeneration === replayWrittenGeneration && replaySnapshotPersistTimer === null) return; } - // No pending timer: still await any in-flight write so shutdown does not race it. - await replaySnapshotPersistGate; } /** @@ -563,7 +588,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts // instead of retaining an unevictable shell. if (existing) { deleteReplaySession(key); - scheduleReplaySnapshotPersist(); + markReplayDirty(); } else { replayBytes -= REPLAY_SESSION_KEY_BYTES; } @@ -574,7 +599,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts refreshReplaySessionCandidate(key, entry); evictIfNeeded(); enforceAppOwnedMemoryBudget(); - scheduleReplaySnapshotPersist(); + markReplayDirty(); } /** @@ -612,7 +637,7 @@ export function applyAntigravityReplay(model: string, sessionId: string, content } if (touched) { refreshReplaySessionCandidate(replayKey(model, sessionId), entry); - scheduleReplaySnapshotPersist(); + markReplayDirty(); } return contents; } @@ -620,7 +645,7 @@ export function applyAntigravityReplay(model: string, sessionId: string, content /** Drop the cache entry when upstream rejects a signature (clear-on-invalid). */ export function clearAntigravityReplay(model: string, sessionId: string): void { ensureReplaySnapshotLoaded(); - if (deleteReplaySession(replayKey(model, sessionId)) > 0) scheduleReplaySnapshotPersist(); + if (deleteReplaySession(replayKey(model, sessionId)) > 0) markReplayDirty(); } export function antigravityReplayMetrics(): { @@ -659,7 +684,7 @@ export function antigravityReplayRetainedStoreSnapshot(): { export function evictOldestAntigravityReplayForBudget(): number { ensureReplaySnapshotLoaded(); const removed = replayOldestSessionKey === undefined ? 0 : deleteReplaySession(replayOldestSessionKey); - if (removed > 0) scheduleReplaySnapshotPersist(); + if (removed > 0) markReplayDirty(); return removed; } @@ -668,6 +693,11 @@ export function setAntigravityReplayLimitsForTests(limits?: Partial { const now = Date.now(); const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); const trimKey = antigravityFunctionCallKeyForTests("get_y", {}); + // 100-byte signatures: one call (64+64+100 = 228) fits the 300-byte cap, + // two calls (392) do not, so the oldest is evicted on load. + const sigA = "a".repeat(100); + const sigB = "b".repeat(100); writeFileSync(snapshotPath(), JSON.stringify({ version: 1, sessions: [ @@ -603,20 +608,57 @@ describe("durable antigravity replay snapshot", () => { expiresAtMs: now + 60_000, // Oldest (get_y) is evicted first on load: only get_x survives the cap. byCall: [ - [trimKey, { signature: SIG, sizeBytes: 200, touchedAtMs: now }], - [callKey, { signature: SIG, sizeBytes: 200, touchedAtMs: now }], + [trimKey, { signature: sigB, sizeBytes: 200, touchedAtMs: now }], + [callKey, { signature: sigA, sizeBytes: 200, touchedAtMs: now }], ], }], ], })); 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); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe(sigA); const metrics = antigravityReplayMetrics(); expect(metrics.sessions).toBe(1); expect(metrics.calls).toBe(1); - // 300-byte cap: one 200-byte call plus the 64-byte fixed session key. - expect(metrics.totalBytes).toBe(64 + 200); + // Recomputed from the signature: 64-byte session key + 64-byte call key + 100-byte signature. + expect(metrics.totalBytes).toBe(64 + 64 + 100); + }); + + test("load recomputes call sizes from the signature and ignores serialized sizeBytes", () => { + const now = Date.now(); + const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); + writeFileSync(snapshotPath(), JSON.stringify({ + version: 1, + sessions: [ + [antigravityReplayKeyForTests(MODEL, SESSION), { + expiresAtMs: now + 60_000, + // Forged tiny sizeBytes must not shrink the accounted bytes. + byCall: [[callKey, { signature: SIG, sizeBytes: 1, touchedAtMs: now }]], + }], + [antigravityReplayKeyForTests(MODEL, "-forged"), { + expiresAtMs: now + 60_000, + // Oversized signature with a forged tiny sizeBytes must be dropped. + byCall: [[antigravityFunctionCallKeyForTests("get_y", {}), { + signature: "y".repeat(70_000), + sizeBytes: 1, + touchedAtMs: now, + }]], + }], + ], + })); + 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); + const metrics = antigravityReplayMetrics(); + expect(metrics.sessions).toBe(1); + // 64-byte session key + 64-byte call key + 20-byte signature. + expect(metrics.totalBytes).toBe(64 + 64 + SIG.length); + }); + + test("snapshot never serializes sizeBytes", async () => { + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + await flushAntigravityReplay(); + expect(readFileSync(snapshotPath(), "utf8")).not.toContain("sizeBytes"); }); test("corrupt or unknown-version snapshots start the cache empty", () => { @@ -644,4 +686,55 @@ describe("durable antigravity replay snapshot", () => { applyAntigravityReplay(MODEL, SESSION, contents); expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); }); + + test("flush waits out a blocked writer and persists mutations that land during it", async () => { + let firstWriteReached!: () => void; + const reached = new Promise(resolve => { firstWriteReached = resolve; }); + let firstWriteRelease!: () => void; + const release = new Promise(resolve => { firstWriteRelease = resolve; }); + let writes = 0; + setAntigravityReplayWriteSeamForTests({ + afterTempWrite: async () => { + writes += 1; + if (writes === 1) { + firstWriteReached(); + await release; + } + }, + }); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + const flushing = flushAntigravityReplay(); + await reached; + // A mutation lands while the first snapshot write is still blocked. + observeAntigravityReplay(MODEL, "-second", [fcPart("get_y", {}, SIG)]); + firstWriteRelease(); + await flushing; + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; + const keys = raw.sessions!.map(session => (session as unknown[])[0]); + expect(keys).toHaveLength(2); + expect(keys).toContain(antigravityReplayKeyForTests(MODEL, SESSION)); + expect(keys).toContain(antigravityReplayKeyForTests(MODEL, "-second")); + }); + + test("flush surfaces snapshot write failures for shutdown diagnostics", async () => { + setAntigravityReplayWriteSeamForTests({ + afterTempWrite: async () => { throw new Error("simulated disk failure"); }, + }); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + await expect(flushAntigravityReplay()).rejects.toThrow("simulated disk failure"); + }); + + test("background snapshot failures log a redacted warning and keep the service running", async () => { + const warn = spyOn(console, "warn"); + setAntigravityReplayWriteSeamForTests({ + afterTempWrite: async () => { throw new Error("simulated disk failure"); }, + }); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + // Let the 2s debounce timer fire and fail. + await Bun.sleep(2_100); + expect(warn.mock.calls.some(call => String(call[0]).includes("antigravity"))).toBe(true); + // The failure must not wedge the cache: a later mutation still works. + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_y", {}, SIG)]); + expect(antigravityReplayMetrics().calls).toBe(2); + }); }); From 2bb294bc0fc642184bb804c7b60fc0cdf4bcd5e1 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:23:50 +0800 Subject: [PATCH 03/10] docs(adapters): describe thoughtSignature values as opaque --- docs-site/src/content/docs/ja/reference/adapters.md | 2 +- docs-site/src/content/docs/ko/reference/adapters.md | 2 +- docs-site/src/content/docs/reference/adapters.md | 2 +- docs-site/src/content/docs/ru/reference/adapters.md | 2 +- docs-site/src/content/docs/zh-cn/reference/adapters.md | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs-site/src/content/docs/ja/reference/adapters.md b/docs-site/src/content/docs/ja/reference/adapters.md index 6fc781dc3..c0211e6d0 100644 --- a/docs-site/src/content/docs/ja/reference/adapters.md +++ b/docs-site/src/content/docs/ja/reference/adapters.md @@ -74,7 +74,7 @@ HTTP リトライ ループの対象外です。 - システムプロンプト → `systemInstruction`；メッセージ → `contents[]`（assistant → `model`）；ツール → `functionDeclarations`。data URL 画像 → `inline_data`。 -- Gemini が tool-call id を省略すると合成します。Vertex と Antigravity では実際の `thoughtSignature` 値を保存・再利用し、tool-result の継続ターンでも reasoning の連続性を保ちます。 +- Gemini が tool-call id を省略すると合成します。Vertex と Antigravity では不透明な `thoughtSignature` 値を保存・再利用し、tool-result の継続ターンでも reasoning の連続性を保ちます。 署名キャッシュは設定ディレクトリにスナップショットされるため、プロキシ再起動後も継続ターンを維持できます。 ## `kiro` diff --git a/docs-site/src/content/docs/ko/reference/adapters.md b/docs-site/src/content/docs/ko/reference/adapters.md index 5995053e6..1c9e75ee2 100644 --- a/docs-site/src/content/docs/ko/reference/adapters.md +++ b/docs-site/src/content/docs/ko/reference/adapters.md @@ -86,7 +86,7 @@ interface ProviderAdapter { - 시스템 프롬프트 → `systemInstruction`; 메시지 → `contents[]`(assistant → `model`); 툴 → `functionDeclarations`. data URL 이미지 → `inline_data`. -- Gemini가 tool-call id를 생략하면 합성합니다. Vertex와 Antigravity에서는 실제 +- Gemini가 tool-call id를 생략하면 합성합니다. Vertex와 Antigravity에서는 불투명한 `thoughtSignature` 값을 보존하고 재사용해 tool-result 후속 턴에서도 reasoning 연속성을 유지합니다. 서명 캐시는 설정 디렉터리에 스냅샷되므로 프록시 재시작 후에도 후속 턴이 유지됩니다. diff --git a/docs-site/src/content/docs/reference/adapters.md b/docs-site/src/content/docs/reference/adapters.md index 4bfa661d7..84b47b367 100644 --- a/docs-site/src/content/docs/reference/adapters.md +++ b/docs-site/src/content/docs/reference/adapters.md @@ -91,7 +91,7 @@ of the HTTP retry loop. - System prompt → `systemInstruction`; messages → `contents[]` (assistant → `model`); tools → `functionDeclarations`. Data-URL images → `inline_data`. - Tool-call ids are synthesized when Gemini omits them. Vertex and Antigravity preserve and replay - real `thoughtSignature` values so tool-result continuations retain Gemini reasoning continuity. + opaque `thoughtSignature` values so tool-result continuations retain Gemini reasoning continuity. The signature cache is snapshotted to the config directory, so continuations also survive proxy restarts. - **Inline image output:** when the model is one of the explicit image-capable chat IDs diff --git a/docs-site/src/content/docs/ru/reference/adapters.md b/docs-site/src/content/docs/ru/reference/adapters.md index b9269cedb..d18892c5f 100644 --- a/docs-site/src/content/docs/ru/reference/adapters.md +++ b/docs-site/src/content/docs/ru/reference/adapters.md @@ -92,7 +92,7 @@ interface ProviderAdapter { - Системный промпт → `systemInstruction`; сообщения → `contents[]` (assistant → `model`); инструменты → `functionDeclarations`. Изображения из data-URL → `inline_data`. - Идентификаторы вызовов инструментов синтезируются, когда Gemini их опускает. Vertex и Antigravity - сохраняют и повторно передают настоящие значения `thoughtSignature`, чтобы непрерывность + сохраняют и повторно передают непрозрачные значения `thoughtSignature`, чтобы непрерывность рассуждений сохранялась после возврата результата инструмента. Кэш подписей сохраняется в каталог конфигурации, поэтому последующие ходы переживают и перезапуск прокси. diff --git a/docs-site/src/content/docs/zh-cn/reference/adapters.md b/docs-site/src/content/docs/zh-cn/reference/adapters.md index c5511524a..a7c65825e 100644 --- a/docs-site/src/content/docs/zh-cn/reference/adapters.md +++ b/docs-site/src/content/docs/zh-cn/reference/adapters.md @@ -80,7 +80,7 @@ interface ProviderAdapter { - 系统提示词 → `systemInstruction`；消息 → `contents[]`（assistant → `model`）；工具 → `functionDeclarations`；data URL 图像 → `inline_data`。 -- Gemini 省略 tool-call id 时会合成 id。Vertex 与 Antigravity 会保留并重放真实 +- Gemini 省略 tool-call id 时会合成 id。Vertex 与 Antigravity 会保留并重放不透明 `thoughtSignature`，使 tool-result 后续 turn 保持 reasoning continuity。签名缓存会快照到配置 目录，因此代理重启后后续 turn 仍可继续。 From aa7da3f8ad4999d5936f9612711a18808e26f662 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:36:45 +0800 Subject: [PATCH 04/10] fix(google-antigravity): bound flush retry loop; restore console.warn spy in test --- src/adapters/google-antigravity-replay.ts | 5 ++++- tests/google-antigravity-replay.test.ts | 24 +++++++++++++---------- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 9b9902b5d..a904843ff 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -44,6 +44,9 @@ const REPLAY_SESSION_KEY_BYTES = 64; const REPLAY_SNAPSHOT_FILE = "antigravity-replay.json"; const REPLAY_SNAPSHOT_VERSION = 1; const REPLAY_SNAPSHOT_DEBOUNCE_MS = 2_000; +/** Upper bound on flush retries: mutations arriving faster than writes complete + * must not let shutdown hang; the last completed write is already on disk. */ +const REPLAY_FLUSH_MAX_ATTEMPTS = 8; /** Write bound for the durable snapshot (mirrors the responses-state cap). */ const REPLAY_SNAPSHOT_MAX_BYTES = 24 * 1024 * 1024; /** Refuse-to-parse ceiling for an existing snapshot file. */ @@ -215,7 +218,7 @@ export async function flushAntigravityReplay(): Promise { // Persist until the durable generation catches up with the latest mutation: // a change that lands while a writer is in flight must not be lost when // shutdown exits right after the first write completes. - for (;;) { + for (let attempt = 0; attempt < REPLAY_FLUSH_MAX_ATTEMPTS; attempt += 1) { if (replaySnapshotPersistTimer || replayMutationGeneration > replayWrittenGeneration) { await persistReplaySnapshotNow(); } diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 6b51f2236..1852fcfa2 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -726,15 +726,19 @@ describe("durable antigravity replay snapshot", () => { test("background snapshot failures log a redacted warning and keep the service running", async () => { const warn = spyOn(console, "warn"); - setAntigravityReplayWriteSeamForTests({ - afterTempWrite: async () => { throw new Error("simulated disk failure"); }, - }); - observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); - // Let the 2s debounce timer fire and fail. - await Bun.sleep(2_100); - expect(warn.mock.calls.some(call => String(call[0]).includes("antigravity"))).toBe(true); - // The failure must not wedge the cache: a later mutation still works. - observeAntigravityReplay(MODEL, SESSION, [fcPart("get_y", {}, SIG)]); - expect(antigravityReplayMetrics().calls).toBe(2); + try { + setAntigravityReplayWriteSeamForTests({ + afterTempWrite: async () => { throw new Error("simulated disk failure"); }, + }); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + // Let the 2s debounce timer fire and fail. + await Bun.sleep(2_100); + expect(warn.mock.calls.some(call => String(call[0]).includes("antigravity"))).toBe(true); + // The failure must not wedge the cache: a later mutation still works. + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_y", {}, SIG)]); + expect(antigravityReplayMetrics().calls).toBe(2); + } finally { + warn.mockRestore(); + } }); }); From eabaf27ff73b65afdd1b24a02bc5b4f3e115798c Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:51:38 +0800 Subject: [PATCH 05/10] fix(google-antigravity): persist load-time cleanups, order snapshot cap by activity, deterministic failure test --- src/adapters/google-antigravity-replay.ts | 52 +++++++++++++++++++---- tests/google-antigravity-replay.test.ts | 39 +++++++++++++++-- 2 files changed, 79 insertions(+), 12 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index a904843ff..a2a9111d7 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -29,6 +29,8 @@ interface ReplayEntry { bytes: number; expiresAtMs: number; oldestAtMs: number | null; + /** Most recent observe/apply activity; orders sessions under the snapshot cap. */ + lastActiveAtMs: number; } const MIN_SIGNATURE_LEN = 16; @@ -78,6 +80,8 @@ let replayMutationGeneration = 0; /** Generation of the data the last successful snapshot write actually captured. */ let replayWrittenGeneration = 0; let replaySnapshotWriteSeam: AtomicWriteAsyncTestSeam | undefined; +let replaySnapshotLoadDiscarded = false; +let replaySnapshotMaxBytes = REPLAY_SNAPSHOT_MAX_BYTES; function replaySnapshotPath(): string { return join(getConfigDir(), REPLAY_SNAPSHOT_FILE); @@ -86,10 +90,13 @@ function replaySnapshotPath(): string { function loadReplaySnapshotEntry(key: string, value: unknown): void { if (!/^[0-9a-f]{64}$/.test(key)) return; if (!value || typeof value !== "object" || Array.isArray(value)) return; - const rec = value as { byCall?: unknown; expiresAtMs?: unknown }; + const rec = value as { byCall?: unknown; expiresAtMs?: unknown; lastActiveAtMs?: unknown }; if (typeof rec.expiresAtMs !== "number" || !Number.isFinite(rec.expiresAtMs)) return; // Expired sessions are dropped at load; a stale snapshot is self-healing. - if (rec.expiresAtMs <= Date.now()) return; + if (rec.expiresAtMs <= Date.now()) { + replaySnapshotLoadDiscarded = true; + return; + } if (!Array.isArray(rec.byCall)) return; const byCall = new Map(); let bytes = REPLAY_SESSION_KEY_BYTES; @@ -114,14 +121,22 @@ function loadReplaySnapshotEntry(key: string, value: unknown): void { }); bytes += callBytes; } - if (byCall.size === 0) return; - const entry: ReplayEntry = { byCall, bytes, expiresAtMs: rec.expiresAtMs, oldestAtMs: null }; + if (byCall.size === 0) { + replaySnapshotLoadDiscarded = true; + return; + } + const lastActiveAtMs = typeof rec.lastActiveAtMs === "number" && Number.isFinite(rec.lastActiveAtMs) + ? rec.lastActiveAtMs + : rec.expiresAtMs; + const entry: ReplayEntry = { byCall, bytes, expiresAtMs: rec.expiresAtMs, oldestAtMs: null, lastActiveAtMs }; + const loadedCallCount = entry.byCall.size; // Account BEFORE trimming: evictInnerCalls decrements the global byte count // through deleteReplayCall, so the entry must already be on the books. replayCache.set(key, entry); replayBytes += entry.bytes; // Same per-session caps as live writes, in case limits changed across versions. evictInnerCalls(entry); + if (entry.byCall.size < loadedCallCount) replaySnapshotLoadDiscarded = true; if (entry.byCall.size === 0) { deleteReplaySession(key); return; @@ -138,6 +153,7 @@ function loadReplaySnapshotEntry(key: string, value: unknown): void { function ensureReplaySnapshotLoaded(): void { if (replaySnapshotLoaded) return; replaySnapshotLoaded = true; + replaySnapshotLoadDiscarded = false; try { const path = replaySnapshotPath(); if (!existsSync(path)) return; @@ -154,8 +170,13 @@ function ensureReplaySnapshotLoaded(): void { } catch { // Missing/corrupt snapshot: start empty. } + const sessionsAfterLoad = replayCache.size; // Load-time admission must respect the same global caps as live writes. evictIfNeeded(); + if (replayCache.size < sessionsAfterLoad) replaySnapshotLoadDiscarded = true; + // Persist the cleaned state once: expired/over-limit sessions must not stay + // on disk and be re-parsed (and re-dropped) after every restart. + if (replaySnapshotLoadDiscarded) markReplayDirty(); enforceAppOwnedMemoryBudget(); } @@ -186,15 +207,20 @@ async function persistReplaySnapshotNow(): Promise { try { const sessions: Array<[string, unknown]> = []; let total = 0; - // Newest-first so the most recent sessions survive the snapshot byte cap. - for (const [key, entry] of [...replayCache].reverse()) { + // Most-recently-active first so the sessions that survive the snapshot + // byte cap are the ones actually in use (Map order is insertion order). + for (const [key, entry] of [...replayCache].sort((a, b) => b[1].lastActiveAtMs - a[1].lastActiveAtMs)) { const byCall = [...entry.byCall].map(([callKey, call]) => [ callKey, { signature: call.signature, touchedAtMs: call.touchedAtMs }, ]); - const persistEntry: [string, unknown] = [key, { byCall, expiresAtMs: entry.expiresAtMs }]; + const persistEntry: [string, unknown] = [key, { + byCall, + expiresAtMs: entry.expiresAtMs, + lastActiveAtMs: entry.lastActiveAtMs, + }]; const size = Buffer.byteLength(JSON.stringify(persistEntry), "utf8"); - if (total + size > REPLAY_SNAPSHOT_MAX_BYTES) break; + if (total + size > replaySnapshotMaxBytes) break; total += size; sessions.push(persistEntry); } @@ -552,6 +578,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts bytes: REPLAY_SESSION_KEY_BYTES, expiresAtMs: 0, oldestAtMs: null, + lastActiveAtMs: 0, }; let inserted = false; let pendingThoughtSig: string | undefined; @@ -598,6 +625,7 @@ export function observeAntigravityReplay(model: string, sessionId: string, parts return; } entry.expiresAtMs = now + REPLAY_TTL_MS; + entry.lastActiveAtMs = now; replayCache.set(key, entry); refreshReplaySessionCandidate(key, entry); evictIfNeeded(); @@ -639,6 +667,7 @@ export function applyAntigravityReplay(model: string, sessionId: string, content } } if (touched) { + entry.lastActiveAtMs = now; refreshReplaySessionCandidate(replayKey(model, sessionId), entry); markReplayDirty(); } @@ -701,6 +730,11 @@ export function setAntigravityReplayWriteSeamForTests(seam: AtomicWriteAsyncTest replaySnapshotWriteSeam = seam; } +/** Test-only snapshot byte cap (mirrors the limits seam). */ +export function setAntigravityReplaySnapshotMaxBytesForTests(maxBytes: number): void { + replaySnapshotMaxBytes = maxBytes; +} + /** Test seam. */ export function __resetAntigravityReplayCache(): void { if (replaySnapshotPersistTimer) { @@ -716,4 +750,6 @@ export function __resetAntigravityReplayCache(): void { replayMutationGeneration = 0; replayWrittenGeneration = 0; replaySnapshotWriteSeam = undefined; + replaySnapshotLoadDiscarded = false; + replaySnapshotMaxBytes = REPLAY_SNAPSHOT_MAX_BYTES; } diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 1852fcfa2..fc0afb4e0 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -19,6 +19,7 @@ import { flushAntigravityReplay, observeAntigravityReplay, setAntigravityReplayLimitsForTests, + setAntigravityReplaySnapshotMaxBytesForTests, setAntigravityReplayWriteSeamForTests, sweepExpiredAntigravityReplay, } from "../src/adapters/google-antigravity-replay"; @@ -567,7 +568,7 @@ describe("durable antigravity replay snapshot", () => { expect((raw.sessions![0] as unknown[])[0]).toBe(antigravityReplayKeyForTests(MODEL, SESSION)); }); - test("expired sessions are dropped when the snapshot is loaded", () => { + test("expired sessions are dropped when the snapshot is loaded", async () => { const now = Date.now(); const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); const staleCallKey = antigravityFunctionCallKeyForTests("get_stale", {}); @@ -590,6 +591,31 @@ describe("durable antigravity replay snapshot", () => { 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); + // Load-time drops are persisted back, so the stale key leaves the file too. + await flushAntigravityReplay(); + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; + const keys = raw.sessions!.map(session => (session as unknown[])[0]); + expect(keys).not.toContain(antigravityReplayKeyForTests(MODEL, "-stale")); + expect(keys).toContain(antigravityReplayKeyForTests(MODEL, SESSION)); + }); + + test("snapshot cap keeps the most recently active sessions", async () => { + // A (two calls, ~899 serialized bytes) fits under the cap; A+B (~1,347) + // does not, so the cap must pick exactly one session by activity. + setAntigravityReplaySnapshotMaxBytesForTests(1_200); + const sigA = "a".repeat(200); + const sigB = "b".repeat(200); + // A is inserted first, B second, then A is re-observed: A is now the + // active session, so it must be the one retained under the byte cap. + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, sigA)]); + observeAntigravityReplay(MODEL, "-second", [fcPart("get_y", {}, sigB)]); + await Bun.sleep(2); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_z", { c: 2 }, sigA)]); + await flushAntigravityReplay(); + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; + const keys = raw.sessions!.map(session => (session as unknown[])[0]); + expect(keys).toContain(antigravityReplayKeyForTests(MODEL, SESSION)); + expect(keys).not.toContain(antigravityReplayKeyForTests(MODEL, "-second")); }); test("loaded sessions are trimmed to the current per-session caps", () => { @@ -725,14 +751,19 @@ describe("durable antigravity replay snapshot", () => { }); test("background snapshot failures log a redacted warning and keep the service running", async () => { - const warn = spyOn(console, "warn"); + let warned!: () => void; + const warningSeen = new Promise(resolve => { warned = resolve; }); + const warn = spyOn(console, "warn").mockImplementation((...args) => { + if (String(args[0]).includes("antigravity")) warned(); + }); try { setAntigravityReplayWriteSeamForTests({ afterTempWrite: async () => { throw new Error("simulated disk failure"); }, }); observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); - // Let the 2s debounce timer fire and fail. - await Bun.sleep(2_100); + // Deterministic: wait until the debounced write actually failed and + // warned (bounded), instead of racing a fixed sleep against the timer. + await Promise.race([warningSeen, Bun.sleep(5_000)]); expect(warn.mock.calls.some(call => String(call[0]).includes("antigravity"))).toBe(true); // The failure must not wedge the cache: a later mutation still works. observeAntigravityReplay(MODEL, SESSION, [fcPart("get_y", {}, SIG)]); From 22d93ab0361f231200857b23e18d27c8fd305e6b Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:59:03 +0800 Subject: [PATCH 06/10] fix(google-antigravity): collapse duplicate call keys during snapshot load --- src/adapters/google-antigravity-replay.ts | 6 +++++ tests/google-antigravity-replay.test.ts | 29 +++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index a2a9111d7..48d5fe5fb 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -114,6 +114,12 @@ function loadReplaySnapshotEntry(key: string, value: unknown): void { if (signatureBytes > replayLimits.maxSignatureBytes) continue; const callBytes = utf8.encode(pair[0]).byteLength + signatureBytes; if (callBytes > replayLimits.maxBytesPerSession) continue; + // A duplicated call key would overstate entry.bytes (the map keeps only the + // last value) and could evict valid sessions; keep the first occurrence. + if (byCall.has(pair[0])) { + replaySnapshotLoadDiscarded = true; + continue; + } byCall.set(pair[0], { signature: call.signature, sizeBytes: callBytes, diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index fc0afb4e0..ea0d16bfe 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -681,6 +681,35 @@ describe("durable antigravity replay snapshot", () => { expect(metrics.totalBytes).toBe(64 + 64 + SIG.length); }); + test("duplicate call keys in a snapshot are collapsed to one admitted call", async () => { + const now = Date.now(); + const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); + writeFileSync(snapshotPath(), JSON.stringify({ + version: 1, + sessions: [ + [antigravityReplayKeyForTests(MODEL, SESSION), { + expiresAtMs: now + 60_000, + byCall: [ + [callKey, { signature: SIG, touchedAtMs: now }], + [callKey, { signature: "sig-abcdef0123456789", touchedAtMs: now }], + ], + }], + ], + })); + 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); + const metrics = antigravityReplayMetrics(); + expect(metrics.calls).toBe(1); + // 64-byte session key + one 64-byte call key + 20-byte signature: no double count. + expect(metrics.totalBytes).toBe(64 + 64 + SIG.length); + // The duplicate is cleaned from disk too. + await flushAntigravityReplay(); + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; + const session = raw.sessions![0] as [unknown, { byCall?: unknown[] }]; + expect(session[1].byCall).toHaveLength(1); + }); + test("snapshot never serializes sizeBytes", async () => { observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); await flushAntigravityReplay(); From 8e171920d17bd8ef726b505075f33d2a61ab1d97 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:11:02 +0800 Subject: [PATCH 07/10] test(google-antigravity): verify duplicate-key cleanup at load before apply --- tests/google-antigravity-replay.test.ts | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index ea0d16bfe..04db7ff4e 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -696,18 +696,22 @@ describe("durable antigravity replay snapshot", () => { }], ], })); - 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); + // Metrics first: the load itself must collapse the duplicate and count once. const metrics = antigravityReplayMetrics(); expect(metrics.calls).toBe(1); // 64-byte session key + one 64-byte call key + 20-byte signature: no double count. expect(metrics.totalBytes).toBe(64 + 64 + SIG.length); - // The duplicate is cleaned from disk too. + // Load-time cleanup is persisted back to the snapshot. await flushAntigravityReplay(); const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; const session = raw.sessions![0] as [unknown, { byCall?: unknown[] }]; expect(session[1].byCall).toHaveLength(1); + // The surviving call still replays. + 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); + // Persist the apply-time touch so no debounced write is left pending. + await flushAntigravityReplay(); }); test("snapshot never serializes sizeBytes", async () => { From b52932e20aa47115b17bca2cb0e479cc3e8d883e Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:18:17 +0800 Subject: [PATCH 08/10] test(google-antigravity): assert warning redaction and load-only expiry cleanup --- tests/google-antigravity-replay.test.ts | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 04db7ff4e..4799aa992 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -586,17 +586,22 @@ describe("durable antigravity replay snapshot", () => { ], })); const staleContents = [{ role: "model", parts: [{ functionCall: { name: "get_stale", args: {} } }] }]; + // Load alone drops the expired session and persists the cleanup, before any + // replay-induced dirty state exists. + expect(antigravityReplayMetrics().sessions).toBe(1); + await flushAntigravityReplay(); + const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; + const keys = raw.sessions!.map(session => (session as unknown[])[0]); + expect(keys).not.toContain(antigravityReplayKeyForTests(MODEL, "-stale")); + expect(keys).toContain(antigravityReplayKeyForTests(MODEL, SESSION)); + // The surviving session still replays. applyAntigravityReplay(MODEL, "-stale", staleContents); expect((staleContents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); 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); - // Load-time drops are persisted back, so the stale key leaves the file too. + // Persist the apply-time touch so no debounced write is left pending. await flushAntigravityReplay(); - const raw = JSON.parse(readFileSync(snapshotPath(), "utf8")) as { sessions?: unknown[] }; - const keys = raw.sessions!.map(session => (session as unknown[])[0]); - expect(keys).not.toContain(antigravityReplayKeyForTests(MODEL, "-stale")); - expect(keys).toContain(antigravityReplayKeyForTests(MODEL, SESSION)); }); test("snapshot cap keeps the most recently active sessions", async () => { @@ -784,6 +789,7 @@ describe("durable antigravity replay snapshot", () => { }); test("background snapshot failures log a redacted warning and keep the service running", async () => { + const FAKE_SECRET = "FAKE_SECRET_antigravity_replay_12345"; let warned!: () => void; const warningSeen = new Promise(resolve => { warned = resolve; }); const warn = spyOn(console, "warn").mockImplementation((...args) => { @@ -791,13 +797,15 @@ describe("durable antigravity replay snapshot", () => { }); try { setAntigravityReplayWriteSeamForTests({ - afterTempWrite: async () => { throw new Error("simulated disk failure"); }, + afterTempWrite: async () => { throw new Error(`simulated disk failure ${FAKE_SECRET}`); }, }); observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); // Deterministic: wait until the debounced write actually failed and // warned (bounded), instead of racing a fixed sleep against the timer. await Promise.race([warningSeen, Bun.sleep(5_000)]); expect(warn.mock.calls.some(call => String(call[0]).includes("antigravity"))).toBe(true); + // The warning is redacted: error payload must never reach the log. + expect(warn.mock.calls.every(call => !call.some(arg => String(arg).includes(FAKE_SECRET)))).toBe(true); // The failure must not wedge the cache: a later mutation still works. observeAntigravityReplay(MODEL, SESSION, [fcPart("get_y", {}, SIG)]); expect(antigravityReplayMetrics().calls).toBe(2); From c9c14e0c907dcf5bbc7a7d7c340567bb365039ae Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:28:43 +0800 Subject: [PATCH 09/10] test(google-antigravity): cover UTF-8 replay byte limits --- tests/google-antigravity-replay.test.ts | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index 4799aa992..cd1427a01 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -658,19 +658,21 @@ describe("durable antigravity replay snapshot", () => { test("load recomputes call sizes from the signature and ignores serialized sizeBytes", () => { const now = Date.now(); const callKey = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); + const unicodeSig = "\u00e9".repeat(16); + const oversizedUnicodeSig = "\u{1f600}".repeat(16_500); writeFileSync(snapshotPath(), JSON.stringify({ version: 1, sessions: [ [antigravityReplayKeyForTests(MODEL, SESSION), { expiresAtMs: now + 60_000, - // Forged tiny sizeBytes must not shrink the accounted bytes. - byCall: [[callKey, { signature: SIG, sizeBytes: 1, touchedAtMs: now }]], + // Forged tiny sizeBytes must not shrink the accounted UTF-8 bytes. + byCall: [[callKey, { signature: unicodeSig, sizeBytes: 1, touchedAtMs: now }]], }], [antigravityReplayKeyForTests(MODEL, "-forged"), { expiresAtMs: now + 60_000, - // Oversized signature with a forged tiny sizeBytes must be dropped. + // The UTF-16 length is below 64 KiB, but the UTF-8 byte length is not. byCall: [[antigravityFunctionCallKeyForTests("get_y", {}), { - signature: "y".repeat(70_000), + signature: oversizedUnicodeSig, sizeBytes: 1, touchedAtMs: now, }]], @@ -679,11 +681,11 @@ describe("durable antigravity replay snapshot", () => { })); 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); + expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBe(unicodeSig); const metrics = antigravityReplayMetrics(); expect(metrics.sessions).toBe(1); - // 64-byte session key + 64-byte call key + 20-byte signature. - expect(metrics.totalBytes).toBe(64 + 64 + SIG.length); + // 64-byte session key + 64-byte call key + 32-byte UTF-8 signature. + expect(metrics.totalBytes).toBe(64 + 64 + new TextEncoder().encode(unicodeSig).byteLength); }); test("duplicate call keys in a snapshot are collapsed to one admitted call", async () => { From 3dad8a87c5f3e152d510cd6b96d2fb9d3d8525e1 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:58:50 +0800 Subject: [PATCH 10/10] fix: bound antigravity replay snapshot flush --- src/adapters/google-antigravity-replay.ts | 21 ++++++++-- tests/google-antigravity-replay.test.ts | 49 +++++++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/src/adapters/google-antigravity-replay.ts b/src/adapters/google-antigravity-replay.ts index 48d5fe5fb..71bfa8c38 100644 --- a/src/adapters/google-antigravity-replay.ts +++ b/src/adapters/google-antigravity-replay.ts @@ -53,6 +53,7 @@ const REPLAY_FLUSH_MAX_ATTEMPTS = 8; const REPLAY_SNAPSHOT_MAX_BYTES = 24 * 1024 * 1024; /** Refuse-to-parse ceiling for an existing snapshot file. */ const REPLAY_SNAPSHOT_REFUSE_BYTES = 32 * 1024 * 1024; +const REPLAY_FLUSH_INCOMPLETE_ERROR = "antigravity replay snapshot flush did not converge before the attempt limit"; interface ReplayLimits { maxCallsPerSession: number; @@ -212,7 +213,13 @@ async function persistReplaySnapshotNow(): Promise { const writeGeneration = replayMutationGeneration; try { const sessions: Array<[string, unknown]> = []; - let total = 0; + // JSON.stringify's empty document gives us the fixed envelope bytes. Each + // additional session contributes its own bytes plus one comma. + const emptySnapshotBytes = Buffer.byteLength(JSON.stringify({ + version: REPLAY_SNAPSHOT_VERSION, + sessions: [], + }), "utf8"); + let total = emptySnapshotBytes - 2; // Remove the empty [] from the envelope. // Most-recently-active first so the sessions that survive the snapshot // byte cap are the ones actually in use (Map order is insertion order). for (const [key, entry] of [...replayCache].sort((a, b) => b[1].lastActiveAtMs - a[1].lastActiveAtMs)) { @@ -226,16 +233,21 @@ async function persistReplaySnapshotNow(): Promise { lastActiveAtMs: entry.lastActiveAtMs, }]; const size = Buffer.byteLength(JSON.stringify(persistEntry), "utf8"); - if (total + size > replaySnapshotMaxBytes) break; - total += size; + const candidateBytes = total + (sessions.length === 0 ? 0 : 1) + size; + if (candidateBytes > replaySnapshotMaxBytes) break; + total = candidateBytes; sessions.push(persistEntry); } sessions.reverse(); + const snapshot = JSON.stringify({ version: REPLAY_SNAPSHOT_VERSION, sessions }); + if (Buffer.byteLength(snapshot, "utf8") > replaySnapshotMaxBytes) { + throw new Error("antigravity replay snapshot exceeds its size limit"); + } mkdirSync(dirname(replaySnapshotPath()), { recursive: true, mode: 0o700 }); try { chmodSync(dirname(replaySnapshotPath()), 0o700); } catch { /* best-effort (e.g. Windows) */ } await atomicWriteFileAsync( replaySnapshotPath(), - JSON.stringify({ version: REPLAY_SNAPSHOT_VERSION, sessions }), + snapshot, undefined, replaySnapshotWriteSeam, ); @@ -258,6 +270,7 @@ export async function flushAntigravityReplay(): Promise { await replaySnapshotPersistGate; if (replayMutationGeneration === replayWrittenGeneration && replaySnapshotPersistTimer === null) return; } + throw new Error(REPLAY_FLUSH_INCOMPLETE_ERROR); } /** diff --git a/tests/google-antigravity-replay.test.ts b/tests/google-antigravity-replay.test.ts index cd1427a01..e9bfbd7cd 100644 --- a/tests/google-antigravity-replay.test.ts +++ b/tests/google-antigravity-replay.test.ts @@ -623,6 +623,39 @@ describe("durable antigravity replay snapshot", () => { expect(keys).not.toContain(antigravityReplayKeyForTests(MODEL, "-second")); }); + test("snapshot cap includes the complete serialized document", async () => { + const now = Date.now(); + const callA = antigravityFunctionCallKeyForTests("get_x", { a: 1 }); + const callB = antigravityFunctionCallKeyForTests("get_y", { b: 2 }); + const sessionA: [string, unknown] = [antigravityReplayKeyForTests(MODEL, SESSION), { + byCall: [[callA, { signature: SIG, touchedAtMs: now }]], + expiresAtMs: now + 60_000, + lastActiveAtMs: now, + }]; + const sessionB: [string, unknown] = [antigravityReplayKeyForTests(MODEL, "-second"), { + byCall: [[callB, { signature: SIG, touchedAtMs: now }]], + expiresAtMs: now + 60_000, + lastActiveAtMs: now + 1, + }]; + // This admits both entries under the old entry-only calculation, but not + // the JSON envelope, commas, and entries under the corrected calculation. + const entryOnlyBytes = [sessionA, sessionB] + .reduce((total, session) => total + Buffer.byteLength(JSON.stringify(session), "utf8"), 0); + setAntigravityReplaySnapshotMaxBytesForTests(entryOnlyBytes); + writeFileSync(snapshotPath(), JSON.stringify({ version: 1, sessions: [sessionA, sessionB] })); + + // Loading is lazy; touching a loaded call marks the exact snapshot dirty. + applyAntigravityReplay(MODEL, SESSION, [{ + role: "model", + parts: [{ functionCall: { name: "get_x", args: { a: 1 } } }], + }]); + await flushAntigravityReplay(); + + const snapshot = readFileSync(snapshotPath(), "utf8"); + expect(Buffer.byteLength(snapshot, "utf8")).toBeLessThanOrEqual(entryOnlyBytes); + expect((JSON.parse(snapshot) as { sessions?: unknown[] }).sessions).toHaveLength(1); + }); + test("loaded sessions are trimmed to the current per-session caps", () => { setAntigravityReplayLimitsForTests({ maxBytesPerSession: 300 }); const now = Date.now(); @@ -782,6 +815,22 @@ describe("durable antigravity replay snapshot", () => { expect(keys).toContain(antigravityReplayKeyForTests(MODEL, "-second")); }); + test("flush rejects when mutations prevent convergence through the attempt bound", async () => { + let writes = 0; + setAntigravityReplayWriteSeamForTests({ + afterTempWrite: async () => { + writes += 1; + observeAntigravityReplay(MODEL, `-churn-${writes}`, [fcPart(`get_${writes}`, {}, SIG)]); + }, + }); + observeAntigravityReplay(MODEL, SESSION, [fcPart("get_x", { a: 1 }, SIG)]); + + await expect(flushAntigravityReplay()).rejects.toThrow( + "antigravity replay snapshot flush did not converge before the attempt limit", + ); + expect(writes).toBe(8); + }); + test("flush surfaces snapshot write failures for shutdown diagnostics", async () => { setAntigravityReplayWriteSeamForTests({ afterTempWrite: async () => { throw new Error("simulated disk failure"); },