@iislee/opencodex 2.11.0 → 2.31.0

package/src/adapters/google-antigravity-replay.ts

@@ -1,14 +1,19 @@
+import { createHash } from "node:crypto";
+import { chmodSync, existsSync, mkdirSync, readFileSync, statSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { atomicWriteFileAsync, getConfigDir, type AtomicWriteAsyncTestSeam } from "../config";
 import { enforceAppOwnedMemoryBudget } from "../lib/app-owned-memory";
 
 /**
- * Antigravity (Cloud Code Assist) thoughtSignature reasoning-replay cache.
+ * Google-family thoughtSignature reasoning-replay cache.
  *
  * Gemini-3 interleaved thinking is stateless upstream: each model content part carries a
  * `thoughtSignature` that MUST be echoed back on the matching part in the next request, or the
  * upstream rejects the turn (HTTP 400). We observe signatures on the response stream, cache them
  * per `model + session`, and re-inject them into the outgoing `request.contents` on the next turn.
  *
- * Mirrors CLIProxyAPI `internal/runtime/executor/antigravity_reasoning_replay.go`. Gemini-only;
+ * Mirrors CLIProxyAPI `internal/runtime/executor/antigravity_reasoning_replay.go` and is also used
+ * by Vertex with a transport/project/location-prefixed model identity. Gemini-only;
  * Claude-on-Antigravity uses inline signature sanitization instead (see google-antigravity-wire).
  */
 
@@ -24,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;
@@ -34,6 +41,18 @@ const REPLAY_MAX_CALLS_PER_SESSION = 256;
 export const ANTIGRAVITY_REPLAY_MAX_BYTES_PER_SESSION = 2 * 1024 * 1024;
 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;
+/** 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. */
+const REPLAY_SNAPSHOT_REFUSE_BYTES = 32 * 1024 * 1024;
 
 interface ReplayLimits {
   maxCallsPerSession: number;
@@ -53,30 +72,413 @@ let replayLimits = { ...DEFAULT_REPLAY_LIMITS };
 let replayBytes = 0;
 let replayOldestSessionKey: string | undefined;
 let replayOldestAt: number | null = null;
+let replaySnapshotLoaded = false;
+let replaySnapshotPersistTimer: ReturnType<typeof setTimeout> | null = null;
+let replaySnapshotPersistGate: Promise<void> = 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;
+let replaySnapshotLoadDiscarded = false;
+let replaySnapshotMaxBytes = REPLAY_SNAPSHOT_MAX_BYTES;
+
+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; 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()) {
+    replaySnapshotLoadDiscarded = true;
+    return;
+  }
+  if (!Array.isArray(rec.byCall)) return;
+  const byCall = new Map<string, ReplayCall>();
+  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; 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.touchedAtMs !== "number" || !Number.isFinite(call.touchedAtMs)) 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;
+    // 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,
+      touchedAtMs: call.touchedAtMs,
+    });
+    bytes += callBytes;
+  }
+  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;
+  }
+  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;
+  replaySnapshotLoadDiscarded = false;
+  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.
+  }
+  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();
+}
+
+function markReplayDirty(): void {
+  replayMutationGeneration += 1;
+  if (replaySnapshotPersistTimer) return;
+  replaySnapshotPersistTimer = setTimeout(() => {
+    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?.();
+}
+
+async function persistReplaySnapshotNow(): Promise<void> {
+  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<void>(resolve => { release = resolve; });
+  await previous;
+  const writeGeneration = replayMutationGeneration;
+  try {
+    const sessions: Array<[string, unknown]> = [];
+    // Account for the bytes actually written, not just the entries: the document
+    // is `{"version":N,"sessions":[...]}`, so the framing and the comma between
+    // entries count against the cap too. Summing per-entry sizes alone let the
+    // written file exceed replaySnapshotMaxBytes by a margin that grew with every
+    // additional session.
+    const prefix = `{"version":${JSON.stringify(REPLAY_SNAPSHOT_VERSION)},"sessions":[`;
+    const suffix = "]}";
+    const serialized: string[] = [];
+    let total = Buffer.byteLength(prefix, "utf8") + Buffer.byteLength(suffix, "utf8");
+    // 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,
+        lastActiveAtMs: entry.lastActiveAtMs,
+      }];
+      const encoded = JSON.stringify(persistEntry);
+      const size = Buffer.byteLength(encoded, "utf8") + (serialized.length > 0 ? 1 : 0);
+      if (total + size > replaySnapshotMaxBytes) break;
+      total += size;
+      serialized.push(encoded);
+      sessions.push(persistEntry);
+    }
+    serialized.reverse();
+    const document = `${prefix}${serialized.join(",")}${suffix}`;
+    // Defensive: the admitted entries are what we serialize, so this can only
+    // trip if the accounting above and the payload below ever drift apart.
+    if (Buffer.byteLength(document, "utf8") > replaySnapshotMaxBytes) {
+      throw new Error("Antigravity replay snapshot exceeded its configured byte cap.");
+    }
+    mkdirSync(dirname(replaySnapshotPath()), { recursive: true, mode: 0o700 });
+    try { chmodSync(dirname(replaySnapshotPath()), 0o700); } catch { /* best-effort (e.g. Windows) */ }
+    await atomicWriteFileAsync(
+      replaySnapshotPath(),
+      document,
+      undefined,
+      replaySnapshotWriteSeam,
+    );
+    replayWrittenGeneration = writeGeneration;
+  } finally {
+    release();
+  }
+}
+
+/** Flush any pending debounced snapshot write (graceful shutdown / deterministic tests). */
+export async function flushAntigravityReplay(): Promise<void> {
+  // 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 (let attempt = 0; attempt < REPLAY_FLUSH_MAX_ATTEMPTS; attempt += 1) {
+    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;
+  }
+  // Budget exhausted without convergence. Resolving here would tell
+  // drainAndShutdown() the snapshot is durable while the latest thought
+  // signature may never have reached disk, so shutdown diagnostics would claim
+  // a durability we cannot demonstrate. Reject instead, with fixed text: the
+  // shutdown path logs this message, so it must not carry session, model, or
+  // signature detail.
+  throw new Error("Antigravity replay snapshot flush did not converge.");
+}
+
+/**
+ * Fixed-size identity for a (model, sessionId) pair: SHA-256 over
+ * length-prefixed UTF-16 code units fed incrementally (no separator ambiguity
+ * — `("a\0b","c")` and `("a","b\0c")` derive different keys — and no raw
+ * model/session strings retained as Map keys, which the byte caps never
+ * counted).
+ */
+/**
+ * Injective string feed for key derivation: length-prefixed in CODE UNITS,
+ * then each code unit as two little-endian bytes. TextEncoder/UTF-8 would
+ * fold lone surrogates into U+FFFD, colliding distinct strings (e.g.
+ * "�" and "�") into the same key.
+ */
+function updateHashWithString(hash: ReturnType<typeof createHash>, value: string): void {
+  hash.update(String(value.length));
+  hash.update("\0");
+  const buf = Buffer.allocUnsafe(8192);
+  let offset = 0;
+  for (let index = 0; index < value.length; index += 1) {
+    buf.writeUInt16LE(value.charCodeAt(index), offset);
+    offset += 2;
+    if (offset === buf.length) {
+      hash.update(buf);
+      offset = 0;
+    }
+  }
+  if (offset > 0) hash.update(buf.subarray(0, offset));
+}
 
 function replayKey(model: string, sessionId: string): string {
-  return `${model}::session:${sessionId}`;
+  const hash = createHash("sha256");
+  updateHashWithString(hash, model);
+  updateHashWithString(hash, sessionId);
+  return hash.digest("hex");
+}
+
+/** Canonical output exceeding this budget is rejected DURING the walk — the
+ * pre-fix path materialized an unbounded canonical string before admission. */
+const REPLAY_MAX_CANONICAL_ARGS_BYTES = 64 * 1024;
+const CANONICAL_OVERFLOW = Symbol("canonical-overflow");
+
+/** Byte-identical output to the old recursive canonicalJson, written incrementally. */
+/**
+ * Key-count pre-check: every object key costs at least 4 canonical bytes
+ * (two quotes, colon, separator), so a wider object ALWAYS overflows the
+ * canonical budget — skip the sort and the walk. Object.keys allocation
+ * itself is linear and irreducible in JS (there is no streaming key API),
+ * but it is transient and never sorted or walked past this bound.
+ */
+const CANONICAL_MAX_KEYS_PER_OBJECT = REPLAY_MAX_CANONICAL_ARGS_BYTES / 4;
+
+/** Test-only scan instrumentation: proves overflow aborts the walk near the
+ * cap instead of scanning/materializing the whole input. */
+let canonicalScanUnitsForTests = 0;
+export function canonicalScanUnitsForTestsValue(): number {
+  return canonicalScanUnitsForTests;
+}
+export function resetCanonicalScanUnitsForTests(): void {
+  canonicalScanUnitsForTests = 0;
+}
+
+const MAX_CANONICAL_DEPTH = 128;
+
+function writeCanonicalJson(value: unknown, sink: (chunk: string) => void, depth = 0): void {
+  canonicalScanUnitsForTests += 1;
+  // Depth overflow is the same class as byte overflow: skip replay for this
+  // call instead of exhausting the stack on a pathological argument shape.
+  if (depth > MAX_CANONICAL_DEPTH) throw CANONICAL_OVERFLOW;
+  if (typeof value === "string") {
+    writeJsonStringEscaped(value, sink);
+    return;
+  }
+  if (value === null || typeof value !== "object") {
+    sink(JSON.stringify(value) ?? "null");
+    return;
+  }
+  if (Array.isArray(value)) {
+    sink("[");
+    for (let index = 0; index < value.length; index += 1) {
+      if (index > 0) sink(",");
+      // Array.prototype.map parity: holes produce NOTHING between the commas
+      // (old output `[1,,3]`), while an explicit undefined element is "null".
+      if (index in value) writeCanonicalJson(value[index], sink, depth + 1);
+    }
+    sink("]");
+    return;
+  }
+  const keys = Object.keys(value as Record<string, unknown>);
+  if (keys.length > CANONICAL_MAX_KEYS_PER_OBJECT) throw CANONICAL_OVERFLOW;
+  keys.sort();
+  sink("{");
+  keys.forEach((k, index) => {
+    if (index > 0) sink(",");
+    writeJsonStringEscaped(k, sink);
+    sink(":");
+    writeCanonicalJson((value as Record<string, unknown>)[k], sink, depth + 1);
+  });
+  sink("}");
+}
+
+/**
+ * JSON.stringify string escaping, streamed in small chunks so the budget can
+ * reject mid-string — calling JSON.stringify on a multi-MiB primitive would
+ * materialize its full escaped form before the sink could refuse it.
+ * Semantics mirror JSON.stringify for strings exactly (ES2019 JSON
+ * superset): quotes/backslash and control characters are escaped, LONE
+ * surrogates become \uXXXX, and valid surrogate pairs pass through raw.
+ */
+function writeJsonStringEscaped(value: string, sink: (chunk: string) => void): void {
+  sink('"');
+  let buffer = "";
+  for (const cp of value) {
+    canonicalScanUnitsForTests += 1;
+    const code = cp.codePointAt(0)!;
+    let escaped: string;
+    if (cp === '"') escaped = '\\"';
+    else if (cp === "\\") escaped = "\\\\";
+    else if (cp === "\b") escaped = "\\b";
+    else if (cp === "\f") escaped = "\\f";
+    else if (cp === "\n") escaped = "\\n";
+    else if (cp === "\r") escaped = "\\r";
+    else if (cp === "\t") escaped = "\\t";
+    else if (code < 0x20) escaped = `\\u${code.toString(16).padStart(4, "0")}`;
+    // Lone surrogates: JSON.stringify emits \uXXXX (a raw one would decode
+    // back as U+FFFD and collide with real U+FFFD content).
+    else if (code >= 0xd800 && code <= 0xdfff) escaped = `\\u${code.toString(16).padStart(4, "0")}`;
+    else escaped = cp;
+    buffer += escaped;
+    if (buffer.length >= 4096) {
+      sink(buffer);
+      buffer = "";
+    }
+  }
+  if (buffer.length > 0) sink(buffer);
+  sink('"');
 }
 
-/** Recursively canonicalize a JSON value: object keys sorted, arrays preserved. */
-function canonicalJson(value: unknown): string {
-  if (value === null || typeof value !== "object") return JSON.stringify(value) ?? "null";
-  if (Array.isArray(value)) return `[${value.map(canonicalJson).join(",")}]`;
-  const entries = Object.keys(value as Record<string, unknown>).sort()
-    .map(k => `${JSON.stringify(k)}:${canonicalJson((value as Record<string, unknown>)[k])}`);
-  return `{${entries.join(",")}}`;
+/** Bounded canonicalization: null on overflow (skip replay for that call). */
+function canonicalJsonBounded(value: unknown, maxBytes: number): string | null {
+  let written = 0;
+  const parts: string[] = [];
+  const sink = (chunk: string) => {
+    written += utf8.encode(chunk).byteLength;
+    if (written > maxBytes) throw CANONICAL_OVERFLOW;
+    parts.push(chunk);
+  };
+  try {
+    writeCanonicalJson(value, sink);
+  } catch (error) {
+    if (error === CANONICAL_OVERFLOW) return null;
+    throw error;
+  }
+  return parts.join("");
 }
 
-/** Stable identity for a functionCall part: name + recursively canonicalized args. */
+/**
+ * Stable identity for a functionCall part: fixed-size SHA-256 over
+ * length-prefixed name + canonical args. Overflow during canonicalization
+ * skips replay for that call (never materializes an unbounded string); other
+ * canonicalization failures keep the old name-only fallback semantics.
+ */
 function functionCallKey(name: unknown, args: unknown): string | undefined {
   if (typeof name !== "string" || name.length === 0) return undefined;
-  let argsKey = "";
+  let canonical: string | null;
   try {
-    argsKey = canonicalJson(args ?? {});
+    canonical = canonicalJsonBounded(args ?? {}, REPLAY_MAX_CANONICAL_ARGS_BYTES);
   } catch {
-    argsKey = "";
+    canonical = "";
   }
-  return `${name}::${argsKey}`;
+  if (canonical === null) return undefined;
+  const hash = createHash("sha256");
+  updateHashWithString(hash, name);
+  updateHashWithString(hash, canonical);
+  return hash.digest("hex");
+}
+
+/** Test-only key-derivation seam: the fixed-key regression cannot go red
+ * through snapshot.bytes (that metric never counted raw outer keys). */
+export function antigravityReplayKeyForTests(model: string, sessionId: string): string {
+  return replayKey(model, sessionId);
+}
+
+export function antigravityFunctionCallKeyForTests(name: unknown, args: unknown): string | undefined {
+  return functionCallKey(name, args);
+}
+
+/** Test-only bounded-canonicalization seam: proves mid-walk rejection without
+ * materializing the escaped form (allocation guard). */
+export function antigravityCanonicalJsonBoundedForTests(value: unknown, maxBytes: number): string | null {
+  return canonicalJsonBounded(value, maxBytes);
+}
+
+/** Test-only: the ACTUAL internal session keys, so tests can prove raw
+ * model/session strings are never retained as Map keys. */
+export function antigravityReplaySessionKeysForTests(): string[] {
+  return [...replayCache.keys()];
 }
 
 function extractSignature(part: Record<string, unknown>): string | undefined {
@@ -120,7 +522,31 @@ function refreshReplaySessionCandidate(key: string, entry: ReplayEntry): void {
 }
 
 function deleteExpiredReplaySessions(now: number): void {
-  for (const [key, entry] of replayCache) if (entry.expiresAtMs <= now) deleteReplaySession(key);
+  let deleted = false;
+  for (const [key, entry] of replayCache) {
+    if (entry.expiresAtMs > now) continue;
+    deleteReplaySession(key);
+    deleted = true;
+  }
+  // Expiry is a durable mutation too: rewrite the snapshot so opaque thought
+  // signatures do not remain at rest after their in-memory TTL has elapsed.
+  if (deleted) markReplayDirty();
+}
+
+/**
+ * The lazy per-call expiry scan is O(sessions); at the 10,240-session cap
+ * every observe/apply would rescan the whole map — O(n²) under load. The 60s
+ * state-store sweeper is already the periodic expiry authority, so lazy scans
+ * are throttled to at most one per interval (expired entries may linger a few
+ * extra seconds; TTL is fuzzy at that scale by design).
+ */
+const LAZY_SWEEP_INTERVAL_MS = 30_000;
+let lastLazySweepAt = Number.NEGATIVE_INFINITY;
+
+function deleteExpiredReplaySessionsThrottled(now: number): void {
+  if (now - lastLazySweepAt < LAZY_SWEEP_INTERVAL_MS) return;
+  lastLazySweepAt = now;
+  deleteExpiredReplaySessions(now);
 }
 
 export function sweepExpiredAntigravityReplay(now = Date.now()): number {
@@ -173,44 +599,83 @@ export function antigravityUsesReplayCache(model: string): boolean {
  * Observe a parsed CCA chunk's `candidates[0].content.parts` and record thought signatures keyed by
  * the functionCall identity (name + args). Accumulates across the whole session so a sequential
  * multi-step tool loop keeps EVERY prior call's signature, not just the latest part-index slot.
+ * A signature on a standalone thought part applies to the functionCall parts that follow it in
+ * the same array AND to later arrays of the same turn: streaming splits a thought part and its
+ * calls across SSE chunks, so `carriedThoughtSig` threads the still-unpaired signature from the
+ * previous chunk and the return value hands the remainder to the next one (#897, #2125). A call's
+ * own signature always takes precedence over a carried one.
  * `parts` is the already-unwrapped `response.candidates[0].content.parts`.
  */
-export function observeAntigravityReplay(model: string, sessionId: string, parts: unknown[]): void {
-  if (!antigravityUsesReplayCache(model) || !Array.isArray(parts) || parts.length === 0) return;
+export function observeAntigravityReplay(
+  model: string,
+  sessionId: string,
+  parts: unknown[],
+  carriedThoughtSig?: string,
+): string | undefined {
+  if (!antigravityUsesReplayCache(model) || !Array.isArray(parts) || parts.length === 0) return carriedThoughtSig;
+  ensureReplaySnapshotLoaded();
   const now = Date.now();
-  deleteExpiredReplaySessions(now);
+  deleteExpiredReplaySessionsThrottled(now);
   const key = replayKey(model, sessionId);
-  const entry = replayCache.get(key) ?? {
+  const existing = replayCache.get(key);
+  const entry = existing ?? {
     byCall: new Map<string, ReplayCall>(),
-    bytes: 0,
+    bytes: REPLAY_SESSION_KEY_BYTES,
     expiresAtMs: 0,
     oldestAtMs: null,
+    lastActiveAtMs: 0,
   };
   let inserted = false;
+  let pendingThoughtSig: string | undefined = carriedThoughtSig;
   for (const raw of parts) {
     if (!raw || typeof raw !== "object") continue;
     const part = raw as Record<string, unknown>;
     const sig = extractSignature(part);
-    if (!sig) continue;
     const fc = part.functionCall as { name?: unknown; args?: unknown } | undefined;
-    const ck = fc ? functionCallKey(fc.name, fc.args) : undefined;
+    if (!fc) {
+      // A signature on a standalone thought part pairs with the NEXT functionCall in this
+      // array: thought parts are stripped from replayed history, so the call part is the only
+      // carrier that survives (#897). Non-thought parts keep their own signature in history.
+      if (sig && part.thought === true) pendingThoughtSig = sig;
+      continue;
+    }
+    const callSig = sig ?? pendingThoughtSig; // a signature on the call part itself wins
+    if (!callSig) continue;
+    const ck = functionCallKey(fc.name, fc.args);
     if (!ck) continue; // only function-call signatures are replayable by identity
-    const signatureBytes = utf8.encode(sig).byteLength;
+    const signatureBytes = utf8.encode(callSig).byteLength;
     const sizeBytes = utf8.encode(ck).byteLength + signatureBytes;
     if (signatureBytes > replayLimits.maxSignatureBytes || sizeBytes > replayLimits.maxBytesPerSession) continue;
     deleteReplayCall(entry, ck);
-    entry.byCall.set(ck, { signature: sig, sizeBytes, touchedAtMs: now });
+    entry.byCall.set(ck, { signature: callSig, sizeBytes, touchedAtMs: now });
     entry.bytes += sizeBytes;
     replayBytes += sizeBytes;
     inserted = true;
   }
-  if (!inserted) return;
+  if (!inserted) return pendingThoughtSig;
+  // Charge the fixed outer key only when the session is actually stored.
+  if (!existing) replayBytes += REPLAY_SESSION_KEY_BYTES;
   evictInnerCalls(entry);
+  if (entry.byCall.size === 0) {
+    // 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);
+      markReplayDirty();
+    } else {
+      replayBytes -= REPLAY_SESSION_KEY_BYTES;
+    }
+    return pendingThoughtSig;
+  }
   entry.expiresAtMs = now + REPLAY_TTL_MS;
+  entry.lastActiveAtMs = now;
   replayCache.set(key, entry);
   refreshReplaySessionCandidate(key, entry);
   evictIfNeeded();
   enforceAppOwnedMemoryBudget();
+  markReplayDirty();
+  return pendingThoughtSig;
 }
 
 /**
@@ -220,8 +685,9 @@ 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();
-  deleteExpiredReplaySessions(now);
+  deleteExpiredReplaySessionsThrottled(now);
   const entry = replayCache.get(replayKey(model, sessionId));
   if (!entry) {
     return contents;
@@ -236,22 +702,54 @@ export function applyAntigravityReplay(model: string, sessionId: string, content
       if (!fc) continue;
       if (part.thoughtSignature !== undefined || part.thought_signature !== undefined) continue;
       const ck = functionCallKey(fc.name, fc.args);
-      const call = ck ? entry.byCall.get(ck) : undefined;
-      if (call && ck) {
+      let call = ck ? entry.byCall.get(ck) : undefined;
+      let matchedKey = ck;
+      if (!call && typeof fc.name === "string" && typeof fc.args === "object" && fc.args !== null) {
+        // Freeform / custom tool replay unwrap:
+        // The client replays custom_tool_call with arguments: { input: "..." }.
+        // Upstream was invoked with args: { input: "..." } or raw string or parsed JSON.
+        const argsObj = fc.args as Record<string, unknown>;
+        if (typeof argsObj.input === "string") {
+          const trimmedInput = argsObj.input.trim();
+          if (
+            trimmedInput.length <= REPLAY_MAX_CANONICAL_ARGS_BYTES
+            && utf8.encode(trimmedInput).byteLength <= REPLAY_MAX_CANONICAL_ARGS_BYTES
+          ) {
+            try {
+              const parsedInput = JSON.parse(trimmedInput);
+            if (parsedInput && typeof parsedInput === "object") {
+              const altKey = functionCallKey(fc.name, parsedInput);
+              if (altKey && entry.byCall.has(altKey)) {
+                call = entry.byCall.get(altKey);
+                matchedKey = altKey;
+              }
+            }
+            } catch {
+              // not JSON, keep default
+            }
+          }
+        }
+      }
+      if (call && matchedKey) {
         part.thoughtSignature = call.signature;
-        entry.byCall.delete(ck);
-        entry.byCall.set(ck, { ...call, touchedAtMs: now });
+        entry.byCall.delete(matchedKey);
+        entry.byCall.set(matchedKey, { ...call, touchedAtMs: now });
         touched = true;
       }
     }
   }
-  if (touched) refreshReplaySessionCandidate(replayKey(model, sessionId), entry);
+  if (touched) {
+    entry.lastActiveAtMs = now;
+    refreshReplaySessionCandidate(replayKey(model, sessionId), entry);
+    markReplayDirty();
+  }
   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) markReplayDirty();
 }
 
 export function antigravityReplayMetrics(): {
@@ -260,6 +758,7 @@ export function antigravityReplayMetrics(): {
   totalBytes: number;
   largestSessionBytes: number;
 } {
+  ensureReplaySnapshotLoaded();
   let calls = 0;
   let largestSessionBytes = 0;
   for (const entry of replayCache.values()) {
@@ -276,6 +775,7 @@ export function antigravityReplayRetainedStoreSnapshot(): {
   pinnedBytes: number;
   oldestAt: number | null;
 } {
+  ensureReplaySnapshotLoaded();
   return {
     count: replayCache.size,
     bytes: replayBytes,
@@ -286,7 +786,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) markReplayDirty();
+  return removed;
 }
 
 export function setAntigravityReplayLimitsForTests(limits?: Partial<ReplayLimits>): void {
@@ -294,10 +797,31 @@ export function setAntigravityReplayLimitsForTests(limits?: Partial<ReplayLimits
   replayLimits = limits ? { ...DEFAULT_REPLAY_LIMITS, ...limits } : { ...DEFAULT_REPLAY_LIMITS };
 }
 
+/** Test-only write seam (mirrors AtomicWriteAsyncTestSeam usage elsewhere). */
+export function setAntigravityReplayWriteSeamForTests(seam: AtomicWriteAsyncTestSeam | undefined): void {
+  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) {
+    clearTimeout(replaySnapshotPersistTimer);
+    replaySnapshotPersistTimer = null;
+  }
+  replaySnapshotLoaded = false;
+  lastLazySweepAt = Number.NEGATIVE_INFINITY;
   replayCache.clear();
   replayBytes = 0;
   replayOldestSessionKey = undefined;
   replayOldestAt = null;
+  replayMutationGeneration = 0;
+  replayWrittenGeneration = 0;
+  replaySnapshotWriteSeam = undefined;
+  replaySnapshotLoadDiscarded = false;
+  replaySnapshotMaxBytes = REPLAY_SNAPSHOT_MAX_BYTES;
 }