diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index c800dc7c60..ec2f85d0eb 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1,8 +1,20 @@ import OpenAI from "openai" import { Anthropic } from "@anthropic-ai/sdk" +import { providerIdentifiers } from "@roo-code/types/provider-identifiers" + import { LiteLLMHandler } from "../lite-llm" import { getModels } from "../fetchers/modelCache" +import { prepareApiConversationMessage } from "../../../core/task/apiConversationHistory" +import { sanitizeToolUseId } from "../../../utils/tool-id" +import { + SANITIZED_LITELLM_TOOL_CALL_ID, + TEST_THOUGHT_SIGNATURE, + signedAssistantMessage, + sentToolCalls, + toolResultsMessage, +} from "../../../test-utils/thought-signature" +import { thoughtSignatureOrigin } from "../utils/thought-signature" import { ApiHandlerOptions } from "../../../shared/api" import { type ModelRecord, litellmDefaultModelId, litellmDefaultModelInfo } from "@roo-code/types" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" @@ -561,129 +573,6 @@ describe("LiteLLMHandler", () => { }) }) - describe("injectThoughtSignatureForGemini", () => { - // Base64 encoded "skip_thought_signature_validator" - const dummySignature = Buffer.from("skip_thought_signature_validator").toString("base64") - - it("should inject provider_specific_fields.thought_signature for assistant messages with tool_calls", () => { - const handler = new LiteLLMHandler(mockOptions) - const injectThoughtSignature = (handler as any).injectThoughtSignatureForGemini.bind(handler) - - const messages = [ - { role: "user", content: "Hello" }, - { - role: "assistant", - content: "", - tool_calls: [ - { id: "call_123", type: "function", function: { name: "test_tool", arguments: "{}" } }, - ], - }, - { role: "tool", tool_call_id: "call_123", content: "result" }, - ] - - const result = injectThoughtSignature(messages) - - // The first tool call should have provider_specific_fields.thought_signature injected - expect(result[1].tool_calls[0].provider_specific_fields).toBeDefined() - expect(result[1].tool_calls[0].provider_specific_fields.thought_signature).toBe(dummySignature) - }) - - it("should not inject if assistant message has no tool_calls", () => { - const handler = new LiteLLMHandler(mockOptions) - const injectThoughtSignature = (handler as any).injectThoughtSignatureForGemini.bind(handler) - - const messages = [ - { role: "user", content: "Hello" }, - { role: "assistant", content: "Hi there!" }, - ] - - const result = injectThoughtSignature(messages) - - // No changes should be made - expect(result[1].tool_calls).toBeUndefined() - }) - - it("should always overwrite existing thought_signature", () => { - const handler = new LiteLLMHandler(mockOptions) - const injectThoughtSignature = (handler as any).injectThoughtSignatureForGemini.bind(handler) - - const existingSignature = "existing_signature_base64" - - const messages = [ - { role: "user", content: "Hello" }, - { - role: "assistant", - content: "", - tool_calls: [ - { - id: "call_123", - type: "function", - function: { name: "test_tool", arguments: "{}" }, - provider_specific_fields: { thought_signature: existingSignature }, - }, - ], - }, - ] - - const result = injectThoughtSignature(messages) - - // Should overwrite with dummy signature (always inject to ensure compatibility) - expect(result[1].tool_calls[0].provider_specific_fields.thought_signature).toBe(dummySignature) - }) - - it("should inject signature into ALL tool calls for parallel calls", () => { - const handler = new LiteLLMHandler(mockOptions) - const injectThoughtSignature = (handler as any).injectThoughtSignatureForGemini.bind(handler) - - const messages = [ - { role: "user", content: "Hello" }, - { - role: "assistant", - content: "", - tool_calls: [ - { id: "call_first", type: "function", function: { name: "tool1", arguments: "{}" } }, - { id: "call_second", type: "function", function: { name: "tool2", arguments: "{}" } }, - { id: "call_third", type: "function", function: { name: "tool3", arguments: "{}" } }, - ], - }, - ] - - const result = injectThoughtSignature(messages) - - // ALL tool calls should have the signature - expect(result[1].tool_calls[0].provider_specific_fields.thought_signature).toBe(dummySignature) - expect(result[1].tool_calls[1].provider_specific_fields.thought_signature).toBe(dummySignature) - expect(result[1].tool_calls[2].provider_specific_fields.thought_signature).toBe(dummySignature) - }) - - it("should preserve existing provider_specific_fields when adding thought_signature", () => { - const handler = new LiteLLMHandler(mockOptions) - const injectThoughtSignature = (handler as any).injectThoughtSignatureForGemini.bind(handler) - - const messages = [ - { role: "user", content: "Hello" }, - { - role: "assistant", - content: "", - tool_calls: [ - { - id: "call_123", - type: "function", - function: { name: "test_tool", arguments: "{}" }, - provider_specific_fields: { other_field: "value" }, - }, - ], - }, - ] - - const result = injectThoughtSignature(messages) - - // Should have both existing field and new thought_signature - expect(result[1].tool_calls[0].provider_specific_fields.other_field).toBe("value") - expect(result[1].tool_calls[0].provider_specific_fields.thought_signature).toBe(dummySignature) - }) - }) - describe("createMessage integration with Gemini models", () => { // Base64 encoded "skip_thought_signature_validator" const dummySignature = Buffer.from("skip_thought_signature_validator").toString("base64") @@ -759,6 +648,188 @@ describe("LiteLLMHandler", () => { expect(assistantMessage.tool_calls[0].provider_specific_fields.thought_signature).toBe(dummySignature) }) + describe("real thought signature round trip", () => { + const origin = thoughtSignatureOrigin("http://localhost:4000", "gemini-3-pro") + const metadata = makeCreateMessageMetadata({ + tools: [ + { + type: "function", + function: { name: "read_file", description: "Read a file", parameters: {} }, + }, + ], + }) + + function mockStream(chunks: unknown[]) { + mockCreate.mockReturnValue({ + withResponse: vi.fn().mockResolvedValue({ data: asyncStreamFrom(chunks) }), + }) + } + + function useModel(id: string, preserveReasoning?: boolean) { + vi.spyOn(handler, "fetchModel").mockResolvedValue({ + id, + info: { ...litellmDefaultModelInfo, preserveReasoning }, + }) + } + + const signedToolCallStream = [ + { choices: [{ delta: { content: "Reading the file." } }] }, + { + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: `call_1__thought__${TEST_THOUGHT_SIGNATURE}`, + function: { name: "read_file", arguments: "" }, + provider_specific_fields: { thought_signature: TEST_THOUGHT_SIGNATURE }, + }, + ], + }, + }, + ], + }, + { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '{"path":' } }] } }] }, + { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '"a.ts"}' } }] } }] }, + { choices: [{ delta: {} }], usage: { prompt_tokens: 1, completion_tokens: 1 } }, + ] + + it("captures the signature from a streamed tool call and resets it per request", async () => { + useModel("gemini-3-pro") + mockStream(signedToolCallStream) + + await collectStream(handler.createMessage("system", [{ role: "user", content: "Hi" }], metadata)) + expect(handler.getThoughtSignature()).toBe(TEST_THOUGHT_SIGNATURE) + expect(handler.getThoughtSignatureOrigin()).toBe(origin) + + mockStream([{ choices: [{ delta: { content: "ok" } }] }]) + await collectStream(handler.createMessage("system", [{ role: "user", content: "Hi" }], metadata)) + expect(handler.getThoughtSignature()).toBeUndefined() + }) + + it("stores the streamed signature in the history and sends it on the next request", async () => { + useModel("gemini-4-pro") + mockStream(signedToolCallStream) + await collectStream(handler.createMessage("system", [{ role: "user", content: "Hi" }], metadata)) + + const historyId = sanitizeToolUseId(`call_1__thought__${TEST_THOUGHT_SIGNATURE}`) + const assistantMessage = prepareApiConversationMessage({ + message: { + role: "assistant", + content: [{ type: "tool_use", id: historyId, name: "read_file", input: { path: "a.ts" } }], + }, + api: handler, + apiConfiguration: { apiProvider: providerIdentifiers.litellm, litellmModelId: "gemini-4-pro" }, + apiConversationHistory: [], + }) + mockStream([{ choices: [{ delta: { content: "ok" } }] }]) + await collectStream( + handler.createMessage( + "system", + [{ role: "user", content: "Hi" }, assistantMessage, toolResultsMessage([historyId])], + metadata, + ), + ) + + expect(sentToolCalls(mockCreate)).toEqual([ + expect.objectContaining({ + provider_specific_fields: { thought_signature: TEST_THOUGHT_SIGNATURE }, + }), + ]) + }) + + const bypassed = expect.objectContaining({ + provider_specific_fields: { thought_signature: dummySignature }, + }) + const untouched = expect.not.objectContaining({ provider_specific_fields: expect.anything() }) + const real = expect.objectContaining({ + provider_specific_fields: { thought_signature: TEST_THOUGHT_SIGNATURE }, + }) + + it.each([ + { + modelId: "gemini-3-pro", + preserveReasoning: undefined, + storedBy: "this backend", + storedOrigin: origin, + expected: [real, bypassed], + }, + { + modelId: "gemini-3-pro", + preserveReasoning: true, + storedBy: "this backend", + storedOrigin: origin, + expected: [real, bypassed], + }, + { + modelId: "gemini-4-pro", + preserveReasoning: undefined, + storedBy: "this backend", + storedOrigin: thoughtSignatureOrigin("http://localhost:4000/", "gemini-4-pro"), + expected: [real, untouched], + }, + { + modelId: "gemini-3-pro", + preserveReasoning: undefined, + storedBy: "another backend", + storedOrigin: "another-backend", + expected: [bypassed, bypassed], + }, + { + modelId: "gpt-4", + preserveReasoning: undefined, + storedBy: "another backend", + storedOrigin: "another-backend", + expected: [untouched, untouched], + }, + ])( + "writes signatures for $modelId with a signature from $storedBy (preserveReasoning=$preserveReasoning)", + async ({ modelId, preserveReasoning, storedOrigin, expected }) => { + useModel(modelId, preserveReasoning) + const historyIds = ["call_1", "call_2"] + mockStream([{ choices: [{ delta: { content: "ok" } }] }]) + + await collectStream( + handler.createMessage( + "system", + [ + { role: "user", content: "Hello" }, + signedAssistantMessage(historyIds, { origin: storedOrigin }), + toolResultsMessage(historyIds), + ], + metadata, + ), + ) + + expect(sentToolCalls(mockCreate)).toEqual(expected) + }, + ) + + it("sends the bypass for a sanitized __thought__ id on an alias that isGeminiModel does not match", async () => { + useModel("my-gateway-model") + const historyId = SANITIZED_LITELLM_TOOL_CALL_ID + mockStream([{ choices: [{ delta: { content: "ok" } }] }]) + + await collectStream( + handler.createMessage( + "system", + [ + { role: "user", content: "Hello" }, + { + role: "assistant", + content: [{ type: "tool_use", id: historyId, name: "read_file", input: {} }], + }, + toolResultsMessage([historyId]), + ], + metadata, + ), + ) + + expect(sentToolCalls(mockCreate)).toEqual([bypassed]) + }) + }) + it("should not inject thought signatures for non-Gemini models", async () => { const optionsWithGPT4: ApiHandlerOptions = { ...mockOptions, diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index aea27805ea..65e40be0b2 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -19,6 +19,12 @@ import { convertToOpenAiMessages } from "../transform/openai-format" import { convertToR1Format } from "../transform/r1-format" import { GEMINI_THOUGHT_SIGNATURE_BYPASS } from "../transform/gemini-format" import { sanitizeOpenAiCallId } from "../../utils/tool-id" +import { + LITELLM_THOUGHT_SIGNATURE_ID_MARKER, + applyThoughtSignatures, + readThoughtSignature, + thoughtSignatureOrigin, +} from "./utils/thought-signature" import type { SingleCompletionHandler, @@ -41,6 +47,8 @@ import { import { RequestConfigBuilder } from "./config-builder/request-config-builder" import { createTimeoutError, getRequestTimeoutMs } from "./utils/request-timeout" +const LITELLM_DEFAULT_BASE_URL = "http://localhost:4000" + /** * LiteLLM provider handler * @@ -48,11 +56,14 @@ import { createTimeoutError, getRequestTimeoutMs } from "./utils/request-timeout * It follows the OpenAI API format for compatibility. */ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHandler { + private lastThoughtSignature: string | undefined + private thoughtSignatureOrigin: string | undefined + constructor(options: ApiHandlerOptions) { super({ options, name: providerIdentifiers.litellm, - baseURL: `${options.litellmBaseUrl || "http://localhost:4000"}`, + baseURL: options.litellmBaseUrl || LITELLM_DEFAULT_BASE_URL, apiKey: options.litellmApiKey || "dummy-key", modelId: options.litellmModelId, defaultModelId: litellmDefaultModelId, @@ -111,50 +122,6 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa ) } - /** - * Inject thought signatures for Gemini models via provider_specific_fields. - * This is required when switching from other models to Gemini to satisfy API validation - * for function calls that weren't generated by Gemini (and thus lack thought signatures). - * - * Per LiteLLM documentation: - * - Thought signatures are stored in provider_specific_fields.thought_signature of tool calls - * - The bypass token (GEMINI_THOUGHT_SIGNATURE_BYPASS) skips signature validation - * - * We inject the dummy signature on EVERY tool call unconditionally to ensure Gemini - * doesn't complain about missing/corrupted signatures when conversation history - * contains tool calls from other models (like Claude). - */ - private injectThoughtSignatureForGemini( - openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[], - ): OpenAI.Chat.ChatCompletionMessageParam[] { - const dummySignature = GEMINI_THOUGHT_SIGNATURE_BYPASS - - return openAiMessages.map((msg) => { - if (msg.role === "assistant") { - const toolCalls = (msg as any).tool_calls as any[] | undefined - - // Only process if there are tool calls - if (toolCalls && toolCalls.length > 0) { - // Inject dummy signature into ALL tool calls' provider_specific_fields - // This ensures Gemini doesn't reject tool calls from other models - const updatedToolCalls = toolCalls.map((tc) => ({ - ...tc, - provider_specific_fields: { - ...(tc.provider_specific_fields || {}), - thought_signature: dummySignature, - }, - })) - - return { - ...msg, - tool_calls: updatedToolCalls, - } - } - } - return msg - }) - } - override async *createMessage( systemPrompt: string, messages: Anthropic.Messages.MessageParam[], @@ -164,6 +131,7 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa // already-aborted request fails with AbortError before provider model // discovery (getModels/refreshModels) begins. throwIfAborted(metadata?.abortSignal) + this.lastThoughtSignature = undefined // Model discovery is shared single-flight (RouterProvider.fetchModel): // concurrent callers join one in-flight fetch, so the fetch itself @@ -272,16 +240,22 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa const usesMaxCompletionTokens = this.isGpt5(modelId) || info.requiresResponsesApi const reasoningEffort = this.getReasoningEffort(info) - // For Gemini models with native protocol: inject fake reasoning.encrypted block for tool calls - // This is required when switching from other models to Gemini to satisfy API validation. - // Gemini 3 models validate thought signatures for function calls, and when conversation - // history contains tool calls from other models (like Claude), they lack the required - // signatures. The "skip_thought_signature_validator" value bypasses this validation. + // Gemini validates thought signatures on function calls. A tool call without a signature + // from this backend gets the bypass token: calls from other models (like Claude) for + // Gemini models, and `__thought__` ids whose embedded signature the history sanitized. const isGemini = this.isGeminiModel(modelId) - let processedMessages = enhancedMessages - if (isGemini) { - processedMessages = this.injectThoughtSignatureForGemini(enhancedMessages) - } + this.thoughtSignatureOrigin = thoughtSignatureOrigin( + this.options.litellmBaseUrl || LITELLM_DEFAULT_BASE_URL, + modelId, + ) + const processedMessages = applyThoughtSignatures(messages, enhancedMessages, { + origin: this.thoughtSignatureOrigin, + normalizeToolCallId: sanitizeOpenAiCallId, + fallback: (toolCallId) => + isGemini || toolCallId.includes(LITELLM_THOUGHT_SIGNATURE_ID_MARKER) + ? GEMINI_THOUGHT_SIGNATURE_BYPASS + : undefined, + }) const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { model: modelId, @@ -360,6 +334,7 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa // Handle tool calls in stream - emit partial chunks for NativeToolCallParser if (delta?.tool_calls) { + this.lastThoughtSignature ??= readThoughtSignature(delta.tool_calls) for (const toolCall of delta.tool_calls) { throwIfAborted(requestSignal) yield { @@ -428,6 +403,14 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa } } + getThoughtSignature(): string | undefined { + return this.lastThoughtSignature + } + + getThoughtSignatureOrigin(): string | undefined { + return this.lastThoughtSignature ? this.thoughtSignatureOrigin : undefined + } + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { // Fast-fail if the request was already aborted before building, so an // already-aborted request fails with AbortError before provider model diff --git a/src/api/providers/utils/__tests__/thought-signature.spec.ts b/src/api/providers/utils/__tests__/thought-signature.spec.ts new file mode 100644 index 0000000000..f6ec6f89cf --- /dev/null +++ b/src/api/providers/utils/__tests__/thought-signature.spec.ts @@ -0,0 +1,135 @@ +import OpenAI from "openai" + +import { sanitizeOpenAiCallId } from "../../../../utils/tool-id" +import { signedAssistantMessage, TEST_THOUGHT_SIGNATURE as signature } from "../../../../test-utils/thought-signature" +import { applyThoughtSignatures, readThoughtSignature, thoughtSignatureOrigin } from "../thought-signature" + +const origin = "litellm|http://localhost:4000|gemini-3-pro" +const bypass = "bypass" + +function toolCall(id: string): OpenAI.Chat.ChatCompletionMessageFunctionToolCall { + return { id, type: "function", function: { name: "read_file", arguments: "{}" } } +} + +function signed(id: string, value: string) { + return { ...toolCall(id), provider_specific_fields: { thought_signature: value } } +} + +function assistantWith(...ids: string[]): OpenAI.Chat.ChatCompletionMessageParam { + return { role: "assistant", content: null, tool_calls: ids.map(toolCall) } +} + +describe("readThoughtSignature", () => { + it("returns the first signature in provider_specific_fields", () => { + const toolCalls = [ + { index: 0, provider_specific_fields: {} }, + { index: 1, provider_specific_fields: { thought_signature: signature } }, + { index: 2, provider_specific_fields: { thought_signature: "later" } }, + ] + + expect(readThoughtSignature(toolCalls)).toBe(signature) + }) + + it.each([ + ["no tool calls", undefined], + ["null tool calls", null], + ["no provider_specific_fields", [{ index: 0 }]], + ["an empty signature", [{ index: 0, provider_specific_fields: { thought_signature: "" } }]], + ["a non-string signature", [{ index: 0, provider_specific_fields: { thought_signature: 42 } }]], + ])("returns undefined for %s", (_case, toolCalls) => { + expect(readThoughtSignature(toolCalls)).toBeUndefined() + }) +}) + +describe("thoughtSignatureOrigin", () => { + const origin = thoughtSignatureOrigin("http://localhost:4000", "gemini-3-pro") + + it("ignores a trailing slash on the base URL", () => { + expect(thoughtSignatureOrigin("http://localhost:4000/", "gemini-3-pro")).toBe(origin) + }) + + it.each([ + ["model", "http://localhost:4000", "gemini-3-pro-preview"], + ["base URL", "http://other:4000", "gemini-3-pro"], + ])("differs for another %s", (_case, baseUrl, modelId) => { + expect(thoughtSignatureOrigin(baseUrl, modelId)).not.toBe(origin) + }) + + it("does not contain the base URL", () => { + expect(origin).toMatch(/^[0-9a-f]{16}$/) + }) +}) + +describe("applyThoughtSignatures", () => { + it("puts a stored signature with the same origin on the first tool call of its message", () => { + const history = [signedAssistantMessage(["call_1", "call_2"], { origin })] + + const result = applyThoughtSignatures(history, [assistantWith("call_1", "call_2")], { origin }) + + expect(result).toEqual([ + { role: "assistant", content: null, tool_calls: [signed("call_1", signature), toolCall("call_2")] }, + ]) + }) + + it("sends each signed turn its own signature", () => { + const history = [ + signedAssistantMessage(["call_1"], { origin, signature: "first" }), + signedAssistantMessage(["call_2"], { origin, signature: "second" }), + ] + + const result = applyThoughtSignatures(history, [assistantWith("call_1"), assistantWith("call_2")], { origin }) + + expect(result).toEqual([ + { role: "assistant", content: null, tool_calls: [signed("call_1", "first")] }, + { role: "assistant", content: null, tool_calls: [signed("call_2", "second")] }, + ]) + }) + + it.each([ + ["another backend", "litellm|http://other:4000|gemini-3-pro"], + ["native Gemini (no origin)", undefined], + ["an origin that is a prefix of this one", "litellm|http://localhost:4000|gemini-3"], + ["an origin that extends this one", "litellm|http://localhost:4000|gemini-3-pro-preview"], + ])("does not send a signature minted by %s", (_case, storedOrigin) => { + const history = [signedAssistantMessage(["call_1"], { origin: storedOrigin })] + const fallback = vi.fn(() => bypass) + + const result = applyThoughtSignatures(history, [assistantWith("call_1")], { origin, fallback }) + + expect(result).toEqual([{ role: "assistant", content: null, tool_calls: [signed("call_1", bypass)] }]) + expect(fallback).toHaveBeenCalledWith("call_1") + }) + + it("matches tool calls by the normalized id", () => { + const longId = "call_" + "x".repeat(100) + const history = [signedAssistantMessage([longId], { origin })] + + const result = applyThoughtSignatures(history, [assistantWith(sanitizeOpenAiCallId(longId))], { + origin, + normalizeToolCallId: sanitizeOpenAiCallId, + }) + + expect(result).toEqual([ + { role: "assistant", content: null, tool_calls: [signed(sanitizeOpenAiCallId(longId), signature)] }, + ]) + }) + + it("uses the fallback for every tool call without a stored signature", () => { + const fallback = (id: string) => (id === "call_2" ? bypass : undefined) + + const result = applyThoughtSignatures([], [assistantWith("call_1", "call_2")], { origin, fallback }) + + expect(result).toEqual([ + { role: "assistant", content: null, tool_calls: [toolCall("call_1"), signed("call_2", bypass)] }, + ]) + }) + + it("leaves messages without tool calls unchanged", () => { + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi" }, + ] + + expect(applyThoughtSignatures([], messages, { origin, fallback: () => bypass })).toEqual(messages) + }) +}) diff --git a/src/api/providers/utils/thought-signature.ts b/src/api/providers/utils/thought-signature.ts new file mode 100644 index 0000000000..3b84ab0a84 --- /dev/null +++ b/src/api/providers/utils/thought-signature.ts @@ -0,0 +1,99 @@ +import { createHash } from "crypto" + +import type { Anthropic } from "@anthropic-ai/sdk" +import type OpenAI from "openai" + +import { type ExtendedContentBlockParam, isThoughtSignatureContentBlock } from "../../transform/gemini-format" + +/** LiteLLM embeds a Gemini thought signature in the tool call id after this marker. */ +export const LITELLM_THOUGHT_SIGNATURE_ID_MARKER = "__thought__" + +/** + * Name the backend that minted a signature, so it is sent back only to that backend. The hash + * keeps the endpoint URL out of the task history and the diagnostics export. + */ +export function thoughtSignatureOrigin(baseUrl: string, modelId: string): string { + const endpoint = baseUrl.replace(/\/+$/, "") + return createHash("sha256").update(`${endpoint}|${modelId}`).digest("hex").slice(0, 16) +} + +function signatureOf(toolCall: unknown): string | undefined { + const fields = (toolCall as { provider_specific_fields?: { thought_signature?: unknown } } | undefined) + ?.provider_specific_fields + return typeof fields?.thought_signature === "string" && fields.thought_signature.length > 0 + ? fields.thought_signature + : undefined +} + +/** + * Read the thought signature from `provider_specific_fields` of tool calls, the slot that + * LiteLLM fills for Gemini. LiteLLM also embeds it in the tool call id, but the history + * sanitizes ids, so the id copy does not survive. + */ +export function readThoughtSignature(toolCalls: readonly unknown[] | null | undefined): string | undefined { + return toolCalls?.map(signatureOf).find((signature) => signature !== undefined) +} + +function collectThoughtSignatures( + history: Anthropic.Messages.MessageParam[], + origin: string, + normalizeToolCallId: (id: string) => string, +): Map { + const signatureByToolCallId = new Map() + for (const message of history) { + if (message.role !== "assistant" || !Array.isArray(message.content)) { + continue + } + let signature: string | undefined + let firstToolUseId: string | undefined + const blocks: readonly ExtendedContentBlockParam[] = message.content + for (const block of blocks) { + if (isThoughtSignatureContentBlock(block)) { + signature = block.origin === origin ? block.thoughtSignature : undefined + } else if (block.type === "tool_use") { + firstToolUseId ??= block.id + } + } + if (signature && firstToolUseId) { + signatureByToolCallId.set(normalizeToolCallId(firstToolUseId), signature) + } + } + return signatureByToolCallId +} + +interface ApplyThoughtSignaturesOptions { + /** Only signatures stored with this origin are sent back. */ + origin: string + /** The same id normalizer the caller used to build the OpenAI messages. */ + normalizeToolCallId?: (id: string) => string + /** Signature for a tool call without a stored one, for example the Gemini bypass token. */ + fallback?: (toolCallId: string) => string | undefined +} + +/** + * Write `provider_specific_fields.thought_signature` on assistant tool calls, where LiteLLM + * reads it before the id. A stored signature goes on the first tool call of its message, + * because Gemini signs only the first call of a parallel set. + */ +export function applyThoughtSignatures( + history: Anthropic.Messages.MessageParam[], + openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[], + { origin, normalizeToolCallId = (id) => id, fallback }: ApplyThoughtSignaturesOptions, +): OpenAI.Chat.ChatCompletionMessageParam[] { + const storedSignatures = collectThoughtSignatures(history, origin, normalizeToolCallId) + + return openAiMessages.map((message) => { + if (message.role !== "assistant" || !message.tool_calls?.length) { + return message + } + return { + ...message, + tool_calls: message.tool_calls.map((toolCall) => { + const signature = storedSignatures.get(toolCall.id) ?? fallback?.(toolCall.id) + return signature + ? { ...toolCall, provider_specific_fields: { thought_signature: signature } } + : toolCall + }), + } + }) +} diff --git a/src/api/transform/__tests__/gemini-format.spec.ts b/src/api/transform/__tests__/gemini-format.spec.ts index 407984956a..a15287e815 100644 --- a/src/api/transform/__tests__/gemini-format.spec.ts +++ b/src/api/transform/__tests__/gemini-format.spec.ts @@ -2,7 +2,8 @@ import { Anthropic } from "@anthropic-ai/sdk" -import { convertAnthropicMessageToGemini } from "../gemini-format" +import { signedAssistantMessage } from "../../../test-utils/thought-signature" +import { GEMINI_THOUGHT_SIGNATURE_BYPASS, convertAnthropicMessageToGemini } from "../gemini-format" describe("convertAnthropicMessageToGemini", () => { it("should convert a simple text message", () => { @@ -166,6 +167,22 @@ describe("convertAnthropicMessageToGemini", () => { expect(functionCallParts[1].thoughtSignature).toBeUndefined() }) + it("should send the bypass instead of a signature that another backend minted", () => { + const result = convertAnthropicMessageToGemini( + signedAssistantMessage(["call-1"], { + signature: "sig-123", + origin: "litellm|http://localhost:4000|gemini-3-pro", + }), + ) + + expect(result[0]?.parts).toEqual([ + expect.objectContaining({ + functionCall: expect.anything(), + thoughtSignature: GEMINI_THOUGHT_SIGNATURE_BYPASS, + }), + ]) + }) + it("should convert a message with tool result as string", () => { const toolIdToName = new Map() toolIdToName.set("calculator-123", "calculator") diff --git a/src/api/transform/gemini-format.ts b/src/api/transform/gemini-format.ts index d45e9f27ec..7d88d26aa3 100644 --- a/src/api/transform/gemini-format.ts +++ b/src/api/transform/gemini-format.ts @@ -6,9 +6,14 @@ import { Content, Part } from "@google/genai" // This bypass token tells Gemini to skip signature validation for cross-model history entries. export const GEMINI_THOUGHT_SIGNATURE_BYPASS = Buffer.from("skip_thought_signature_validator").toString("base64") -type ThoughtSignatureContentBlock = { +export type ThoughtSignatureContentBlock = { type: "thoughtSignature" thoughtSignature?: string + /** + * Backend that minted the signature, set by handlers that are not native Gemini. + * A signature validates only on the backend that minted it. + */ + origin?: string } type ReasoningContentBlock = { @@ -16,7 +21,8 @@ type ReasoningContentBlock = { text: string } -type ExtendedContentBlockParam = Anthropic.ContentBlockParam | ThoughtSignatureContentBlock | ReasoningContentBlock +export type ExtendedContentBlockParam = + Anthropic.ContentBlockParam | ThoughtSignatureContentBlock | ReasoningContentBlock type ExtendedAnthropicContent = string | ExtendedContentBlockParam[] // Extension type to safely add thoughtSignature to Part @@ -24,7 +30,9 @@ type PartWithThoughtSignature = Part & { thoughtSignature?: string } -function isThoughtSignatureContentBlock(block: ExtendedContentBlockParam): block is ThoughtSignatureContentBlock { +export function isThoughtSignatureContentBlock( + block: ExtendedContentBlockParam, +): block is ThoughtSignatureContentBlock { return block.type === "thoughtSignature" } @@ -38,7 +46,9 @@ export function convertAnthropicContentToGemini( // First pass: find thoughtSignature if it exists in the content blocks let activeThoughtSignature: string | undefined if (Array.isArray(content)) { - const sigBlock = content.find((block) => isThoughtSignatureContentBlock(block)) as ThoughtSignatureContentBlock + const sigBlock = content.find( + (block) => isThoughtSignatureContentBlock(block) && block.origin === undefined, + ) as ThoughtSignatureContentBlock if (sigBlock?.thoughtSignature) { activeThoughtSignature = sigBlock.thoughtSignature } diff --git a/src/core/task/__tests__/apiConversationHistory.spec.ts b/src/core/task/__tests__/apiConversationHistory.spec.ts index 7a1ef2adeb..5e737c4ef1 100644 --- a/src/core/task/__tests__/apiConversationHistory.spec.ts +++ b/src/core/task/__tests__/apiConversationHistory.spec.ts @@ -137,6 +137,30 @@ describe("prepareApiConversationMessage", () => { ]) }) + it("stores the origin of a thought signature when the handler reports one", () => { + const api = { + getThoughtSignature: () => "signature-1", + getThoughtSignatureOrigin: () => "litellm|http://localhost:4000|gemini-3-pro", + } + + const result = prepareApiConversationMessage({ + message: { role: "assistant", content: "answer" }, + // prepareApiConversationMessage reads only the thought signature hooks of the handler. + api: api as unknown as ApiHandler, + apiConfiguration: { apiProvider: providerIdentifiers.litellm, litellmModelId: "gemini-3-pro" }, + apiConversationHistory: [], + }) + + expect(result.content).toEqual([ + { type: "text", text: "answer" }, + { + type: "thoughtSignature", + thoughtSignature: "signature-1", + origin: "litellm|http://localhost:4000|gemini-3-pro", + }, + ]) + }) + it("validates user tool_result blocks against the last effective assistant message", () => { const result = prepareApiConversationMessage({ message: { diff --git a/src/core/task/apiConversationHistory.ts b/src/core/task/apiConversationHistory.ts index c75fb4f1c5..cb358949f5 100644 --- a/src/core/task/apiConversationHistory.ts +++ b/src/core/task/apiConversationHistory.ts @@ -12,6 +12,7 @@ type ApiHistoryHandler = ApiHandler & { getResponseId?: () => string | undefined getEncryptedContent?: () => { encrypted_content: string; id?: string } | undefined getThoughtSignature?: () => string | undefined + getThoughtSignatureOrigin?: () => string | undefined getReasoningDetails?: () => any[] | undefined } @@ -46,6 +47,7 @@ function prepareAssistantMessage( const responseId = handler.getResponseId?.() const reasoningData = handler.getEncryptedContent?.() const thoughtSignature = handler.getThoughtSignature?.() + const thoughtSignatureOrigin = handler.getThoughtSignatureOrigin?.() const reasoningDetails = handler.getReasoningDetails?.() const modelId = getModelId(apiConfiguration) @@ -107,6 +109,7 @@ function prepareAssistantMessage( const thoughtSignatureBlock = { type: "thoughtSignature", thoughtSignature, + ...(thoughtSignatureOrigin ? { origin: thoughtSignatureOrigin } : {}), } appendContentBlock(messageWithTs, thoughtSignatureBlock) diff --git a/src/eslint-suppressions.json b/src/eslint-suppressions.json index e59c4b6cfc..4aaf09719d 100644 --- a/src/eslint-suppressions.json +++ b/src/eslint-suppressions.json @@ -81,7 +81,7 @@ }, "api/providers/__tests__/lite-llm.spec.ts": { "@typescript-eslint/no-explicit-any": { - "count": 36 + "count": 31 } }, "api/providers/__tests__/mimo.spec.ts": { @@ -276,7 +276,7 @@ }, "api/providers/lite-llm.ts": { "@typescript-eslint/no-explicit-any": { - "count": 8 + "count": 6 } }, "api/providers/mimo.ts": { diff --git a/src/test-utils/thought-signature.ts b/src/test-utils/thought-signature.ts new file mode 100644 index 0000000000..9eb048fa1d --- /dev/null +++ b/src/test-utils/thought-signature.ts @@ -0,0 +1,36 @@ +import type { Anthropic } from "@anthropic-ai/sdk" +import type OpenAI from "openai" +import type { Mock } from "vitest" + +import type { ExtendedContentBlockParam } from "../api/transform/gemini-format" + +export const TEST_THOUGHT_SIGNATURE = "AY89a18ErUKBAfkTYnDbzq6mGnRG5QS8VFogncpDzYVXzC0i+aDwY5o/0123456789abcdef==" + +/** A LiteLLM `__thought__` id after the history sanitizer changed `+`, `/` and `=` to `_`. */ +export const SANITIZED_LITELLM_TOOL_CALL_ID = + "call_1__thought__AY89a18ErUKBAfkTYnDbzq6mGnRG5QS8VFogncpDzYVXzC0i_aDwY5o_0" + +export function signedAssistantMessage( + toolUseIds: string[], + { signature = TEST_THOUGHT_SIGNATURE, origin }: { signature?: string; origin?: string } = {}, +): Anthropic.Messages.MessageParam { + const content: ExtendedContentBlockParam[] = [ + ...toolUseIds.map((id) => ({ type: "tool_use" as const, id, name: "read_file", input: {} })), + { type: "thoughtSignature", thoughtSignature: signature, ...(origin ? { origin } : {}) }, + ] + // The history stores the non-Anthropic thoughtSignature block next to the tool_use blocks. + return { role: "assistant", content: content as Anthropic.Messages.ContentBlockParam[] } +} + +export function toolResultsMessage(toolUseIds: string[]): Anthropic.Messages.MessageParam { + return { + role: "user", + content: toolUseIds.map((id) => ({ type: "tool_result", tool_use_id: id, content: "ok" })), + } +} + +export function sentToolCalls(create: Mock): OpenAI.Chat.ChatCompletionMessageToolCall[] | undefined { + const sent: OpenAI.Chat.ChatCompletionMessageParam[] = create.mock.calls.at(-1)?.[0].messages + const assistant = sent.find((msg) => msg.role === "assistant" && msg.tool_calls?.length) + return assistant?.role === "assistant" ? assistant.tool_calls : undefined +}