fix(antigravity): scope cached thought signatures to the model family

This commit is contained in:
Louis Phạm
2026-09-18 17:06:50 +07:00
parent 092c84eac9
commit bc3be0cb28
5 changed files with 149 additions and 16 deletions

View File

@@ -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) {

View File

@@ -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;

View File

@@ -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;

View File

@@ -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));

View File

@@ -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"]);
});
});