fix(antigravity): scope cached thought signatures to the model family
This commit is contained in:
@@ -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) {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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));
|
||||
|
||||
104
tests/unit/antigravity-thought-signature-family.test.js
Normal file
104
tests/unit/antigravity-thought-signature-family.test.js
Normal 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"]);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user