diff --git a/open-sse/executors/antigravity.js b/open-sse/executors/antigravity.js index 35ec1006..f2ad630a 100644 --- a/open-sse/executors/antigravity.js +++ b/open-sse/executors/antigravity.js @@ -3,10 +3,11 @@ import { BaseExecutor } from "./base.js"; import { PROVIDERS } from "../config/providers.js"; import { OAUTH_ENDPOINTS, ANTIGRAVITY_HEADERS, AG_DEFAULT_TOOLS, AG_TOOL_SUFFIX, ANTIGRAVITY_PROMPT_REWRITES } from "../config/appConstants.js"; import { HTTP_STATUS } from "../config/runtimeConfig.js"; -import { resolveSessionId } from "../utils/sessionManager.js"; +import { resolveSessionId, toNumericSessionId } from "../utils/sessionManager.js"; import { proxyAwareFetch } from "../utils/proxyFetch.js"; import { cleanJSONSchemaForAntigravity } from "../translator/formats/gemini.js"; import { DEFAULT_THINKING_AG_SIGNATURE } from "../config/defaultThinkingSignature.js"; +import { getGeminiThoughtSignatureSync } from "../services/thoughtSignatureStore.js"; // Sanitize function name: Gemini requires [a-zA-Z_][a-zA-Z0-9_.:\-]{0,63} function sanitizeFunctionName(name) { @@ -187,6 +188,9 @@ export class AntigravityExecutor extends BaseExecutor { }; } + const rawSessionId = body.request?.sessionId || resolveSessionId({ headers: credentials?.rawHeaders, body, connectionId: credentials?.email || credentials?.connectionId, scope: "antigravity" }); + const sessionId = toNumericSessionId(rawSessionId) || rawSessionId; + // ─── Standard (non-image) request ─── // Fix contents for Claude models via Antigravity const contents = body.request?.contents?.map(c => { @@ -202,17 +206,31 @@ export class AntigravityExecutor extends BaseExecutor { return true; }); // Gemini 3+ rejects functionCall parts without thoughtSignature. Clients (Claude Code, IDE) - // don't persist thoughtSignature in their history, so backfill the default signature on any - // functionCall part that arrives without one. - const needsBackfill = parts?.some(p => p.functionCall && !p.thoughtSignature) ?? false; - if (role !== c.role || parts?.length !== c.parts?.length || needsBackfill) { + // don't persist thoughtSignature in their history, so backfill from cache or default signature. + // In parallel function calls, only the first call needs a signature; siblings stay unsigned. + let firstFunctionCallSeen = false; + const modifiedParts = parts?.map(p => { + if (!p.functionCall) return p; + const callId = p.functionCall.id; + const cachedSig = callId ? getGeminiThoughtSignatureSync(callId, sessionId) : null; + const callSig = p.thoughtSignature || cachedSig || (!firstFunctionCallSeen ? DEFAULT_THINKING_AG_SIGNATURE : undefined); + firstFunctionCallSeen = true; + if (callSig) { + return { ...p, thoughtSignature: callSig }; + } + if (p.thoughtSignature && !cachedSig) { + // Unsigned sibling call + const { thoughtSignature: _, ...rest } = p; + return rest; + } + return p; + }); + + const partsChanged = parts?.length !== c.parts?.length || modifiedParts?.some((p, idx) => p !== c.parts[idx]); + if (role !== c.role || partsChanged) { return { ...c, role, - parts: needsBackfill - ? parts.map(p => (p.functionCall && !p.thoughtSignature) - ? { ...p, thoughtSignature: DEFAULT_THINKING_AG_SIGNATURE } - : p) - : parts, + parts: modifiedParts || parts, }; } return c; @@ -267,7 +285,7 @@ export class AntigravityExecutor extends BaseExecutor { generationConfig, ...(contents && { contents }), ...(tools && { tools }), - sessionId: body.request?.sessionId || resolveSessionId({ headers: credentials?.rawHeaders, body, connectionId: credentials?.email || credentials?.connectionId, scope: "antigravity" }), + sessionId, safetySettings: undefined, ...(tools?.length > 0 && { toolConfig: { functionCallingConfig: { mode: "VALIDATED" } } }) }; diff --git a/open-sse/handlers/chatCore.js b/open-sse/handlers/chatCore.js index 67e377ea..3c083978 100644 --- a/open-sse/handlers/chatCore.js +++ b/open-sse/handlers/chatCore.js @@ -469,7 +469,7 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred // Streaming response const { onStreamComplete, streamDetailId } = buildOnStreamComplete({ ...sharedCtx }); - return handleStreamingResponse({ ...sharedCtx, providerResponse, sourceFormat, targetFormat: providerResponseFormat, userAgent, reqLogger, toolNameMap, customToolNames, streamController, onStreamComplete, streamDetailId }); + return handleStreamingResponse({ ...sharedCtx, providerResponse, sourceFormat, targetFormat: providerResponseFormat, userAgent, reqLogger, toolNameMap, customToolNames, streamController, onStreamComplete, streamDetailId, credentials }); } export function isTokenExpiringSoon(expiresAt, bufferMs = 5 * 60 * 1000) { diff --git a/open-sse/handlers/chatCore/streamingHandler.js b/open-sse/handlers/chatCore/streamingHandler.js index be008e7d..7b7634e2 100644 --- a/open-sse/handlers/chatCore/streamingHandler.js +++ b/open-sse/handlers/chatCore/streamingHandler.js @@ -22,7 +22,7 @@ const CODEX_SOURCE_TO_TARGET = { /** * Determine which SSE transform stream to use based on provider/format. */ -function buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, reqLogger, toolNameMap, customToolNames, model, connectionId, body, onStreamComplete, apiKey }) { +function buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, reqLogger, toolNameMap, customToolNames, model, connectionId, body, onStreamComplete, apiKey, credentials }) { const isDroidCLI = userAgent?.toLowerCase().includes("droid") || userAgent?.toLowerCase().includes("codex-cli"); // Responses-API providers (e.g. codex) emit Responses SSE → translate into client format const isResponsesProvider = PROVIDERS[provider]?.format === FORMATS.OPENAI_RESPONSES; @@ -30,11 +30,11 @@ function buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, if (needsCodexTranslation) { const codexTarget = CODEX_SOURCE_TO_TARGET[sourceFormat] || FORMATS.OPENAI; - return createSSETransformStreamWithLogger(FORMATS.OPENAI_RESPONSES, codexTarget, provider, reqLogger, toolNameMap, model, connectionId, body, onStreamComplete, apiKey, customToolNames); + return createSSETransformStreamWithLogger(FORMATS.OPENAI_RESPONSES, codexTarget, provider, reqLogger, toolNameMap, model, connectionId, body, onStreamComplete, apiKey, customToolNames, credentials); } if (needsTranslation(targetFormat, sourceFormat)) { - return createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider, reqLogger, toolNameMap, model, connectionId, body, onStreamComplete, apiKey, customToolNames); + return createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider, reqLogger, toolNameMap, model, connectionId, body, onStreamComplete, apiKey, customToolNames, credentials); } return createPassthroughStreamWithLogger(provider, reqLogger, model, connectionId, body, onStreamComplete, apiKey); @@ -43,7 +43,7 @@ function buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, /** * Handle streaming response — pipe provider SSE through transform stream to client. */ -export async function handleStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, userAgent, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, customToolNames, streamController, onStreamComplete, streamDetailId, pxpipe, reqTag, log }) { +export async function handleStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, userAgent, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, customToolNames, streamController, onStreamComplete, streamDetailId, pxpipe, reqTag, log, credentials }) { if (onRequestSuccess) { Promise.resolve() .then(onRequestSuccess) @@ -79,7 +79,7 @@ export async function handleStreamingResponse({ providerResponse, provider, mode }; } - const transformStream = buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, reqLogger, toolNameMap, customToolNames, model, connectionId, body, onStreamComplete, apiKey }); + const transformStream = buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, reqLogger, toolNameMap, customToolNames, model, connectionId, body, onStreamComplete, apiKey, credentials }); // Responses passthrough: synthesize response.failed + [DONE] if the stream aborts/stalls before a terminal event const isResponsesPassthrough = sourceFormat === FORMATS.OPENAI_RESPONSES && targetFormat === FORMATS.OPENAI_RESPONSES; diff --git a/open-sse/services/thoughtSignatureStore.js b/open-sse/services/thoughtSignatureStore.js new file mode 100644 index 00000000..eb8a60e4 --- /dev/null +++ b/open-sse/services/thoughtSignatureStore.js @@ -0,0 +1,170 @@ +import { makeKv } from "../../src/lib/db/helpers/kvStore.js"; + +const MAX_SIGNATURES = 2000; +const MAX_PERSISTED_SIGNATURES = 10_000; +const MEMORY_TTL_MS = 1000 * 60 * 60; // 1 hour +const PERSISTED_TTL_MS = 1000 * 60 * 60 * 24 * 7; // 7 days +const SCOPE = "gemini_thought_signatures"; + +const signatureKv = makeKv(SCOPE); +const memorySignatures = new Map(); +let pruneCounter = 0; + +function pruneMemoryExpired() { + const now = Date.now(); + for (const [key, value] of memorySignatures.entries()) { + if (value.expiresAt <= now) { + memorySignatures.delete(key); + } + } + + while (memorySignatures.size > MAX_SIGNATURES) { + const oldestKey = memorySignatures.keys().next().value; + if (!oldestKey) break; + memorySignatures.delete(oldestKey); + } +} + +async function maybePrunePersisted() { + pruneCounter++; + if (pruneCounter % 100 !== 0) return; + + try { + const all = await signatureKv.getAll(); + const keys = Object.keys(all); + const now = Date.now(); + const expiredKeys = []; + const valid = []; + + for (const k of keys) { + const entry = all[k]; + if (!entry || typeof entry.signature !== "string" || (entry.expiresAt && entry.expiresAt <= now)) { + expiredKeys.push(k); + } else { + valid.push({ key: k, createdAt: entry.createdAt || 0 }); + } + } + + for (const k of expiredKeys) { + await signatureKv.remove(k).catch(() => {}); + } + + if (valid.length > MAX_PERSISTED_SIGNATURES) { + valid.sort((a, b) => b.createdAt - a.createdAt); + const toRemove = valid.slice(MAX_PERSISTED_SIGNATURES); + for (const item of toRemove) { + await signatureKv.remove(item.key).catch(() => {}); + } + } + } catch { + // Fail-open + } +} + +/** + * Store a thought signature for a tool_call_id with optional sessionId namespace (RAM + SQLite async) + */ +export function storeGeminiThoughtSignature(toolCallId, signature, sessionId = null) { + if (typeof toolCallId !== "string" || !toolCallId) return; + if (typeof signature !== "string" || !signature) return; + + const now = Date.now(); + pruneMemoryExpired(); + + const keys = []; + if (sessionId && typeof sessionId === "string") { + keys.push(`${sessionId}:${toolCallId}`); + } + keys.push(toolCallId); + + for (const k of keys) { + memorySignatures.set(k, { + signature, + expiresAt: now + MEMORY_TTL_MS, + }); + + // Async persist to SQLite kv table without blocking + signatureKv.set(k, { + signature, + createdAt: now, + expiresAt: now + PERSISTED_TTL_MS, + }).catch(() => {}); + } + + maybePrunePersisted().catch(() => {}); +} + +/** + * Retrieve a thought signature by tool_call_id (RAM first, then SQLite fallback) + */ +export async function getGeminiThoughtSignature(toolCallId, sessionId = null) { + if (typeof toolCallId !== "string" || !toolCallId) return null; + + pruneMemoryExpired(); + + if (sessionId && typeof sessionId === "string") { + const sessionKey = `${sessionId}:${toolCallId}`; + const sessionEntry = memorySignatures.get(sessionKey); + if (sessionEntry && sessionEntry.expiresAt > Date.now()) { + return sessionEntry.signature; + } + } + + const entry = memorySignatures.get(toolCallId); + if (entry && entry.expiresAt > Date.now()) { + return entry.signature; + } + + try { + 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())) { + memorySignatures.set(sessionKey, { + signature: sessionRow.signature, + expiresAt: Date.now() + MEMORY_TTL_MS, + }); + return sessionRow.signature; + } + } + + const row = await signatureKv.get(toolCallId); + if (row && typeof row.signature === "string") { + if (row.expiresAt && row.expiresAt <= Date.now()) { + signatureKv.remove(toolCallId).catch(() => {}); + return null; + } + memorySignatures.set(toolCallId, { + signature: row.signature, + expiresAt: Date.now() + MEMORY_TTL_MS, + }); + return row.signature; + } + } catch { + // Fail-open + } + + return null; +} + +/** + * Synchronous get from RAM cache only (for sync translators) + */ +export function getGeminiThoughtSignatureSync(toolCallId, sessionId = null) { + if (typeof toolCallId !== "string" || !toolCallId) return null; + pruneMemoryExpired(); + + if (sessionId && typeof sessionId === "string") { + const sessionKey = `${sessionId}:${toolCallId}`; + const sessionEntry = memorySignatures.get(sessionKey); + if (sessionEntry && sessionEntry.expiresAt > Date.now()) { + return sessionEntry.signature; + } + } + + const entry = memorySignatures.get(toolCallId); + if (entry && entry.expiresAt > Date.now()) { + 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 9e029e2c..2b05670d 100644 --- a/open-sse/translator/request/openai-to-gemini.js +++ b/open-sse/translator/request/openai-to-gemini.js @@ -2,6 +2,7 @@ import { register } from "../index.js"; import { FORMATS } from "../formats.js"; import { DEFAULT_THINKING_AG_SIGNATURE, DEFAULT_THINKING_GEMINI_CLI_SIGNATURE } from "../../config/defaultThinkingSignature.js"; import { openaiToClaudeRequestForAntigravity } from "./openai-to-claude.js"; +import { getGeminiThoughtSignatureSync } from "../../services/thoughtSignatureStore.js"; function generateUUID() { return crypto.randomUUID(); } @@ -46,7 +47,7 @@ function normalizeGeminiContents(contents) { } // Core: Convert OpenAI request to Gemini format (base for all variants) -function openaiToGeminiBase(model, body, stream, signature = DEFAULT_THINKING_AG_SIGNATURE) { +function openaiToGeminiBase(model, body, stream, signature = DEFAULT_THINKING_AG_SIGNATURE, sessionId = null) { const result = { model: model, contents: [], @@ -133,18 +134,27 @@ function openaiToGeminiBase(model, body, stream, signature = DEFAULT_THINKING_AG if (msg.tool_calls && Array.isArray(msg.tool_calls)) { const toolCallIds = []; + let firstFunctionCallSeen = false; for (const tc of msg.tool_calls) { if (tc.type !== OPENAI_BLOCK.FUNCTION) continue; const args = tryParseJSON(tc.function?.arguments || "{}"); - parts.push({ - thoughtSignature: signature, + const cachedSig = tc.id ? getGeminiThoughtSignatureSync(tc.id, sessionId) : null; + // First call gets cached signature or fallback; sibling calls remain unsigned if no cached sig + const callSig = cachedSig || (!firstFunctionCallSeen ? signature : undefined); + firstFunctionCallSeen = true; + + const part = { functionCall: { id: tc.id, name: sanitizeGeminiFunctionName(tc.function.name), args: args } - }); + }; + if (callSig) { + part.thoughtSignature = callSig; + } + parts.push(part); toolCallIds.push(tc.id); } @@ -232,13 +242,13 @@ function openaiToGeminiBase(model, body, stream, signature = DEFAULT_THINKING_AG } // OpenAI -> Gemini (standard API) -export function openaiToGeminiRequest(model, body, stream) { - return openaiToGeminiBase(model, body, stream); +export function openaiToGeminiRequest(model, body, stream, credentials = null) { + return openaiToGeminiBase(model, body, stream, DEFAULT_THINKING_AG_SIGNATURE, credentials?._clientSessionId); } // OpenAI -> Gemini CLI (Cloud Code Assist) -export function openaiToGeminiCLIRequest(model, body, stream) { - const gemini = openaiToGeminiBase(model, body, stream, DEFAULT_THINKING_GEMINI_CLI_SIGNATURE); +export function openaiToGeminiCLIRequest(model, body, stream, credentials = null) { + const gemini = openaiToGeminiBase(model, body, stream, DEFAULT_THINKING_GEMINI_CLI_SIGNATURE, credentials?._clientSessionId); // Thinking is normalized centrally by applyThinking (thinkingUnified.js) after translation. // Clean schema for tools @@ -335,18 +345,26 @@ function wrapInCloudCodeEnvelopeForClaude(model, claudeRequest, credentials = nu const parts = []; if (Array.isArray(msg.content)) { + let firstToolUseSeen = false; for (const block of msg.content) { if (block.type === CLAUDE_BLOCK.TEXT) { parts.push({ text: block.text }); } else if (block.type === CLAUDE_BLOCK.TOOL_USE) { - parts.push({ - thoughtSignature: signature, + const cachedSig = block.id ? getGeminiThoughtSignatureSync(block.id, credentials?._clientSessionId) : null; + const callSig = cachedSig || (!firstToolUseSeen ? signature : undefined); + firstToolUseSeen = true; + + const part = { functionCall: { id: block.id, name: sanitizeGeminiFunctionName(block.name), args: block.input || {} } - }); + }; + if (callSig) { + part.thoughtSignature = callSig; + } + parts.push(part); } else if (block.type === CLAUDE_BLOCK.TOOL_RESULT) { let content = block.content; if (Array.isArray(content)) { diff --git a/open-sse/translator/response/gemini-to-openai.js b/open-sse/translator/response/gemini-to-openai.js index 8f04ffd3..66bf1022 100644 --- a/open-sse/translator/response/gemini-to-openai.js +++ b/open-sse/translator/response/gemini-to-openai.js @@ -6,6 +6,7 @@ import { toOpenAIUsage } from "../concerns/usage.js"; import { reasoningDelta } from "../concerns/reasoning.js"; import { encodeDataUri } from "../concerns/image.js"; import { toOpenAIFinish } from "../concerns/finishReason.js"; +import { storeGeminiThoughtSignature } from "../../services/thoughtSignatureStore.js"; // Build chunk meta for current gemini state function chunkMeta(state) { @@ -13,14 +14,18 @@ function chunkMeta(state) { } // Build a tool_call chunk from a gemini functionCall part (shared by sig/non-sig branches) -function emitFunctionCall(functionCall, state) { +function emitFunctionCall(functionCall, state, signature = null) { const rawName = functionCall.name; // Restore original tool name from mapping (AG cloaking) const fcName = state.toolNameMap?.get(rawName) || rawName; const fcArgs = functionCall.args || {}; const toolCallIndex = state.functionIndex++; + const callId = functionCall.id || `${fcName}-${Date.now()}-${toolCallIndex}`; + if (signature) { + storeGeminiThoughtSignature(callId, signature, state.sessionId); + } const toolCall = { - id: `${fcName}-${Date.now()}-${toolCallIndex}`, + id: callId, index: toolCallIndex, type: OPENAI_BLOCK.FUNCTION, function: { name: fcName, arguments: JSON.stringify(fcArgs) }, @@ -57,13 +62,21 @@ export function geminiToOpenAIResponse(chunk, state) { if (content?.parts) { for (const part of content.parts) { const hasThoughtSig = part.thoughtSignature || part.thought_signature; + if (hasThoughtSig && typeof hasThoughtSig === "string") { + state.pendingThoughtSignature = hasThoughtSig; + } const isThought = part.thought === true; - + // Handle thought signature (thinking mode) if (hasThoughtSig) { const hasTextContent = part.text !== undefined && part.text !== ""; const hasFunctionCall = !!part.functionCall; - + + // Standalone thoughtSignature part (no text, no functionCall): keep pending for next functionCall + if (!hasTextContent && !hasFunctionCall) { + continue; + } + if (hasTextContent) { results.push(buildChunk( chunkMeta(state), @@ -71,9 +84,10 @@ export function geminiToOpenAIResponse(chunk, state) { null )); } - + if (hasFunctionCall) { - results.push(emitFunctionCall(part.functionCall, state)); + results.push(emitFunctionCall(part.functionCall, state, hasThoughtSig)); + state.pendingThoughtSignature = null; } continue; } @@ -92,7 +106,9 @@ export function geminiToOpenAIResponse(chunk, state) { // Function call if (part.functionCall) { - results.push(emitFunctionCall(part.functionCall, state)); + const sig = state.pendingThoughtSignature || null; + results.push(emitFunctionCall(part.functionCall, state, sig)); + state.pendingThoughtSignature = null; } // Inline data (images) diff --git a/open-sse/utils/stream.js b/open-sse/utils/stream.js index daca8d6a..15ee37d8 100644 --- a/open-sse/utils/stream.js +++ b/open-sse/utils/stream.js @@ -49,7 +49,8 @@ export function createSSEStream(options = {}) { connectionId = null, body = null, onStreamComplete = null, - apiKey = null + apiKey = null, + credentials = null } = options; let buffer = ""; @@ -59,7 +60,7 @@ export function createSSEStream(options = {}) { const decoder = new TextDecoder("utf-8", { fatal: false }); const state = mode === STREAM_MODE.TRANSLATE - ? { ...initState(sourceFormat), provider, toolNameMap, customToolNames: new Set(customToolNames || []), model } + ? { ...initState(sourceFormat), provider, toolNameMap, customToolNames: new Set(customToolNames || []), model, sessionId: credentials?._clientSessionId || null } : null; let totalContentLength = 0; @@ -487,7 +488,7 @@ export function createSSEStream(options = {}) { }); } -export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider = null, reqLogger = null, toolNameMap = null, model = null, connectionId = null, body = null, onStreamComplete = null, apiKey = null, customToolNames = null) { +export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, provider = null, reqLogger = null, toolNameMap = null, model = null, connectionId = null, body = null, onStreamComplete = null, apiKey = null, customToolNames = null, credentials = null) { return createSSEStream({ mode: STREAM_MODE.TRANSLATE, targetFormat, @@ -500,7 +501,8 @@ export function createSSETransformStreamWithLogger(targetFormat, sourceFormat, p connectionId, body, onStreamComplete, - apiKey + apiKey, + credentials }); }