From bc3be0cb284f9edd711aabd6c7a8d84c1a4a8f17 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Louis=20Ph=E1=BA=A1m?= Date: Fri, 18 Sep 2026 17:06:50 +0700 Subject: [PATCH] fix(antigravity): scope cached thought signatures to the model family --- open-sse/executors/antigravity.js | 2 +- open-sse/services/thoughtSignatureStore.js | 51 +++++++-- .../translator/request/openai-to-gemini.js | 4 +- .../translator/response/gemini-to-openai.js | 4 +- ...tigravity-thought-signature-family.test.js | 104 ++++++++++++++++++ 5 files changed, 149 insertions(+), 16 deletions(-) create mode 100644 tests/unit/antigravity-thought-signature-family.test.js diff --git a/open-sse/executors/antigravity.js b/open-sse/executors/antigravity.js index fb7beee0..da4a7981 100644 --- a/open-sse/executors/antigravity.js +++ b/open-sse/executors/antigravity.js @@ -212,7 +212,7 @@ export class AntigravityExecutor extends BaseExecutor { const modifiedParts = parts?.map(p => { if (!p.functionCall) return p; const callId = p.functionCall.id; - const cachedSig = callId ? getGeminiThoughtSignatureSync(callId, sessionId) : null; + const cachedSig = callId ? getGeminiThoughtSignatureSync(callId, sessionId, body.model || model) : null; const callSig = p.thoughtSignature || cachedSig || (!firstFunctionCallSeen ? DEFAULT_THINKING_AG_SIGNATURE : undefined); firstFunctionCallSeen = true; if (callSig) { diff --git a/open-sse/services/thoughtSignatureStore.js b/open-sse/services/thoughtSignatureStore.js index eb8a60e4..71fb94b8 100644 --- a/open-sse/services/thoughtSignatureStore.js +++ b/open-sse/services/thoughtSignatureStore.js @@ -10,6 +10,24 @@ const signatureKv = makeKv(SCOPE); const memorySignatures = new Map(); let pruneCounter = 0; +/** + * Model family that produced / will consume a signature. Antigravity serves Gemini and Claude + * models behind the same API, and each backend only accepts its own signatures: a Claude + * signature replayed to Gemini fails with 400 "Corrupted thought signature." (and vice versa). + */ +export function signatureFamily(model) { + const m = typeof model === "string" ? model.toLowerCase() : ""; + if (!m) return null; + if (m.includes("claude")) return "claude"; + if (m.includes("gemini")) return "gemini"; + return m; +} + +// Entries stored before families were recorded (no `family`) stay usable for any model. +function isCompatible(entry, family) { + return !entry.family || !family || entry.family === family; +} + function pruneMemoryExpired() { const now = Date.now(); for (const [key, value] of memorySignatures.entries()) { @@ -62,13 +80,15 @@ async function maybePrunePersisted() { } /** - * Store a thought signature for a tool_call_id with optional sessionId namespace (RAM + SQLite async) + * Store a thought signature for a tool_call_id with optional sessionId namespace (RAM + SQLite async). + * `model` is the model that produced the signature; lookups for another model family skip it. */ -export function storeGeminiThoughtSignature(toolCallId, signature, sessionId = null) { +export function storeGeminiThoughtSignature(toolCallId, signature, sessionId = null, model = null) { if (typeof toolCallId !== "string" || !toolCallId) return; if (typeof signature !== "string" || !signature) return; const now = Date.now(); + const family = signatureFamily(model); pruneMemoryExpired(); const keys = []; @@ -80,12 +100,14 @@ export function storeGeminiThoughtSignature(toolCallId, signature, sessionId = n for (const k of keys) { memorySignatures.set(k, { signature, + family, expiresAt: now + MEMORY_TTL_MS, }); // Async persist to SQLite kv table without blocking signatureKv.set(k, { signature, + family, createdAt: now, expiresAt: now + PERSISTED_TTL_MS, }).catch(() => {}); @@ -95,23 +117,25 @@ export function storeGeminiThoughtSignature(toolCallId, signature, sessionId = n } /** - * Retrieve a thought signature by tool_call_id (RAM first, then SQLite fallback) + * Retrieve a thought signature by tool_call_id (RAM first, then SQLite fallback). + * `model` is the target model; signatures produced by another model family are ignored. */ -export async function getGeminiThoughtSignature(toolCallId, sessionId = null) { +export async function getGeminiThoughtSignature(toolCallId, sessionId = null, model = null) { if (typeof toolCallId !== "string" || !toolCallId) return null; + const family = signatureFamily(model); pruneMemoryExpired(); if (sessionId && typeof sessionId === "string") { const sessionKey = `${sessionId}:${toolCallId}`; const sessionEntry = memorySignatures.get(sessionKey); - if (sessionEntry && sessionEntry.expiresAt > Date.now()) { + if (sessionEntry && sessionEntry.expiresAt > Date.now() && isCompatible(sessionEntry, family)) { return sessionEntry.signature; } } const entry = memorySignatures.get(toolCallId); - if (entry && entry.expiresAt > Date.now()) { + if (entry && entry.expiresAt > Date.now() && isCompatible(entry, family)) { return entry.signature; } @@ -119,9 +143,10 @@ export async function getGeminiThoughtSignature(toolCallId, sessionId = null) { if (sessionId && typeof sessionId === "string") { const sessionKey = `${sessionId}:${toolCallId}`; const sessionRow = await signatureKv.get(sessionKey); - if (sessionRow && typeof sessionRow.signature === "string" && (!sessionRow.expiresAt || sessionRow.expiresAt > Date.now())) { + if (sessionRow && typeof sessionRow.signature === "string" && (!sessionRow.expiresAt || sessionRow.expiresAt > Date.now()) && isCompatible(sessionRow, family)) { memorySignatures.set(sessionKey, { signature: sessionRow.signature, + family: sessionRow.family || null, expiresAt: Date.now() + MEMORY_TTL_MS, }); return sessionRow.signature; @@ -134,8 +159,10 @@ export async function getGeminiThoughtSignature(toolCallId, sessionId = null) { signatureKv.remove(toolCallId).catch(() => {}); return null; } + if (!isCompatible(row, family)) return null; memorySignatures.set(toolCallId, { signature: row.signature, + family: row.family || null, expiresAt: Date.now() + MEMORY_TTL_MS, }); return row.signature; @@ -148,22 +175,24 @@ export async function getGeminiThoughtSignature(toolCallId, sessionId = null) { } /** - * Synchronous get from RAM cache only (for sync translators) + * Synchronous get from RAM cache only (for sync translators). + * `model` is the target model; signatures produced by another model family are ignored. */ -export function getGeminiThoughtSignatureSync(toolCallId, sessionId = null) { +export function getGeminiThoughtSignatureSync(toolCallId, sessionId = null, model = null) { if (typeof toolCallId !== "string" || !toolCallId) return null; + const family = signatureFamily(model); pruneMemoryExpired(); if (sessionId && typeof sessionId === "string") { const sessionKey = `${sessionId}:${toolCallId}`; const sessionEntry = memorySignatures.get(sessionKey); - if (sessionEntry && sessionEntry.expiresAt > Date.now()) { + if (sessionEntry && sessionEntry.expiresAt > Date.now() && isCompatible(sessionEntry, family)) { return sessionEntry.signature; } } const entry = memorySignatures.get(toolCallId); - if (entry && entry.expiresAt > Date.now()) { + if (entry && entry.expiresAt > Date.now() && isCompatible(entry, family)) { return entry.signature; } return null; diff --git a/open-sse/translator/request/openai-to-gemini.js b/open-sse/translator/request/openai-to-gemini.js index 24e7f262..305eaca0 100644 --- a/open-sse/translator/request/openai-to-gemini.js +++ b/open-sse/translator/request/openai-to-gemini.js @@ -129,7 +129,7 @@ function openaiToGeminiBase(model, body, stream, signature = DEFAULT_THINKING_AG if (tc.type !== OPENAI_BLOCK.FUNCTION) continue; const args = tryParseJSON(tc.function?.arguments || "{}"); - const cachedSig = tc.id ? getGeminiThoughtSignatureSync(tc.id, sessionId) : null; + const cachedSig = tc.id ? getGeminiThoughtSignatureSync(tc.id, sessionId, model) : null; // First call gets cached signature or fallback; sibling calls remain unsigned if no cached sig const callSig = cachedSig || (!firstFunctionCallSeen ? signature : undefined); firstFunctionCallSeen = true; @@ -341,7 +341,7 @@ function wrapInCloudCodeEnvelopeForClaude(model, claudeRequest, credentials = nu if (block.type === CLAUDE_BLOCK.TEXT) { parts.push({ text: block.text }); } else if (block.type === CLAUDE_BLOCK.TOOL_USE) { - const cachedSig = block.id ? getGeminiThoughtSignatureSync(block.id, credentials?._clientSessionId) : null; + const cachedSig = block.id ? getGeminiThoughtSignatureSync(block.id, credentials?._clientSessionId, model) : null; const callSig = cachedSig || (!firstToolUseSeen ? signature : undefined); firstToolUseSeen = true; diff --git a/open-sse/translator/response/gemini-to-openai.js b/open-sse/translator/response/gemini-to-openai.js index 66bf1022..14c47939 100644 --- a/open-sse/translator/response/gemini-to-openai.js +++ b/open-sse/translator/response/gemini-to-openai.js @@ -22,7 +22,7 @@ function emitFunctionCall(functionCall, state, signature = null) { const toolCallIndex = state.functionIndex++; const callId = functionCall.id || `${fcName}-${Date.now()}-${toolCallIndex}`; if (signature) { - storeGeminiThoughtSignature(callId, signature, state.sessionId); + storeGeminiThoughtSignature(callId, signature, state.sessionId, state.model); } const toolCall = { id: callId, @@ -52,7 +52,7 @@ export function geminiToOpenAIResponse(chunk, state) { // Initialize state if (!state.messageId) { state.messageId = response.responseId || `msg_${Date.now()}`; - state.model = response.modelVersion || "gemini"; + state.model = response.modelVersion || state.model || "gemini"; state.functionIndex = 0; state.geminiToolCallCount = 0; results.push(buildChunk(chunkMeta(state), { role: ROLE.ASSISTANT }, null)); diff --git a/tests/unit/antigravity-thought-signature-family.test.js b/tests/unit/antigravity-thought-signature-family.test.js new file mode 100644 index 00000000..7e34a17c --- /dev/null +++ b/tests/unit/antigravity-thought-signature-family.test.js @@ -0,0 +1,104 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// Keep the signature store in RAM only; the SQLite kv layer is not under test here. +vi.mock("@/lib/db/helpers/kvStore.js", () => ({ + makeKv: () => ({ + get: async () => null, + set: async () => {}, + remove: async () => {}, + getAll: async () => ({}), + }), +})); + +const { + storeGeminiThoughtSignature, + getGeminiThoughtSignatureSync, + signatureFamily, +} = await import("../../open-sse/services/thoughtSignatureStore.js"); +const { openaiToAntigravityRequest } = await import("../../open-sse/translator/request/openai-to-gemini.js"); +const { geminiToOpenAIResponse } = await import("../../open-sse/translator/response/gemini-to-openai.js"); +const { DEFAULT_THINKING_GEMINI_CLI_SIGNATURE } = await import("../../open-sse/config/defaultThinkingSignature.js"); + +let n = 0; +const uid = (p) => `${p}_${Date.now()}_${n++}`; + +function toolHistory(callId) { + return { + messages: [ + { role: "user", content: "list files" }, + { role: "assistant", content: null, tool_calls: [{ id: callId, type: "function", function: { name: "ls", arguments: "{}" } }] }, + { role: "tool", tool_call_id: callId, content: "a.txt" }, + ], + }; +} + +function functionCallSignatures(req) { + return req.request.contents.flatMap((c) => c.parts || []).filter((p) => p.functionCall).map((p) => p.thoughtSignature); +} + +describe("antigravity thought signatures are scoped to the model family", () => { + beforeEach(() => { n++; }); + + it("classifies model families", () => { + expect(signatureFamily("claude-opus-4-6-thinking")).toBe("claude"); + expect(signatureFamily("gemini-3.8-flash-tiered")).toBe("gemini"); + expect(signatureFamily("gpt-oss-120b-medium")).toBe("gpt-oss-120b-medium"); + expect(signatureFamily(null)).toBe(null); + }); + + it("does not return a Claude signature for a Gemini target (and vice versa)", () => { + const claudeCall = uid("toolu"); + const geminiCall = uid("call"); + storeGeminiThoughtSignature(claudeCall, "CLAUDE_SIG", "sess", "claude-opus-4-6-thinking"); + storeGeminiThoughtSignature(geminiCall, "GEMINI_SIG", "sess", "gemini-3.8-flash-tiered"); + + expect(getGeminiThoughtSignatureSync(claudeCall, "sess", "gemini-3.8-flash")).toBe(null); + expect(getGeminiThoughtSignatureSync(claudeCall, "sess", "claude-opus-4-6-thinking")).toBe("CLAUDE_SIG"); + expect(getGeminiThoughtSignatureSync(geminiCall, "sess", "gemini-3.7-flash")).toBe("GEMINI_SIG"); + expect(getGeminiThoughtSignatureSync(geminiCall, null, "claude-sonnet-4-6")).toBe(null); + }); + + it("keeps old behaviour for untagged entries and untargeted lookups", () => { + const call = uid("call"); + storeGeminiThoughtSignature(call, "LEGACY_SIG", "sess"); + expect(getGeminiThoughtSignatureSync(call, "sess", "gemini-3.8-flash")).toBe("LEGACY_SIG"); + + const tagged = uid("toolu"); + storeGeminiThoughtSignature(tagged, "CLAUDE_SIG", "sess", "claude-opus-4-6-thinking"); + expect(getGeminiThoughtSignatureSync(tagged, "sess")).toBe("CLAUDE_SIG"); + }); + + it("records the producing model from the Gemini response stream", () => { + const call = uid("toolu_vrtx"); + const state = { model: "claude-opus-4-6-thinking", sessionId: null, toolNameMap: null }; + geminiToOpenAIResponse({ + response: { + responseId: "r1", + candidates: [{ content: { role: "model", parts: [{ functionCall: { id: call, name: "ls", args: {} }, thoughtSignature: "CLAUDE_SIG" }] } }], + }, + }, state); + expect(getGeminiThoughtSignatureSync(call, null, "gemini-3.8-flash-tiered")).toBe(null); + expect(getGeminiThoughtSignatureSync(call, null, "claude-opus-4-6-thinking")).toBe("CLAUDE_SIG"); + }); + + it("switching Claude -> Gemini mid-conversation sends the default signature, not Claude's", () => { + const call = uid("toolu_vrtx"); + storeGeminiThoughtSignature(call, "CLAUDE_SIG", null, "claude-opus-4-6-thinking"); + const req = openaiToAntigravityRequest("gemini-3.8-flash-tiered", toolHistory(call), true); + expect(functionCallSignatures(req)).toEqual([DEFAULT_THINKING_GEMINI_CLI_SIGNATURE]); + }); + + it("switching Gemini -> Claude mid-conversation does not replay Gemini's signature", () => { + const call = uid("call"); + storeGeminiThoughtSignature(call, "GEMINI_SIG", null, "gemini-3.8-flash-tiered"); + const req = openaiToAntigravityRequest("claude-opus-4-6-thinking", toolHistory(call), true); + expect(functionCallSignatures(req)).not.toContain("GEMINI_SIG"); + }); + + it("same model family still reuses the cached signature", () => { + const call = uid("call"); + storeGeminiThoughtSignature(call, "GEMINI_SIG", null, "gemini-3.8-flash-tiered"); + const req = openaiToAntigravityRequest("gemini-3.8-flash-tiered", toolHistory(call), true); + expect(functionCallSignatures(req)).toEqual(["GEMINI_SIG"]); + }); +});