antigravity-claude-proxy 2.9.2 → 2.9.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "antigravity-claude-proxy",
-  "version": "2.9.2",
+  "version": "2.9.3",
   "description": "Proxy server to use Antigravity's Claude models with Claude Code CLI",
   "main": "src/index.js",
   "type": "module",

package/src/cloudcode/session-manager.js

@@ -1,36 +1,92 @@
 /**
  * Session Management for Cloud Code
  *
- * Handles session ID derivation for prompt caching continuity.
- * Session IDs are derived from the first user message to ensure
- * the same conversation uses the same session across turns.
+ * Session IDs give the Cloud Code backend a cache-affinity namespace.
+ * A per-conversation ID (stable hash of the first user message) keeps every
+ * conversation in its own namespace so interleaved conversations on the same
+ * account stop evicting each other's implicit-cache KV blocks. Requests
+ * without any user text fall back to a per-account process-lifetime ID.
  */
 
 import crypto from 'crypto';
 
 
-// Runtime storage for session IDs (per account)
-// This mimics the behavior of the binary which generates a session ID at startup
-// and keeps it for the process lifetime.
+// Runtime fallback storage for session IDs (per account)
 // Key: accountEmail, Value: sessionId
 const runtimeSessionStore = new Map();
 
 /**
- * Get or create a session ID for the given account.
- * 
- * The binary generates a session ID once at startup: `p.sessionID = rs() + Date.now()`.
- * Since our proxy is long-running, we simulate this "per-launch" behavior by storing
- * a generated ID in memory for each account.
+ * Extract the first meaningful user text from an Anthropic request.
+ * Returns null when no user message carries text.
+ */
+function extractFirstUserText(anthropicRequest) {
+    const messages = anthropicRequest?.messages;
+    if (!Array.isArray(messages)) return null;
+    for (const msg of messages) {
+        if (!msg || msg.role !== 'user') continue;
+        const content = msg.content;
+        let text = null;
+        if (typeof content === 'string') {
+            text = content;
+        } else if (Array.isArray(content)) {
+            const parts = [];
+            for (const block of content) {
+                if (block && block.type === 'text' && typeof block.text === 'string') {
+                    parts.push(block.text);
+                }
+            }
+            if (parts.length > 0) text = parts.join('\n');
+        }
+        if (text && text.trim().length > 0) {
+            // Hash the full text, never a prefix: clients such as Claude Code
+            // open the first user message with a long shared preamble
+            // (<system-reminder>, CLAUDE.md), so a truncated prefix maps
+            // every conversation in a project to the same session ID.
+            return text;
+        }
+    }
+    return null;
+}
+
+/**
+ * Derive a deterministic per-conversation session ID shaped like the binary's
+ * `uuid + Date.now()` format. The same conversation (same first user text,
+ * same account) always maps to the same ID across turns and proxy restarts.
+ */
+function deriveConversationSessionId(anthropicRequest, accountEmail) {
+    const firstUserText = extractFirstUserText(anthropicRequest);
+    if (!firstUserText) return null;
+    const digest = crypto
+        .createHash('sha256')
+        .update(`${accountEmail || 'anonymous'}\u0000${firstUserText}`)
+        .digest('hex');
+    const uuid =
+        `${digest.slice(0, 8)}-${digest.slice(8, 12)}-${digest.slice(12, 16)}-` +
+        `${digest.slice(16, 20)}-${digest.slice(20, 32)}`;
+    const tail = BigInt(`0x${digest.slice(32, 45)}`).toString();
+    return `${uuid}${tail}`;
+}
+
+/**
+ * Get or create a session ID for a request.
  *
- * - If the proxy restarts, the ID changes (matching binary/VS Code restart behavior).
- * - Within a running proxy instance, the ID is stable for that account.
- * - This enables prompt caching while using the EXACT random logic of the binary.
+ * Priority:
+ *   1. Per-conversation deterministic ID from the first user message —
+ *      restores cache continuity per conversation and isolates concurrent
+ *      conversations from each other's cache pressure.
+ *   2. Per-account process-lifetime ID when the request carries no user
+ *      text at all (previous behavior).
  *
- * @param {Object} anthropicRequest - The Anthropic-format request (unused for ID generation now)
- * @param {string} accountEmail - The account email to scope the session ID
+ * @param {Object} anthropicRequest - The Anthropic-format request
+ * @param {string} accountEmail - The account email to scope fallback IDs
  * @returns {string} A stable session ID string matching binary format
  */
 export function deriveSessionId(anthropicRequest, accountEmail) {
+    const conversationId = deriveConversationSessionId(anthropicRequest, accountEmail);
+    if (conversationId) {
+        return conversationId;
+    }
+
     if (!accountEmail) {
         // Fallback for requests without an account (should differ every time)
         return generateBinaryStyleId();

package/src/cloudcode/sse-streamer.js

@@ -55,7 +55,11 @@ export async function* streamSSEResponse(response, originalModel) {
                 const usage = innerResponse.usageMetadata;
                 if (usage) {
                     inputTokens = usage.promptTokenCount || inputTokens;
-                    outputTokens = usage.candidatesTokenCount || outputTokens;
+                    // Anthropic semantics: output_tokens covers ALL generated
+                    // tokens, thinking included (thoughtsTokenCount is billed as
+                    // output by the backend). candidatesTokenCount alone silently
+                    // undercounts thinking-heavy responses.
+                    outputTokens = ((usage.candidatesTokenCount || 0) + (usage.thoughtsTokenCount || 0)) || outputTokens;
                     cacheReadTokens = usage.cachedContentTokenCount || cacheReadTokens;
                 }
 
@@ -279,10 +283,20 @@ export async function* streamSSEResponse(response, originalModel) {
     }
 
     // Emit message_delta and message_stop
+    //
+    // Usage reporting: Google streams usageMetadata progressively; the
+    // early chunks carry an interim promptTokenCount and no
+    // cachedContentTokenCount yet. message_start therefore latches an
+    // inflated input_tokens (the nearly-full prompt minus zero cache) whenever
+    // the request is served mostly from cache. The final usageMetadata is
+    // authoritative, so correct input_tokens here: clients that take the last
+    // non-zero value (or read message_delta.usage.input_tokens) then see the
+    // true uncached input instead of a phantom near-100%-of-prompt figure.
     yield {
         type: 'message_delta',
         delta: { stop_reason: stopReason || 'end_turn', stop_sequence: null },
         usage: {
+            input_tokens: Math.max(0, inputTokens - cacheReadTokens),
             output_tokens: outputTokens,
             cache_read_input_tokens: cacheReadTokens,
             cache_creation_input_tokens: 0

package/src/format/response-converter.js

@@ -112,7 +112,8 @@ export function convertGoogleToAnthropic(googleResponse, model) {
         stop_sequence: null,
         usage: {
             input_tokens: promptTokens - cachedTokens,
-            output_tokens: usageMetadata.candidatesTokenCount || 0,
+            // Thinking tokens are generated (and billed) output — include them.
+            output_tokens: (usageMetadata.candidatesTokenCount || 0) + (usageMetadata.thoughtsTokenCount || 0),
             cache_read_input_tokens: cachedTokens,
             cache_creation_input_tokens: 0
         }

package/src/format/signature-cache.js

@@ -1,82 +1,270 @@
 /**
- * Signature Cache
- * In-memory cache for Gemini thoughtSignatures
+ * Signature Cache — persistent, LRU-bounded
  *
- * Gemini models require thoughtSignature on tool calls, but Claude Code
- * strips non-standard fields. This cache stores signatures by tool_use_id
- * so they can be restored in subsequent requests.
+ * Gemini models require thoughtSignature on tool calls and thinking blocks,
+ * but Anthropic-protocol clients strip non-standard fields. This cache stores
+ * signatures by tool_use_id so they can be restored in subsequent requests.
+ *
+ * Signatures have no fixed time-based expiry — they stay valid for the lifetime
+ * of active conversations, and dropping them mid-conversation silently rewrites
+ * conversation history on replay, which busts the upstream implicit-cache prefix
+ * and drops thinking blocks. The store is persisted to disk and reloaded at boot,
+ * bounded by an LRU cap so memory and on-disk size stay bounded over months of use.
  *
  * Also caches thinking block signatures with model family for cross-model
  * compatibility checking.
+ *
+ * Persistence: JSON file, atomic tmp+rename writes, mode 0600, debounced
+ * async writes (plus a periodic safety flush and synchronous flush on SIGTERM/SIGINT/
+ * exit). All persistence failures fail open: the in-memory cache keeps
+ * working and the proxy never errors because of store I/O.
  */
 
-import { GEMINI_SIGNATURE_CACHE_TTL_MS, MIN_SIGNATURE_LENGTH } from '../constants.js';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { MIN_SIGNATURE_LENGTH } from '../constants.js';
+
+const toolSignatureCache = new Map();      // toolUseId -> signature (insertion-ordered LRU)
+const thinkingSignatureCache = new Map();   // signature -> modelFamily (insertion-ordered LRU)
+
+const STORE_VERSION = 1;
+const DEFAULT_MAX_TOOL_SIGNATURES = 20_000;
+const DEFAULT_MAX_THINKING_SIGNATURES = 10_000;
+const FLUSH_DEBOUNCE_MS = 5000;
+const PERIODIC_FLUSH_MS = 30_000;
+const MAX_SILENT_STORE_ERRORS = 20;
+
+function parsePositiveIntEnv(name, fallback) {
+    const raw = process.env[name];
+    if (!raw) return fallback;
+    const n = Number.parseInt(raw, 10);
+    return Number.isFinite(n) && n > 0 ? n : fallback;
+}
+
+const maxToolSignatures = parsePositiveIntEnv('ANTIGRAVITY_SIGNATURE_CACHE_MAX_TOOL', DEFAULT_MAX_TOOL_SIGNATURES);
+const maxThinkingSignatures = parsePositiveIntEnv('ANTIGRAVITY_SIGNATURE_CACHE_MAX_THINK', DEFAULT_MAX_THINKING_SIGNATURES);
+
+function lruSet(map, key, value, maxEntries) {
+    if (map.has(key)) {
+        map.delete(key);
+    }
+    map.set(key, value);
+    while (map.size > maxEntries) {
+        const oldestKey = map.keys().next().value;
+        if (oldestKey === undefined) break;
+        map.delete(oldestKey);
+    }
+}
+
+function lruGet(map, key) {
+    if (!map.has(key)) return null;
+    const value = map.get(key);
+    map.delete(key);
+    map.set(key, value);
+    return value;
+}
+
+function resolveStorePath() {
+    const fromEnv = process.env.ANTIGRAVITY_SIGNATURE_STORE;
+    if (fromEnv && fromEnv.trim()) return fromEnv;
+    return path.join(os.homedir(), '.config', 'antigravity-proxy', 'signature-store.json');
+}
+
+let storePath = resolveStorePath();
+let flushTimer = null;
+let periodicTimer = null;
+let dirty = false;
+let storeErrorCount = 0;
+let writeInFlight = null;
+let storeGeneration = 0;
 
-const signatureCache = new Map();
-const thinkingSignatureCache = new Map();
+function reportStoreError(context, error) {
+    if (storeErrorCount < MAX_SILENT_STORE_ERRORS) {
+        storeErrorCount++;
+        // Console only: this module is imported before the logger is configured.
+        console.warn(`[SignatureCache] ${context} failed: ${error?.message ?? error}`);
+    }
+}
+
+function loadFromDisk() {
+    try {
+        const raw = fs.readFileSync(storePath, 'utf8');
+        if (!raw.trim()) return;
+        const parsed = JSON.parse(raw);
+        if (parsed?.version !== STORE_VERSION) return;
+        let prunedOnLoad = false;
+        if (parsed.tool && typeof parsed.tool === 'object') {
+            const entries = Object.entries(parsed.tool);
+            if (entries.length > maxToolSignatures) prunedOnLoad = true;
+            for (const [id, sig] of entries) {
+                if (typeof id === 'string' && typeof sig === 'string') {
+                    lruSet(toolSignatureCache, id, sig, maxToolSignatures);
+                }
+            }
+        }
+        if (parsed.think && typeof parsed.think === 'object') {
+            const entries = Object.entries(parsed.think);
+            if (entries.length > maxThinkingSignatures) prunedOnLoad = true;
+            for (const [sig, family] of entries) {
+                if (typeof sig === 'string' && typeof family === 'string') {
+                    lruSet(thinkingSignatureCache, sig, family, maxThinkingSignatures);
+                }
+            }
+        }
+        if (prunedOnLoad) {
+            scheduleFlush();
+        }
+    } catch (error) {
+        if (error?.code !== 'ENOENT') {
+            reportStoreError('load', error);
+        }
+        // Corrupt or unreadable store: start empty rather than failing boot.
+    }
+}
+
+function serializeStore() {
+    return JSON.stringify({
+        version: STORE_VERSION,
+        tool: Object.fromEntries(toolSignatureCache),
+        think: Object.fromEntries(thinkingSignatureCache),
+    });
+}
 
 /**
- * Store a signature for a tool_use_id
+ * Synchronous flush, used only on process exit/signals where async I/O
+ * would never complete. Bumps the store generation so any async write still
+ * in flight discards its (older) snapshot instead of renaming over this one.
+ */
+function writeStoreSync() {
+    if (!dirty) return;
+    storeGeneration++;
+    try {
+        const data = serializeStore();
+        fs.mkdirSync(path.dirname(storePath), { recursive: true });
+        const tmp = `${storePath}.tmp-sync-${process.pid}`;
+        fs.writeFileSync(tmp, data, { mode: 0o600 });
+        fs.renameSync(tmp, storePath);
+        dirty = false;
+    } catch (error) {
+        reportStoreError('flush', error);
+    }
+}
+
+/**
+ * Normal-operation flush. File I/O runs asynchronously so a store near the
+ * LRU cap (tens of MB) never blocks the event loop on disk writes; only the
+ * JSON serialization stays synchronous. At most one write is in flight.
+ */
+function writeStoreAsync() {
+    if (!dirty || writeInFlight) return;
+    let data;
+    try {
+        data = serializeStore();
+    } catch (error) {
+        reportStoreError('flush', error);
+        return;
+    }
+    dirty = false;
+    const generation = storeGeneration;
+    const targetPath = storePath;
+    const tmp = `${targetPath}.tmp-${process.pid}`;
+    let failed = false;
+    writeInFlight = (async () => {
+        try {
+            await fs.promises.mkdir(path.dirname(targetPath), { recursive: true });
+            await fs.promises.writeFile(tmp, data, { mode: 0o600 });
+            if (generation !== storeGeneration) {
+                // Cache was cleared or flushed synchronously meanwhile.
+                await fs.promises.rm(tmp, { force: true });
+                return;
+            }
+            await fs.promises.rename(tmp, targetPath);
+        } catch (error) {
+            failed = true;
+            dirty = true;
+            reportStoreError('flush', error);
+        }
+    })().finally(() => {
+        writeInFlight = null;
+        // Changes made during the write get flushed promptly; after a failed
+        // write, leave the retry to the periodic flush instead of spinning.
+        if (dirty && !failed) scheduleFlush();
+    });
+}
+
+function scheduleFlush() {
+    dirty = true;
+    if (flushTimer) return;
+    flushTimer = setTimeout(() => {
+        flushTimer = null;
+        writeStoreAsync();
+    }, FLUSH_DEBOUNCE_MS);
+    if (typeof flushTimer.unref === 'function') flushTimer.unref();
+}
+
+function removeStoreFile() {
+    try {
+        fs.rmSync(storePath, { force: true });
+    } catch (error) {
+        reportStoreError('remove', error);
+    }
+}
+
+loadFromDisk();
+
+periodicTimer = setInterval(() => writeStoreAsync(), PERIODIC_FLUSH_MS);
+if (typeof periodicTimer.unref === 'function') periodicTimer.unref();
+
+for (const signal of ['SIGTERM', 'SIGINT', 'exit']) {
+    process.on(signal, () => writeStoreSync());
+}
+
+/**
+ * Store a signature for a tool_use_id. Bounded by an LRU cap (no fixed TTL):
+ * replayed signatures stay byte-stable for active conversations.
  * @param {string} toolUseId - The tool use ID
  * @param {string} signature - The thoughtSignature to cache
  */
 export function cacheSignature(toolUseId, signature) {
     if (!toolUseId || !signature) return;
-    signatureCache.set(toolUseId, {
-        signature,
-        timestamp: Date.now()
-    });
+    const prev = toolSignatureCache.get(toolUseId);
+    lruSet(toolSignatureCache, toolUseId, signature, maxToolSignatures);
+    if (prev === signature) return;
+    scheduleFlush();
 }
 
 /**
- * Get a cached signature for a tool_use_id
+ * Get a cached signature for a tool_use_id.
  * @param {string} toolUseId - The tool use ID
- * @returns {string|null} The cached signature or null if not found/expired
+ * @returns {string|null} The cached signature or null if never seen
  */
 export function getCachedSignature(toolUseId) {
     if (!toolUseId) return null;
-    const entry = signatureCache.get(toolUseId);
-    if (!entry) return null;
-
-    // Check TTL
-    if (Date.now() - entry.timestamp > GEMINI_SIGNATURE_CACHE_TTL_MS) {
-        signatureCache.delete(toolUseId);
-        return null;
-    }
-
-    return entry.signature;
+    return lruGet(toolSignatureCache, toolUseId);
 }
 
 /**
- * Cache a thinking block signature with its model family
+ * Cache a thinking block signature with its model family. Bounded by an LRU cap.
  * @param {string} signature - The thinking signature to cache
  * @param {string} modelFamily - The model family ('claude' or 'gemini')
  */
 export function cacheThinkingSignature(signature, modelFamily) {
     if (!signature || signature.length < MIN_SIGNATURE_LENGTH) return;
-    thinkingSignatureCache.set(signature, {
-        modelFamily,
-        timestamp: Date.now()
-    });
+    const prev = thinkingSignatureCache.get(signature);
+    lruSet(thinkingSignatureCache, signature, modelFamily, maxThinkingSignatures);
+    if (prev === modelFamily) return;
+    scheduleFlush();
 }
 
 /**
- * Get the cached model family for a thinking signature
+ * Get the cached model family for a thinking signature.
  * @param {string} signature - The signature to look up
- * @returns {string|null} 'claude', 'gemini', or null if not found/expired
+ * @returns {string|null} 'claude', 'gemini', or null if never seen
  */
 export function getCachedSignatureFamily(signature) {
     if (!signature) return null;
-    const entry = thinkingSignatureCache.get(signature);
-    if (!entry) return null;
-
-    // Check TTL
-    if (Date.now() - entry.timestamp > GEMINI_SIGNATURE_CACHE_TTL_MS) {
-        thinkingSignatureCache.delete(signature);
-        return null;
-    }
-
-    return entry.modelFamily;
+    return lruGet(thinkingSignatureCache, signature);
 }
 
 /**
@@ -85,4 +273,36 @@ export function getCachedSignatureFamily(signature) {
  */
 export function clearThinkingSignatureCache() {
     thinkingSignatureCache.clear();
+    scheduleFlush();
+}
+
+/**
+ * Flush pending changes to disk now and wait for the write to finish.
+ * The debounced flush normally does this; exported for tests and callers
+ * that need the store on disk at a known point.
+ * @returns {Promise<void>}
+ */
+export async function flushSignatureStore() {
+    if (flushTimer) {
+        clearTimeout(flushTimer);
+        flushTimer = null;
+    }
+    if (writeInFlight) await writeInFlight;
+    writeStoreAsync();
+    if (writeInFlight) await writeInFlight;
+}
+
+/**
+ * Clear both signature caches and delete the persisted store.
+ */
+export function clearAllSignatureCaches() {
+    toolSignatureCache.clear();
+    thinkingSignatureCache.clear();
+    if (flushTimer) {
+        clearTimeout(flushTimer);
+        flushTimer = null;
+    }
+    dirty = false;
+    storeGeneration++; // any in-flight async write must not resurrect the store
+    removeStoreFile();
 }

package/src/format/thinking-utils.js

@@ -656,6 +656,21 @@ function stripInvalidThinkingBlocks(messages, targetFamily = null) {
  * @returns {Array<Object>} Modified messages with synthetic messages injected
  */
 export function closeToolLoopForThinking(messages, targetFamily = null) {
+    // Escape hatch (Gemini only): skip the synthetic-message injection.
+    // Google's backend accepts raw mid-tool-loop and interrupted-tool
+    // histories natively (verified 2026-09-03: identical answers with and
+    // without injection), so for Gemini targets the fabricated messages are
+    // pure downside — a fake "[Continue]" becomes the model's latest user
+    // instruction and the count-bearing text is unstable across requests.
+    // Claude targets keep the recovery: the Anthropic API hard-rejects
+    // unclosed tool_use ids, so the interrupted-tool synthesis stays needed.
+    // Set ANTIGRAVITY_DISABLE_THINKING_RECOVERY=1 to activate.
+    if (
+        process.env.ANTIGRAVITY_DISABLE_THINKING_RECOVERY === '1' &&
+        targetFamily === 'gemini'
+    ) {
+        return messages;
+    }
     const state = analyzeConversationState(messages);
 
     // Handle neither tool loop nor interrupted tool

package/src/server.js

@@ -17,7 +17,7 @@ const __dirname = path.dirname(__filename);
 import { forceRefresh } from './auth/token-extractor.js';
 import { REQUEST_BODY_LIMIT, BUILTIN_MODEL_MAPPINGS } from './constants.js';
 import { AccountManager } from './account-manager/index.js';
-import { clearThinkingSignatureCache } from './format/signature-cache.js';
+import { clearAllSignatureCaches } from './format/signature-cache.js';
 import { formatDuration } from './utils/helpers.js';
 import { logger } from './utils/logger.js';
 import usageStats from './modules/usage-stats.js';
@@ -217,13 +217,13 @@ app.post('/', (req, res) => {
 });
 
 /**
- * Test endpoint - Clear thinking signature cache
+ * Test endpoint - Clear all signature caches (in-memory + persisted store)
  * Used for testing cold cache scenarios in cross-model tests
  */
 app.post('/test/clear-signature-cache', (req, res) => {
-    clearThinkingSignatureCache();
-    logger.debug('[Test] Cleared thinking signature cache');
-    res.json({ success: true, message: 'Thinking signature cache cleared' });
+    clearAllSignatureCaches();
+    logger.debug('[Test] Cleared all signature caches and removed persisted store');
+    res.json({ success: true, message: 'All signature caches cleared' });
 });
 
 /**