diff --git a/src/agents/openai-transport-stream.test.ts b/src/agents/openai-transport-stream.test.ts index a8fe415e8b6..9f9012fe58c 100644 --- a/src/agents/openai-transport-stream.test.ts +++ b/src/agents/openai-transport-stream.test.ts @@ -1654,6 +1654,232 @@ describe("openai transport stream", () => { expect(params.tools?.[0]?.function?.strict).toBe(false); }); + describe("Gemini thought_signature round-trip on OpenAI-compatible completions", () => { + const geminiModel = { + id: "gemini-3-flash-preview", + name: "Gemini 3 Flash Preview", + api: "openai-completions", + provider: "google", + baseUrl: "https://generativelanguage.googleapis.com/v1beta/openai", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1_000_000, + maxTokens: 8192, + } satisfies Model<"openai-completions">; + + function makeAssistantOutput(model: Model<"openai-completions">) { + return { + role: "assistant" as const, + content: [] as Array>, + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; + } + + it("captures thought_signature from streamed Google tool_calls", async () => { + const output = makeAssistantOutput(geminiModel); + const chunks = [ + { + id: "chatcmpl-gemini", + object: "chat.completion.chunk" as const, + created: 1, + model: geminiModel.id, + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: "call_abc", + type: "function", + function: { name: "echo_value", arguments: "" }, + extra_content: { google: { thought_signature: "SIG-OPAQUE-ABC==" } }, + }, + ], + }, + logprobs: null, + finish_reason: null, + }, + ], + }, + { + id: "chatcmpl-gemini", + object: "chat.completion.chunk" as const, + created: 1, + model: geminiModel.id, + choices: [ + { + index: 0, + delta: { + tool_calls: [{ index: 0, function: { arguments: '{"value":"repro"}' } }], + }, + logprobs: null, + finish_reason: "tool_calls" as const, + }, + ], + }, + ] as const; + async function* mockStream() { + for (const chunk of chunks) { + yield chunk as never; + } + } + + await __testing.processOpenAICompletionsStream(mockStream(), output, geminiModel, { + push() {}, + }); + + expect(output.content[0]).toMatchObject({ + type: "toolCall", + id: "call_abc", + name: "echo_value", + arguments: { value: "repro" }, + thoughtSignature: "SIG-OPAQUE-ABC==", + }); + }); + + it("re-emits captured thought_signature for same Google route tool-call replay", () => { + const params = buildOpenAICompletionsParams( + geminiModel, + { + messages: [ + { role: "user", content: "echo" }, + { + role: "assistant", + api: geminiModel.api, + provider: geminiModel.provider, + model: geminiModel.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: 1, + content: [ + { + type: "toolCall", + id: "call_abc", + name: "echo_value", + arguments: { value: "repro" }, + thoughtSignature: "SIG-OPAQUE-ABC==", + }, + ], + }, + { + role: "toolResult", + toolCallId: "call_abc", + toolName: "echo_value", + content: [{ type: "text", text: "ok" }], + isError: false, + }, + ], + tools: [], + } as never, + undefined, + ) as { messages: Array> }; + + const assistant = params.messages.find((message) => message.role === "assistant") as + | { tool_calls?: Array<{ extra_content?: { google?: { thought_signature?: string } } }> } + | undefined; + expect(assistant?.tool_calls?.[0]?.extra_content?.google?.thought_signature).toBe( + "SIG-OPAQUE-ABC==", + ); + }); + + it("does not replay thought_signature across a different API surface", () => { + const params = buildOpenAICompletionsParams( + geminiModel, + { + messages: [ + { + role: "assistant", + api: "google-generative-ai", + provider: geminiModel.provider, + model: geminiModel.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: 1, + content: [ + { + type: "toolCall", + id: "call_abc", + name: "echo_value", + arguments: { value: "repro" }, + thoughtSignature: "SIG-OPAQUE-ABC==", + }, + ], + }, + ], + tools: [], + } as never, + undefined, + ) as { messages: Array> }; + + const assistant = params.messages.find((message) => message.role === "assistant") as + | { tool_calls?: Array<{ extra_content?: unknown }> } + | undefined; + expect(assistant?.tool_calls?.[0]?.extra_content).toBeUndefined(); + }); + + it("does not emit extra_content when no thought_signature was captured", () => { + const params = buildOpenAICompletionsParams( + geminiModel, + { + messages: [ + { + role: "assistant", + api: geminiModel.api, + provider: geminiModel.provider, + model: geminiModel.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: 1, + content: [{ type: "toolCall", id: "call_abc", name: "echo_value", arguments: {} }], + }, + ], + tools: [], + } as never, + undefined, + ) as { messages: Array> }; + + const assistant = params.messages.find((message) => message.role === "assistant") as + | { tool_calls?: Array<{ extra_content?: unknown }> } + | undefined; + expect(assistant?.tool_calls?.[0]?.extra_content).toBeUndefined(); + }); + }); + it("uses Mistral compat defaults for direct Mistral completions providers", () => { const params = buildOpenAICompletionsParams( { diff --git a/src/agents/openai-transport-stream.ts b/src/agents/openai-transport-stream.ts index b271ce0cdc7..b6d0fd58305 100644 --- a/src/agents/openai-transport-stream.ts +++ b/src/agents/openai-transport-stream.ts @@ -1156,6 +1156,7 @@ async function processOpenAICompletionsStream( name: string; arguments: Record; partialArgs: string; + thoughtSignature?: string; } | null = null; let pendingPostToolCallDeltas: CompletionsReasoningDelta[] = []; @@ -1320,12 +1321,14 @@ async function processOpenAICompletionsStream( currentBlock = null; flushPendingPostToolCallDeltas(); } + const initialSig = extractGoogleThoughtSignature(toolCall); currentBlock = { type: "toolCall", id: toolCall.id || "", name: toolCall.function?.name || "", arguments: {}, partialArgs: "", + ...(initialSig ? { thoughtSignature: initialSig } : {}), }; currentToolCallArgumentBytes = 0; output.content.push(currentBlock); @@ -1340,6 +1343,10 @@ async function processOpenAICompletionsStream( if (toolCall.function?.name) { currentBlock.name = toolCall.function.name; } + const deltaSig = extractGoogleThoughtSignature(toolCall); + if (deltaSig) { + currentBlock.thoughtSignature = deltaSig; + } if (toolCall.function?.arguments) { const nextArgumentBytes = measureUtf8Bytes(toolCall.function.arguments); if ( @@ -1574,6 +1581,100 @@ function convertTools( })); } +function extractGoogleThoughtSignature(toolCall: unknown): string | undefined { + const tc = toolCall as Record | undefined; + if (!tc) { + return undefined; + } + const extra = (tc.extra_content as Record | undefined)?.google as + | Record + | undefined; + const fromExtra = extra?.thought_signature; + if (typeof fromExtra === "string" && fromExtra.length > 0) { + return fromExtra; + } + const fromFunction = (tc.function as { thought_signature?: unknown } | undefined) + ?.thought_signature; + return typeof fromFunction === "string" && fromFunction.length > 0 ? fromFunction : undefined; +} + +function isGoogleOpenAICompatModel(model: OpenAIModeModel): boolean { + const endpointClass = detectOpenAICompletionsCompat(model as Model<"openai-completions">) + .capabilities.endpointClass; + return ( + model.provider === "google" || + endpointClass === "google-generative-ai" || + endpointClass === "google-vertex" + ); +} + +function injectToolCallThoughtSignatures( + outgoingMessages: unknown[], + context: Context, + model: OpenAIModeModel, +): void { + if (!isGoogleOpenAICompatModel(model)) { + return; + } + const sigById = new Map(); + for (const msg of context.messages ?? []) { + if ((msg as { role?: string }).role !== "assistant") { + continue; + } + const source = msg as { api?: string; provider?: string; model?: string; content?: unknown }; + if ( + source.api !== model.api || + source.provider !== model.provider || + source.model !== model.id + ) { + continue; + } + if (!Array.isArray(source.content)) { + continue; + } + for (const block of source.content as Array>) { + if (block.type !== "toolCall") { + continue; + } + const id = block.id; + const sig = block.thoughtSignature; + if (typeof id === "string" && typeof sig === "string" && sig.length > 0) { + sigById.set(id, sig); + } + } + } + if (sigById.size === 0) { + return; + } + for (const message of outgoingMessages) { + const toolCalls = (message as { tool_calls?: unknown }).tool_calls; + if (!Array.isArray(toolCalls)) { + continue; + } + for (const toolCall of toolCalls as Array>) { + const id = toolCall.id; + if (typeof id !== "string") { + continue; + } + const sig = sigById.get(id); + if (!sig) { + continue; + } + const extra = + toolCall.extra_content && typeof toolCall.extra_content === "object" + ? (toolCall.extra_content as Record) + : {}; + toolCall.extra_content = extra; + const google = + extra.google && typeof extra.google === "object" + ? (extra.google as Record) + : {}; + extra.google = google; + google.thought_signature = sig; + } + } +} + export function buildOpenAICompletionsParams( model: OpenAIModeModel, context: Context, @@ -1587,6 +1688,7 @@ export function buildOpenAICompletionsParams( } : context; const messages = convertMessages(model as never, completionsContext, compat as never); + injectToolCallThoughtSignatures(messages as unknown[], context, model); const params: Record = { model: model.id, messages: compat.requiresStringContent