From 7dd2562ec048ef82390dfdfbaf601252a857e0dc Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 19 Aug 2026 16:48:50 +0200 Subject: [PATCH 1/3] fix(ai-gateway): preserve Gemini thought signatures --- .../api/openrouter/[...path]/route.test.ts | 7 +- .../experiments/build-direct-provider.test.ts | 110 ++++++++++++++---- .../experiments/build-direct-provider.ts | 71 +++++++---- .../apply-provider-specific-logic.test.ts | 12 +- .../apply-provider-specific-logic.ts | 11 +- .../providers/openrouter/request-helpers.ts | 2 +- .../providers/provider-definitions.test.ts | 6 +- .../providers/provider-definitions.ts | 6 +- .../web/src/lib/ai-gateway/providers/types.ts | 10 +- .../ai-gateway/rewriteModelResponse.test.ts | 85 ++++++++++---- .../lib/ai-gateway/rewriteModelResponse.ts | 104 ++++++++++++++--- 11 files changed, 320 insertions(+), 104 deletions(-) diff --git a/apps/web/src/app/api/openrouter/[...path]/route.test.ts b/apps/web/src/app/api/openrouter/[...path]/route.test.ts index 88d5ec0f66..8b73a336b1 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.test.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.test.ts @@ -12,7 +12,7 @@ import { import { emitApiMetricsForResponse } from '@/lib/ai-gateway/o11y/api-metrics.server'; import { accountForMicrodollarUsage } from '@/lib/ai-gateway/llm-proxy-helpers'; import { redisClient } from '@/lib/redis'; -import type { Provider } from '@/lib/ai-gateway/providers/types'; +import { ReasoningDetailsTransform, type Provider } from '@/lib/ai-gateway/providers/types'; import { fetchEfficientAutoDecision } from '@/lib/ai-gateway/auto-routing-decision'; import { collectDeniedAutoRoutingModelIds } from '@/lib/ai-gateway/auto-routing-denied-models'; import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; @@ -368,10 +368,7 @@ describe('POST /api/openrouter/v1/chat/completions rules-engine actions', () => }); it('passes provider response transforms to the response rewriter', async () => { - const responseTransforms = { - mapGeminiThoughtContent: true, - mapReasoningContentToDetails: false, - }; + const responseTransforms = ReasoningDetailsTransform.GeminiThought; mockedGetProvider.mockResolvedValue({ kind: 'provider', provider: { ...provider, responseTransforms }, diff --git a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts index e340101af9..b903843338 100644 --- a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts +++ b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from '@jest/globals'; import { CustomLlmApiConfigSchema } from '@kilocode/db'; import { EmptyFraudDetectionHeaders } from '@/lib/utils'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; +import { ReasoningDetailsTransform } from '@/lib/ai-gateway/providers/types'; import { buildDirectProvider } from './build-direct-provider'; type ChatCompletionRequest = Extract; @@ -46,27 +47,32 @@ function makeRequest(): ChatCompletionRequest { content: 'result', tool_call_id: 'call-1', }; + const assistantMessage = { + role: 'assistant' as const, + content: null, + tool_calls: [toolCall], + }; + Object.assign(assistantMessage, { + reasoning_details: [ + { + type: 'reasoning.encrypted' as const, + data: 'assistant-signature', + id: 'call-1', + index: 0, + format: 'google-gemini-v1' as const, + }, + ], + }); const request: ChatCompletionRequest = { kind: 'chat_completions', body: { model: 'public-model', stream: false, - messages: [ - { - role: 'assistant', - content: null, - tool_calls: [toolCall], - }, - toolMessage, - ], + messages: [assistantMessage, toolMessage], }, }; - Object.assign(toolCall, { - thoughtSignature: 'assistant-signature', - extra_content: { trace_id: 'trace-1' }, - }); - Object.assign(toolMessage, { thoughtSignature: 'tool-signature' }); + Object.assign(toolCall, { extra_content: { trace_id: 'trace-1' } }); return request; } @@ -96,10 +102,7 @@ describe('buildDirectProvider response transforms', () => { use_gemini_reasoning_transform: true, }); - expect(provider.responseTransforms).toEqual({ - mapGeminiThoughtContent: true, - mapReasoningContentToDetails: false, - }); + expect(provider.responseTransforms).toBe(ReasoningDetailsTransform.GeminiThought); }); it('sets response transforms to null when the transform is not enabled', () => { @@ -114,7 +117,7 @@ describe('buildDirectProvider response transforms', () => { }); describe('buildDirectProvider Gemini reasoning transform', () => { - it('maps assistant tool-call signatures and removes camel-case transport fields', async () => { + it('maps encrypted reasoning details to native tool-call signatures', async () => { const request = makeRequest(); await transformRequest(request, { @@ -146,6 +149,65 @@ describe('buildDirectProvider Gemini reasoning transform', () => { ]); }); + it('maps a message-level encrypted detail to a native signature', async () => { + const request: ChatCompletionRequest = { + kind: 'chat_completions', + body: { + model: 'public-model', + messages: [{ role: 'assistant', content: 'answer' }], + }, + }; + Object.assign(request.body.messages[0], { + reasoning_details: [ + { + type: 'reasoning.encrypted', + data: 'root-signature', + format: 'google-gemini-v1', + }, + ], + }); + + await transformRequest(request, { use_gemini_reasoning_transform: true }); + + expect(request.body.messages[0]).toEqual({ + role: 'assistant', + content: 'answer', + extra_content: { google: { thought_signature: 'root-signature' } }, + }); + }); + + it('does not map encrypted details from other reasoning formats', async () => { + const request = makeRequest(); + const assistant = request.body.messages[0] as unknown as Record; + assistant.reasoning_details = [ + { + type: 'reasoning.encrypted', + data: 'anthropic-signature', + format: 'anthropic-claude-v1', + }, + ]; + + await transformRequest(request, { use_gemini_reasoning_transform: true }); + + expect(request.body.messages[0]).not.toHaveProperty('reasoning_details'); + expect(request.body.messages[0]).not.toHaveProperty('extra_content.google.thought_signature'); + }); + + it('continues mapping legacy tool-call signatures', async () => { + const request = makeRequest(); + const assistant = request.body.messages[0] as unknown as Record; + delete assistant.reasoning_details; + Object.assign(assistant.tool_calls as object[], [{ thoughtSignature: 'legacy-signature' }]); + + await transformRequest(request, { use_gemini_reasoning_transform: true }); + + expect(request.body.messages[0]).toHaveProperty( + 'tool_calls.0.extra_content.google.thought_signature', + 'legacy-signature' + ); + expect(request.body.messages[0]).not.toHaveProperty('tool_calls.0.thoughtSignature'); + }); + it('moves reasoning_effort into google.thinking_config and keeps extra_body', async () => { const request = makeRequest(); request.body.reasoning_effort = 'high'; @@ -200,16 +262,22 @@ describe('buildDirectProvider Gemini reasoning transform', () => { expect(request.body).not.toHaveProperty('reasoning_effort'); }); - it('preserves signatures when the transform is not enabled', async () => { + it('preserves reasoning details when the transform is not enabled', async () => { const request = makeRequest(); await transformRequest(request); expect(request.body.messages).toMatchObject([ { - tool_calls: [{ thoughtSignature: 'assistant-signature' }], + reasoning_details: [ + { + type: 'reasoning.encrypted', + data: 'assistant-signature', + id: 'call-1', + }, + ], }, - { thoughtSignature: 'tool-signature' }, + { role: 'tool' }, ]); }); }); diff --git a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts index 3d4410dcab..a8b9eee1f3 100644 --- a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts +++ b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts @@ -1,9 +1,12 @@ import { addCacheBreakpoints } from '@/lib/ai-gateway/providers/openrouter/request-helpers'; +import { ReasoningFormat } from '@/lib/ai-gateway/custom-llm/format'; +import { ReasoningDetailType } from '@/lib/ai-gateway/custom-llm/reasoning-details'; import type { CustomLlmApiConfig } from '@kilocode/db'; -import type { - GatewayChatApiKind, - Provider, - TransformRequestContext, +import { + ReasoningDetailsTransform, + type GatewayChatApiKind, + type Provider, + type TransformRequestContext, } from '@/lib/ai-gateway/providers/types'; /** @@ -21,7 +24,19 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } -function mapGeminiThoughtSignatures(context: TransformRequestContext) { +function setGeminiThoughtSignature(value: Record, signature: string) { + const extraContent = isRecord(value.extra_content) ? value.extra_content : {}; + const google = isRecord(extraContent.google) ? extraContent.google : {}; + value.extra_content = { + ...extraContent, + google: { + ...google, + thought_signature: signature, + }, + }; +} + +function mapGeminiReasoningDetails(context: TransformRequestContext) { if (context.request.kind !== 'chat_completions') { return; } @@ -33,30 +48,36 @@ function mapGeminiThoughtSignatures(context: TransformRequestContext) { delete message.thoughtSignature; - if (!Array.isArray(message.tool_calls)) { - continue; + const toolCalls = Array.isArray(message.tool_calls) ? message.tool_calls.filter(isRecord) : []; + for (const toolCall of toolCalls) { + const legacySignature = toolCall.thoughtSignature; + delete toolCall.thoughtSignature; + if (typeof legacySignature === 'string') { + setGeminiThoughtSignature(toolCall, legacySignature); + } } - for (const toolCall of message.tool_calls) { - if (!isRecord(toolCall)) { - continue; - } + const reasoningDetails = message.reasoning_details; + delete message.reasoning_details; + if (!Array.isArray(reasoningDetails)) { + continue; + } - const signature = toolCall.thoughtSignature; - delete toolCall.thoughtSignature; - if (typeof signature !== 'string') { + for (const detail of reasoningDetails) { + if ( + !isRecord(detail) || + detail.type !== ReasoningDetailType.Encrypted || + typeof detail.data !== 'string' || + detail.format !== ReasoningFormat.GoogleGeminiV1 + ) { continue; } - const extraContent = isRecord(toolCall.extra_content) ? toolCall.extra_content : {}; - const google = isRecord(extraContent.google) ? extraContent.google : {}; - toolCall.extra_content = { - ...extraContent, - google: { - ...google, - thought_signature: signature, - }, - }; + const toolCall = + typeof detail.id === 'string' + ? toolCalls.find(candidate => candidate.id === detail.id) + : toolCalls[0]; + setGeminiThoughtSignature(toolCall ?? message, detail.data); } } } @@ -80,7 +101,7 @@ function applyGeminiReasoningTransform(context: TransformRequestContext, reasoni }; } - mapGeminiThoughtSignatures(context); + mapGeminiReasoningDetails(context); } function renameJsonRefProperties(value: unknown): boolean { @@ -161,7 +182,7 @@ export function buildDirectProvider( apiKey: upstream.api_key, supportedChatApis, responseTransforms: upstream.use_gemini_reasoning_transform - ? { mapGeminiThoughtContent: true, mapReasoningContentToDetails: false } + ? ReasoningDetailsTransform.GeminiThought : null, async transformRequest(context) { const useGeminiReasoning = Boolean(upstream.use_gemini_reasoning_transform); diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts index 5cafb6b803..1165066511 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts @@ -7,7 +7,11 @@ import { applyReasoningDetailsTransform, } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; -import type { Provider, ProviderId } from '@/lib/ai-gateway/providers/types'; +import { + ReasoningDetailsTransform, + type Provider, + type ProviderId, +} from '@/lib/ai-gateway/providers/types'; import { PERPLEXITY_KIMI_PUBLIC_ID } from '@/lib/ai-gateway/providers/moonshotai'; import { FRIENDLI_GLM_PUBLIC_ID } from '@/lib/ai-gateway/providers/zai'; @@ -120,7 +124,7 @@ describe('applyReasoningDetailsTransform', () => { const request = makeReasoningRequest(); applyReasoningDetailsTransform( - makeProvider({ mapGeminiThoughtContent: false, mapReasoningContentToDetails: true }), + makeProvider(ReasoningDetailsTransform.ReasoningContent), request ); @@ -129,7 +133,7 @@ describe('applyReasoningDetailsTransform', () => { expect(assistant.reasoning_content).toBe('thinking hard'); }); - it.each([null, { mapGeminiThoughtContent: false, mapReasoningContentToDetails: false }])( + it.each([null, ReasoningDetailsTransform.GeminiThought])( 'leaves reasoning_details untouched with transforms %p', responseTransforms => { const request = makeReasoningRequest(); @@ -146,7 +150,7 @@ describe('applyReasoningDetailsTransform', () => { const request = makeMessagesRequest('vendor/model'); applyReasoningDetailsTransform( - makeProvider({ mapGeminiThoughtContent: false, mapReasoningContentToDetails: true }), + makeProvider(ReasoningDetailsTransform.ReasoningContent), request ); diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts index 01bcfc806f..e1303e0e51 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts @@ -23,7 +23,12 @@ import { } from '@/lib/ai-gateway/providers/moonshotai'; import { FRIENDLI_GLM_PUBLIC_ID, isGlmModel } from '@/lib/ai-gateway/providers/zai'; import { isMinimaxModel } from '@/lib/ai-gateway/providers/minimax'; -import type { BYOKResult, Provider, ProviderId } from '@/lib/ai-gateway/providers/types'; +import { + ReasoningDetailsTransform, + type BYOKResult, + type Provider, + type ProviderId, +} from '@/lib/ai-gateway/providers/types'; import { isStepModel } from '@/lib/ai-gateway/providers/stepfun'; import { isDeepseekModel } from '@/lib/ai-gateway/providers/deepseek'; import type { FraudDetectionHeaders } from '@/lib/utils'; @@ -143,7 +148,7 @@ export function applyAnthropicThinkingDefault( } /** - * Inverse of the `mapReasoningContentToDetails` response transform: folds + * Inverse of the reasoning-content response transform: folds * client-supplied `reasoning_details` back into the `reasoning_content` string * the upstream speaks, so reasoning survives the round trip. */ @@ -153,7 +158,7 @@ export function applyReasoningDetailsTransform( ) { if ( requestToMutate.kind === 'chat_completions' && - provider.responseTransforms?.mapReasoningContentToDetails + provider.responseTransforms === ReasoningDetailsTransform.ReasoningContent ) { mapReasoningDetailsToReasoningContent(requestToMutate.body); } diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/request-helpers.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/request-helpers.ts index 1c68e1fd75..dbbfe63253 100644 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/request-helpers.ts +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/request-helpers.ts @@ -295,7 +295,7 @@ export function removeChatCompletionsToolNames(request: OpenRouterChatCompletion } /** - * Inverse of the `mapReasoningContentToDetails` response transform: folds + * Inverse of the reasoning-content response transform: folds * OpenRouter-style `reasoning_details` back into the DeepSeek-style * `reasoning_content` string that upstreams like Friendli and Perplexity * expect on chat completions messages. diff --git a/apps/web/src/lib/ai-gateway/providers/provider-definitions.test.ts b/apps/web/src/lib/ai-gateway/providers/provider-definitions.test.ts index fc47f4c46c..6321d4e397 100644 --- a/apps/web/src/lib/ai-gateway/providers/provider-definitions.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/provider-definitions.test.ts @@ -3,6 +3,7 @@ import { describe, expect, test } from '@jest/globals'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; import type { TransformRequestContext } from '@/lib/ai-gateway/providers/types'; +import { ReasoningDetailsTransform } from '@/lib/ai-gateway/providers/types'; import { PERPLEXITY_KIMI_PUBLIC_ID } from '@/lib/ai-gateway/providers/moonshotai'; import { FRIENDLI_GLM_PUBLIC_ID } from '@/lib/ai-gateway/providers/zai'; @@ -58,10 +59,7 @@ describe.each([ }); test('enables the reasoning details response transform', () => { - expect(provider.responseTransforms).toEqual({ - mapGeminiThoughtContent: false, - mapReasoningContentToDetails: true, - }); + expect(provider.responseTransforms).toBe(ReasoningDetailsTransform.ReasoningContent); }); test('hardwires the upstream model and removes provider settings', async () => { diff --git a/apps/web/src/lib/ai-gateway/providers/provider-definitions.ts b/apps/web/src/lib/ai-gateway/providers/provider-definitions.ts index d16d795450..63e6e37abb 100644 --- a/apps/web/src/lib/ai-gateway/providers/provider-definitions.ts +++ b/apps/web/src/lib/ai-gateway/providers/provider-definitions.ts @@ -3,7 +3,7 @@ import { isReasoningExplicitlyDisabled, removeChatCompletionsToolNames, } from '@/lib/ai-gateway/providers/openrouter/request-helpers'; -import type { Provider } from '@/lib/ai-gateway/providers/types'; +import { ReasoningDetailsTransform, type Provider } from '@/lib/ai-gateway/providers/types'; import { applyVercelSettings } from '@/lib/ai-gateway/providers/vercel'; export default { @@ -101,7 +101,7 @@ export default { // 'messages', // supported, not tested // 'responses', // supported, not tested ], - responseTransforms: { mapGeminiThoughtContent: false, mapReasoningContentToDetails: true }, + responseTransforms: ReasoningDetailsTransform.ReasoningContent, async transformRequest(context) { context.request.body.model = 'zai-org/GLM-5.2'; delete context.request.body.provider; @@ -133,7 +133,7 @@ export default { // 'messages', // supported, not tested // 'responses', // supported, not tested ], - responseTransforms: { mapGeminiThoughtContent: false, mapReasoningContentToDetails: true }, + responseTransforms: ReasoningDetailsTransform.ReasoningContent, async transformRequest(context) { context.request.body.model = 'perplexity/kimi-k3'; if (context.request.kind === 'chat_completions') { diff --git a/apps/web/src/lib/ai-gateway/providers/types.ts b/apps/web/src/lib/ai-gateway/providers/types.ts index fddb996534..41b522eb5f 100644 --- a/apps/web/src/lib/ai-gateway/providers/types.ts +++ b/apps/web/src/lib/ai-gateway/providers/types.ts @@ -40,10 +40,12 @@ export type GatewayChatApiKind = GatewayRequest['kind']; export type ProviderApiUrlOverrides = Readonly>>; -export type ProviderResponseTransforms = { - mapGeminiThoughtContent: boolean; - mapReasoningContentToDetails: boolean; -}; +export enum ReasoningDetailsTransform { + GeminiThought = 'gemini-thought', + ReasoningContent = 'reasoning-content', +} + +export type ProviderResponseTransforms = ReasoningDetailsTransform; export type Provider = { id: ProviderId; diff --git a/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts b/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts index bf6118c44f..e6e22d49b0 100644 --- a/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts +++ b/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts @@ -13,6 +13,7 @@ import { isDynamicallyOptedIntoRequestLogging } from '@/lib/ai-gateway/request-l import { QWEN37_PLUS_MODEL_ID } from '@/lib/ai-gateway/custom-pricing'; import { KILO_ORGANIZATION_ID } from '@/lib/organizations/constants'; import { logExceptInTest } from '@/lib/utils.server'; +import { ReasoningDetailsTransform } from '@/lib/ai-gateway/providers/types'; jest.mock('next/server', () => ({ ...(jest.requireActual('next/server') as Record), @@ -297,10 +298,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { removeCost: true, capture: null, vercelRequestId: null, - responseTransforms: { - mapGeminiThoughtContent: false, - mapReasoningContentToDetails: true, - }, + responseTransforms: ReasoningDetailsTransform.ReasoningContent, }); const json = await result.json(); @@ -335,10 +333,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { removeCost: true, capture: null, vercelRequestId: null, - responseTransforms: { - mapGeminiThoughtContent: false, - mapReasoningContentToDetails: true, - }, + responseTransforms: ReasoningDetailsTransform.ReasoningContent, }); const json = await result.json(); @@ -467,9 +462,9 @@ describe('rewriteModelResponse_ChatCompletions', () => { expect(dataPayloads(sse)).toContain('[DONE]'); }); - test('moves marked delta content to reasoning content', async () => { + test('maps Gemini thought content and signatures to reasoning_details', async () => { const upstream = sseResponse( - 'data: {"model":"upstream-model","choices":[{"index":0,"delta":{"content":"first thought","extra_content":{"google":{"thought":true}}}},{"index":1,"delta":{"content":"answer","extra_content":{"google":{"thought":false}}}},{"index":2,"delta":{"content":"more answer"}}]}\n\n' + + 'data: {"model":"upstream-model","choices":[{"index":0,"delta":{"content":"first thought","extra_content":{"google":{"thought":true}}}},{"index":1,"delta":{"content":"answer","extra_content":{"google":{"thought":false}}}},{"index":2,"delta":{"tool_calls":[{"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"},"extra_content":{"trace_id":"trace-1","google":{"thought_signature":"opaque-signature"}}}]}}]}\n\n' + 'data: [DONE]\n\n' ); @@ -478,21 +473,72 @@ describe('rewriteModelResponse_ChatCompletions', () => { removeCost: true, capture: null, vercelRequestId: null, - responseTransforms: { - mapGeminiThoughtContent: true, - mapReasoningContentToDetails: false, - }, + responseTransforms: ReasoningDetailsTransform.GeminiThought, }); const [chunk] = dataObjects(await readOutputStream(result)) as Array<{ choices: Array<{ delta: Record }>; }>; expect(chunk.choices[0].delta).toEqual({ - reasoning_content: 'first thought', - extra_content: { google: { thought: true } }, + reasoning_details: [ + { + type: 'reasoning.text', + text: 'first thought', + index: 0, + format: 'google-gemini-v1', + }, + ], }); expect(chunk.choices[1].delta).toMatchObject({ content: 'answer' }); - expect(chunk.choices[2].delta).toMatchObject({ content: 'more answer' }); + expect(chunk.choices[2].delta).toEqual({ + tool_calls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + extra_content: { trace_id: 'trace-1' }, + }, + ], + reasoning_details: [ + { + type: 'reasoning.encrypted', + data: 'opaque-signature', + id: 'call-1', + index: 0, + format: 'google-gemini-v1', + }, + ], + }); + }); + + test('maps a message-level Gemini signature to encrypted reasoning details', async () => { + const upstream = sseResponse( + 'data: {"model":"upstream-model","choices":[{"index":0,"delta":{"content":"answer","extra_content":{"google":{"thought_signature":"root-signature"}}}}]}\n\n' + + 'data: [DONE]\n\n' + ); + + const result = await rewriteModelResponse_ChatCompletions({ + response: upstream, + removeCost: true, + capture: null, + vercelRequestId: null, + responseTransforms: ReasoningDetailsTransform.GeminiThought, + }); + const [chunk] = dataObjects(await readOutputStream(result)) as Array<{ + choices: Array<{ delta: Record }>; + }>; + + expect(chunk.choices[0].delta).toEqual({ + content: 'answer', + reasoning_details: [ + { + type: 'reasoning.encrypted', + data: 'root-signature', + index: 0, + format: 'google-gemini-v1', + }, + ], + }); }); test('maps delta reasoning_content to reasoning_details', async () => { @@ -506,10 +552,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { removeCost: true, capture: null, vercelRequestId: null, - responseTransforms: { - mapGeminiThoughtContent: false, - mapReasoningContentToDetails: true, - }, + responseTransforms: ReasoningDetailsTransform.ReasoningContent, }); const [chunk] = dataObjects(await readOutputStream(result)) as Array<{ choices: Array<{ delta: Record }>; diff --git a/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts b/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts index 9017b2bdb5..bda62ec317 100644 --- a/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts +++ b/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts @@ -1,5 +1,6 @@ import { api_request_log, type User } from '@kilocode/db/schema'; import { + type ReasoningDetailEncrypted, type ReasoningDetailText, ReasoningDetailType, } from '@/lib/ai-gateway/custom-llm/reasoning-details'; @@ -7,7 +8,11 @@ import { isKiloExclusiveFreeModel } from '@/lib/ai-gateway/models'; import { getCustomPricing } from '@/lib/ai-gateway/custom-pricing'; import { detectToolCallArgumentErrors } from '@/lib/ai-gateway/api-request-log-errors'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; -import type { ProviderId, ProviderResponseTransforms } from '@/lib/ai-gateway/providers/types'; +import { + type ProviderId, + type ProviderResponseTransforms, + ReasoningDetailsTransform, +} from '@/lib/ai-gateway/providers/types'; import { getOutputHeaders } from '@/lib/ai-gateway/llm-proxy-helpers'; import type { ChatCompletionChunk, OpenRouterUsage } from '@/lib/ai-gateway/processUsage.types'; import { isDynamicallyOptedIntoRequestLogging } from '@/lib/ai-gateway/request-logging-opt-ins'; @@ -408,22 +413,90 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } -function isGeminiThoughtDelta(delta: Record) { - const extraContent = delta.extra_content; +function getGoogleExtraContent(value: Record) { + const extraContent = value.extra_content; if (!isRecord(extraContent)) { - return false; + return null; } const google = extraContent.google; - return isRecord(google) && google.thought === true; + return isRecord(google) ? google : null; } -function rewriteGeminiThoughtContent(delta: unknown) { - if (!isRecord(delta) || typeof delta.content !== 'string' || !isGeminiThoughtDelta(delta)) { +function deleteGoogleExtraContentProperty(value: Record, property: string) { + const extraContent = value.extra_content; + const google = getGoogleExtraContent(value); + if (!isRecord(extraContent) || !google) { return; } - delta.reasoning_content = delta.content; - delete delta.content; + delete google[property]; + if (Object.keys(google).length === 0) { + delete extraContent.google; + } + if (Object.keys(extraContent).length === 0) { + delete value.extra_content; + } +} + +/** + * Normalizes Google's OpenAI-compatible thought fields to OpenRouter reasoning + * details so clients can retain opaque Gemini signatures between tool calls. + */ +function rewriteGeminiThoughtToReasoningDetails(delta: unknown) { + if (!isRecord(delta)) { + return; + } + + const details: Array = []; + const google = getGoogleExtraContent(delta); + if (typeof delta.content === 'string' && google?.thought === true) { + details.push({ + type: ReasoningDetailType.Text, + text: delta.content, + index: 0, + format: ReasoningFormat.GoogleGeminiV1, + }); + delete delta.content; + deleteGoogleExtraContentProperty(delta, 'thought'); + } + + if (typeof google?.thought_signature === 'string') { + details.push({ + type: ReasoningDetailType.Encrypted, + data: google.thought_signature, + index: 0, + format: ReasoningFormat.GoogleGeminiV1, + }); + deleteGoogleExtraContentProperty(delta, 'thought_signature'); + } + + if (Array.isArray(delta.tool_calls)) { + for (const toolCall of delta.tool_calls) { + if (!isRecord(toolCall)) { + continue; + } + + const signature = getGoogleExtraContent(toolCall)?.thought_signature; + if (typeof signature !== 'string') { + continue; + } + + details.push({ + type: ReasoningDetailType.Encrypted, + data: signature, + id: typeof toolCall.id === 'string' ? toolCall.id : undefined, + index: 0, + format: ReasoningFormat.GoogleGeminiV1, + }); + deleteGoogleExtraContentProperty(toolCall, 'thought_signature'); + } + } + + if (details.length > 0) { + delta.reasoning_details = Array.isArray(delta.reasoning_details) + ? [...delta.reasoning_details, ...details] + : details; + } } /** @@ -486,7 +559,12 @@ export async function rewriteModelResponse_ChatCompletions({ if (usage) { rewriteUsage(usage, removeCost); } - if (responseTransforms?.mapReasoningContentToDetails) { + if (responseTransforms === ReasoningDetailsTransform.GeminiThought) { + for (const choice of json.choices ?? []) { + rewriteGeminiThoughtToReasoningDetails(choice.message); + } + } + if (responseTransforms === ReasoningDetailsTransform.ReasoningContent) { for (const choice of json.choices ?? []) { rewriteReasoningContentToReasoningDetails(choice.message); } @@ -546,10 +624,10 @@ export async function rewriteModelResponse_ChatCompletions({ if (delta.role === null) { delete delta.role; } - if (responseTransforms?.mapGeminiThoughtContent) { - rewriteGeminiThoughtContent(delta); + if (responseTransforms === ReasoningDetailsTransform.GeminiThought) { + rewriteGeminiThoughtToReasoningDetails(delta); } - if (responseTransforms?.mapReasoningContentToDetails) { + if (responseTransforms === ReasoningDetailsTransform.ReasoningContent) { rewriteReasoningContentToReasoningDetails(delta); } } From 87774b51d4de52a3de8bc679940f1064f74a0ab1 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 19 Aug 2026 17:04:28 +0200 Subject: [PATCH 2/3] refactor(ai-gateway): centralize reasoning transforms --- .../experiments/build-direct-provider.test.ts | 45 ++++-- .../experiments/build-direct-provider.ts | 96 +----------- .../apply-provider-specific-logic.test.ts | 66 ++++++-- .../apply-provider-specific-logic.ts | 101 +++++++++++- .../web/src/lib/ai-gateway/providers/types.ts | 13 +- .../ai-gateway/reasoning-details-transform.ts | 131 ++++++++++++++++ .../lib/ai-gateway/rewriteModelResponse.ts | 147 +----------------- packages/db/src/schema-types.ts | 11 +- 8 files changed, 332 insertions(+), 278 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/reasoning-details-transform.ts diff --git a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts index b903843338..25a1b4012d 100644 --- a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts +++ b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts @@ -1,20 +1,16 @@ import { describe, expect, it } from '@jest/globals'; -import { CustomLlmApiConfigSchema } from '@kilocode/db'; +import { CustomLlmApiConfigSchema, type CustomLlmApiConfig } from '@kilocode/db'; import { EmptyFraudDetectionHeaders } from '@/lib/utils'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; import { ReasoningDetailsTransform } from '@/lib/ai-gateway/providers/types'; +import { applyReasoningDetailsTransform } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; import { buildDirectProvider } from './build-direct-provider'; type ChatCompletionRequest = Extract; async function transformRequest( request: GatewayRequest, - options: { - sanitize_ref_fields?: boolean; - use_gemini_reasoning_transform?: boolean; - extra_body?: Record; - remove_from_body?: string[]; - } = {} + options: Partial> = {} ) { const provider = buildDirectProvider('custom', ['chat_completions'], { internal_id: 'upstream-model', @@ -34,6 +30,7 @@ async function transformRequest( organization_id: null, session_id: null, }); + applyReasoningDetailsTransform(provider, request); } function makeRequest(): ChatCompletionRequest { @@ -83,11 +80,17 @@ describe('custom LLM Gemini reasoning transform configuration', () => { base_url: 'https://llm.example.com/v1', }; - it('accepts the Gemini reasoning transform flag', () => { + it('accepts reasoning detail transform enum values', () => { expect( CustomLlmApiConfigSchema.safeParse({ ...config, - use_gemini_reasoning_transform: true, + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, + }).success + ).toBe(true); + expect( + CustomLlmApiConfigSchema.safeParse({ + ...config, + reasoning_details_transform: ReasoningDetailsTransform.ReasoningContent, }).success ).toBe(true); }); @@ -99,7 +102,7 @@ describe('buildDirectProvider response transforms', () => { internal_id: 'upstream-model', base_url: 'https://llm.example.com/v1', api_key: 'test-key', - use_gemini_reasoning_transform: true, + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, }); expect(provider.responseTransforms).toBe(ReasoningDetailsTransform.GeminiThought); @@ -121,7 +124,7 @@ describe('buildDirectProvider Gemini reasoning transform', () => { const request = makeRequest(); await transformRequest(request, { - use_gemini_reasoning_transform: true, + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, }); expect(request.body.model).toBe('upstream-model'); @@ -167,7 +170,9 @@ describe('buildDirectProvider Gemini reasoning transform', () => { ], }); - await transformRequest(request, { use_gemini_reasoning_transform: true }); + await transformRequest(request, { + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, + }); expect(request.body.messages[0]).toEqual({ role: 'assistant', @@ -187,7 +192,9 @@ describe('buildDirectProvider Gemini reasoning transform', () => { }, ]; - await transformRequest(request, { use_gemini_reasoning_transform: true }); + await transformRequest(request, { + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, + }); expect(request.body.messages[0]).not.toHaveProperty('reasoning_details'); expect(request.body.messages[0]).not.toHaveProperty('extra_content.google.thought_signature'); @@ -199,7 +206,9 @@ describe('buildDirectProvider Gemini reasoning transform', () => { delete assistant.reasoning_details; Object.assign(assistant.tool_calls as object[], [{ thoughtSignature: 'legacy-signature' }]); - await transformRequest(request, { use_gemini_reasoning_transform: true }); + await transformRequest(request, { + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, + }); expect(request.body.messages[0]).toHaveProperty( 'tool_calls.0.extra_content.google.thought_signature', @@ -213,7 +222,7 @@ describe('buildDirectProvider Gemini reasoning transform', () => { request.body.reasoning_effort = 'high'; await transformRequest(request, { - use_gemini_reasoning_transform: true, + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, extra_body: { temperature: 0.2, google: { existing: true } }, remove_from_body: ['stream'], }); @@ -238,7 +247,7 @@ describe('buildDirectProvider Gemini reasoning transform', () => { request.body.reasoning_effort = 'none'; await transformRequest(request, { - use_gemini_reasoning_transform: true, + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, extra_body: { google: { existing: true } }, }); @@ -250,7 +259,9 @@ describe('buildDirectProvider Gemini reasoning transform', () => { it('still sets thinking_config when reasoning_effort is absent', async () => { const request = makeRequest(); - await transformRequest(request, { use_gemini_reasoning_transform: true }); + await transformRequest(request, { + reasoning_details_transform: ReasoningDetailsTransform.GeminiThought, + }); expect(request.body).toMatchObject({ google: { diff --git a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts index a8b9eee1f3..7f58ccd620 100644 --- a/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts +++ b/apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts @@ -1,9 +1,6 @@ import { addCacheBreakpoints } from '@/lib/ai-gateway/providers/openrouter/request-helpers'; -import { ReasoningFormat } from '@/lib/ai-gateway/custom-llm/format'; -import { ReasoningDetailType } from '@/lib/ai-gateway/custom-llm/reasoning-details'; import type { CustomLlmApiConfig } from '@kilocode/db'; import { - ReasoningDetailsTransform, type GatewayChatApiKind, type Provider, type TransformRequestContext, @@ -24,86 +21,6 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } -function setGeminiThoughtSignature(value: Record, signature: string) { - const extraContent = isRecord(value.extra_content) ? value.extra_content : {}; - const google = isRecord(extraContent.google) ? extraContent.google : {}; - value.extra_content = { - ...extraContent, - google: { - ...google, - thought_signature: signature, - }, - }; -} - -function mapGeminiReasoningDetails(context: TransformRequestContext) { - if (context.request.kind !== 'chat_completions') { - return; - } - - for (const message of context.request.body.messages) { - if (!isRecord(message)) { - continue; - } - - delete message.thoughtSignature; - - const toolCalls = Array.isArray(message.tool_calls) ? message.tool_calls.filter(isRecord) : []; - for (const toolCall of toolCalls) { - const legacySignature = toolCall.thoughtSignature; - delete toolCall.thoughtSignature; - if (typeof legacySignature === 'string') { - setGeminiThoughtSignature(toolCall, legacySignature); - } - } - - const reasoningDetails = message.reasoning_details; - delete message.reasoning_details; - if (!Array.isArray(reasoningDetails)) { - continue; - } - - for (const detail of reasoningDetails) { - if ( - !isRecord(detail) || - detail.type !== ReasoningDetailType.Encrypted || - typeof detail.data !== 'string' || - detail.format !== ReasoningFormat.GoogleGeminiV1 - ) { - continue; - } - - const toolCall = - typeof detail.id === 'string' - ? toolCalls.find(candidate => candidate.id === detail.id) - : toolCalls[0]; - setGeminiThoughtSignature(toolCall ?? message, detail.data); - } - } -} - -function applyGeminiReasoningTransform(context: TransformRequestContext, reasoningEffort: unknown) { - if (context.request.kind !== 'chat_completions') { - return; - } - - const extra = context.request.body as typeof context.request.body & { google?: unknown }; - delete extra.reasoning_effort; - - if (reasoningEffort !== 'none') { - const existingGoogle = isRecord(extra.google) ? extra.google : {}; - extra.google = { - ...existingGoogle, - thinking_config: { - ...(reasoningEffort !== undefined ? { thinking_level: reasoningEffort } : {}), - include_thoughts: true, - }, - }; - } - - mapGeminiReasoningDetails(context); -} - function renameJsonRefProperties(value: unknown): boolean { if (Array.isArray(value)) { return value.reduce( @@ -181,16 +98,8 @@ export function buildDirectProvider( apiUrlOverrides: {}, apiKey: upstream.api_key, supportedChatApis, - responseTransforms: upstream.use_gemini_reasoning_transform - ? ReasoningDetailsTransform.GeminiThought - : null, + responseTransforms: upstream.reasoning_details_transform ?? null, async transformRequest(context) { - const useGeminiReasoning = Boolean(upstream.use_gemini_reasoning_transform); - const reasoningEffort = - useGeminiReasoning && context.request.kind === 'chat_completions' - ? context.request.body.reasoning_effort - : undefined; - if (upstream.remove_from_body) { const body = context.request.body as Record; for (const key of upstream.remove_from_body) { @@ -205,9 +114,6 @@ export function buildDirectProvider( if (upstream.add_cache_breakpoints) { addCacheBreakpoints(context.request); } - if (useGeminiReasoning) { - applyGeminiReasoningTransform(context, reasoningEffort); - } if (upstream.sanitize_ref_fields) { sanitizeJsonRefToolResults(context); } diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts index 1165066511..904e65bb45 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts @@ -133,18 +133,64 @@ describe('applyReasoningDetailsTransform', () => { expect(assistant.reasoning_content).toBe('thinking hard'); }); - it.each([null, ReasoningDetailsTransform.GeminiThought])( - 'leaves reasoning_details untouched with transforms %p', - responseTransforms => { - const request = makeReasoningRequest(); + it('leaves reasoning_details untouched without a transform', () => { + const request = makeReasoningRequest(); - applyReasoningDetailsTransform(makeProvider(responseTransforms), request); + applyReasoningDetailsTransform(makeProvider(null), request); - const assistant = request.body.messages[1] as unknown as Record; - expect(assistant.reasoning_details).toBeDefined(); - expect(assistant.reasoning_content).toBeUndefined(); - } - ); + const assistant = request.body.messages[1] as unknown as Record; + expect(assistant.reasoning_details).toBeDefined(); + expect(assistant.reasoning_content).toBeUndefined(); + }); + + it('maps Gemini encrypted details to matching tool-call signatures', () => { + const request: Extract = { + kind: 'chat_completions', + body: { + model: 'vendor/model', + reasoning_effort: 'high', + messages: [ + { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }, + ], + reasoning_details: [ + { + type: 'reasoning.encrypted', + data: 'opaque-signature', + id: 'call-1', + format: 'google-gemini-v1', + }, + ], + } as never, + ], + }, + }; + + applyReasoningDetailsTransform(makeProvider(ReasoningDetailsTransform.GeminiThought), request); + + expect(request.body).toMatchObject({ + google: { thinking_config: { thinking_level: 'high', include_thoughts: true } }, + messages: [ + { + tool_calls: [ + { + id: 'call-1', + extra_content: { google: { thought_signature: 'opaque-signature' } }, + }, + ], + }, + ], + }); + expect(request.body).not.toHaveProperty('reasoning_effort'); + expect(request.body.messages[0]).not.toHaveProperty('reasoning_details'); + }); it('does not touch Messages requests', () => { const request = makeMessagesRequest('vendor/model'); diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts index e1303e0e51..02d1746474 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts @@ -50,6 +50,85 @@ import { import { isQwenExplicitCacheModel, isQwenModel } from '@/lib/ai-gateway/providers/qwen'; import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; import { isOpenAiModel } from '@/lib/ai-gateway/providers/openai'; +import { ReasoningFormat } from '@/lib/ai-gateway/custom-llm/format'; +import { ReasoningDetailType } from '@/lib/ai-gateway/custom-llm/reasoning-details'; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function setGeminiThoughtSignature(value: Record, signature: string) { + const extraContent = isRecord(value.extra_content) ? value.extra_content : {}; + const google = isRecord(extraContent.google) ? extraContent.google : {}; + value.extra_content = { + ...extraContent, + google: { + ...google, + thought_signature: signature, + }, + }; +} + +function mapGeminiReasoningDetails(request: OpenRouterChatCompletionRequest) { + for (const message of request.messages) { + if (!isRecord(message)) { + continue; + } + + delete message.thoughtSignature; + + const toolCalls = Array.isArray(message.tool_calls) ? message.tool_calls.filter(isRecord) : []; + for (const toolCall of toolCalls) { + const legacySignature = toolCall.thoughtSignature; + delete toolCall.thoughtSignature; + if (typeof legacySignature === 'string') { + setGeminiThoughtSignature(toolCall, legacySignature); + } + } + + const reasoningDetails = message.reasoning_details; + delete message.reasoning_details; + if (!Array.isArray(reasoningDetails)) { + continue; + } + + for (const detail of reasoningDetails) { + if ( + !isRecord(detail) || + detail.type !== ReasoningDetailType.Encrypted || + typeof detail.data !== 'string' || + detail.format !== ReasoningFormat.GoogleGeminiV1 + ) { + continue; + } + + const toolCall = + typeof detail.id === 'string' + ? toolCalls.find(candidate => candidate.id === detail.id) + : toolCalls[0]; + setGeminiThoughtSignature(toolCall ?? message, detail.data); + } + } +} + +function applyGeminiReasoningTransform(request: OpenRouterChatCompletionRequest) { + const reasoningEffort = request.reasoning_effort; + const extra = request as typeof request & { google?: unknown }; + delete extra.reasoning_effort; + + if (reasoningEffort !== 'none') { + const existingGoogle = isRecord(extra.google) ? extra.google : {}; + extra.google = { + ...existingGoogle, + thinking_config: { + ...(reasoningEffort !== undefined ? { thinking_level: reasoningEffort } : {}), + include_thoughts: true, + }, + }; + } + + mapGeminiReasoningDetails(request); +} export function getPreferredProviderOrder(requestedModel: string): string[] { if (isOpenAiModel(requestedModel)) { @@ -156,11 +235,19 @@ export function applyReasoningDetailsTransform( provider: Provider, requestToMutate: GatewayRequest ) { - if ( - requestToMutate.kind === 'chat_completions' && - provider.responseTransforms === ReasoningDetailsTransform.ReasoningContent - ) { - mapReasoningDetailsToReasoningContent(requestToMutate.body); + if (requestToMutate.kind !== 'chat_completions') { + return; + } + + switch (provider.responseTransforms) { + case ReasoningDetailsTransform.GeminiThought: + applyGeminiReasoningTransform(requestToMutate.body); + break; + case ReasoningDetailsTransform.ReasoningContent: + mapReasoningDetailsToReasoningContent(requestToMutate.body); + break; + case null: + break; } } @@ -186,8 +273,6 @@ export async function applyProviderSpecificLogic( repairChatCompletionsTools(requestToMutate.body); - applyReasoningDetailsTransform(provider, requestToMutate); - if (isClaudeModel(requestedModel)) { // Workaround for older clients corrupting Claude reasoning, resulting in: // `thinking` or `redacted_thinking` blocks in the latest assistant message cannot be modified @@ -243,4 +328,6 @@ export async function applyProviderSpecificLogic( organization_id: organizationId, session_id: sessionId, }); + + applyReasoningDetailsTransform(provider, requestToMutate); } diff --git a/apps/web/src/lib/ai-gateway/providers/types.ts b/apps/web/src/lib/ai-gateway/providers/types.ts index 41b522eb5f..4549412cc0 100644 --- a/apps/web/src/lib/ai-gateway/providers/types.ts +++ b/apps/web/src/lib/ai-gateway/providers/types.ts @@ -1,6 +1,12 @@ import type { UserByokProviderId } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; import type { FraudDetectionHeaders } from '@/lib/utils'; +import { + ReasoningDetailsTransform, + type ReasoningDetailsTransform as ReasoningDetailsTransformType, +} from '@kilocode/db'; + +export { ReasoningDetailsTransform }; export type ProviderId = | 'openrouter' @@ -40,12 +46,7 @@ export type GatewayChatApiKind = GatewayRequest['kind']; export type ProviderApiUrlOverrides = Readonly>>; -export enum ReasoningDetailsTransform { - GeminiThought = 'gemini-thought', - ReasoningContent = 'reasoning-content', -} - -export type ProviderResponseTransforms = ReasoningDetailsTransform; +export type ProviderResponseTransforms = ReasoningDetailsTransformType; export type Provider = { id: ProviderId; diff --git a/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts b/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts new file mode 100644 index 0000000000..d9f03c187b --- /dev/null +++ b/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts @@ -0,0 +1,131 @@ +import { ReasoningFormat } from '@/lib/ai-gateway/custom-llm/format'; +import { + type ReasoningDetailEncrypted, + type ReasoningDetailText, + ReasoningDetailType, +} from '@/lib/ai-gateway/custom-llm/reasoning-details'; +import { + ReasoningDetailsTransform, + type ProviderResponseTransforms, +} from '@/lib/ai-gateway/providers/types'; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function getGoogleExtraContent(value: Record) { + const extraContent = value.extra_content; + if (!isRecord(extraContent)) { + return null; + } + const google = extraContent.google; + return isRecord(google) ? google : null; +} + +function deleteGoogleExtraContentProperty(value: Record, property: string) { + const extraContent = value.extra_content; + const google = getGoogleExtraContent(value); + if (!isRecord(extraContent) || !google) { + return; + } + + delete google[property]; + if (Object.keys(google).length === 0) { + delete extraContent.google; + } + if (Object.keys(extraContent).length === 0) { + delete value.extra_content; + } +} + +function mapGeminiThoughtToReasoningDetails(value: unknown) { + if (!isRecord(value)) { + return; + } + + const details: Array = []; + const google = getGoogleExtraContent(value); + if (typeof value.content === 'string' && google?.thought === true) { + details.push({ + type: ReasoningDetailType.Text, + text: value.content, + index: 0, + format: ReasoningFormat.GoogleGeminiV1, + }); + delete value.content; + deleteGoogleExtraContentProperty(value, 'thought'); + } + + if (typeof google?.thought_signature === 'string') { + details.push({ + type: ReasoningDetailType.Encrypted, + data: google.thought_signature, + index: 0, + format: ReasoningFormat.GoogleGeminiV1, + }); + deleteGoogleExtraContentProperty(value, 'thought_signature'); + } + + if (Array.isArray(value.tool_calls)) { + for (const toolCall of value.tool_calls) { + if (!isRecord(toolCall)) { + continue; + } + + const signature = getGoogleExtraContent(toolCall)?.thought_signature; + if (typeof signature !== 'string') { + continue; + } + + details.push({ + type: ReasoningDetailType.Encrypted, + data: signature, + id: typeof toolCall.id === 'string' ? toolCall.id : undefined, + index: 0, + format: ReasoningFormat.GoogleGeminiV1, + }); + deleteGoogleExtraContentProperty(toolCall, 'thought_signature'); + } + } + + if (details.length > 0) { + value.reasoning_details = Array.isArray(value.reasoning_details) + ? [...value.reasoning_details, ...details] + : details; + } +} + +function mapReasoningContentToDetails(value: unknown) { + if ( + !isRecord(value) || + typeof value.reasoning_content !== 'string' || + typeof value.reasoning_details !== 'undefined' + ) { + return; + } + + const detail = { + type: ReasoningDetailType.Text, + text: value.reasoning_content, + index: 0, + format: ReasoningFormat.Unknown, + } satisfies ReasoningDetailText; + value.reasoning_details = [detail]; + delete value.reasoning_content; +} + +export function applyReasoningDetailsResponseTransform( + transform: ProviderResponseTransforms | null, + value: unknown +) { + switch (transform) { + case ReasoningDetailsTransform.GeminiThought: + mapGeminiThoughtToReasoningDetails(value); + break; + case ReasoningDetailsTransform.ReasoningContent: + mapReasoningContentToDetails(value); + break; + case null: + break; + } +} diff --git a/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts b/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts index bda62ec317..765159a2aa 100644 --- a/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts +++ b/apps/web/src/lib/ai-gateway/rewriteModelResponse.ts @@ -1,18 +1,9 @@ import { api_request_log, type User } from '@kilocode/db/schema'; -import { - type ReasoningDetailEncrypted, - type ReasoningDetailText, - ReasoningDetailType, -} from '@/lib/ai-gateway/custom-llm/reasoning-details'; import { isKiloExclusiveFreeModel } from '@/lib/ai-gateway/models'; import { getCustomPricing } from '@/lib/ai-gateway/custom-pricing'; import { detectToolCallArgumentErrors } from '@/lib/ai-gateway/api-request-log-errors'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; -import { - type ProviderId, - type ProviderResponseTransforms, - ReasoningDetailsTransform, -} from '@/lib/ai-gateway/providers/types'; +import type { ProviderId, ProviderResponseTransforms } from '@/lib/ai-gateway/providers/types'; import { getOutputHeaders } from '@/lib/ai-gateway/llm-proxy-helpers'; import type { ChatCompletionChunk, OpenRouterUsage } from '@/lib/ai-gateway/processUsage.types'; import { isDynamicallyOptedIntoRequestLogging } from '@/lib/ai-gateway/request-logging-opt-ins'; @@ -25,7 +16,7 @@ import { createParser } from 'eventsource-parser'; import { after, NextResponse } from 'next/server'; import type OpenAI from 'openai'; import type Anthropic from '@anthropic-ai/sdk'; -import { ReasoningFormat } from '@/lib/ai-gateway/custom-llm/format'; +import { applyReasoningDetailsResponseTransform } from '@/lib/ai-gateway/reasoning-details-transform'; /** * Handle passed to the response pipeline so the upstream response body can be @@ -409,122 +400,6 @@ function rewriteUsage(usage: OpenRouterUsage, removeCost: boolean) { } } -function isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null && !Array.isArray(value); -} - -function getGoogleExtraContent(value: Record) { - const extraContent = value.extra_content; - if (!isRecord(extraContent)) { - return null; - } - const google = extraContent.google; - return isRecord(google) ? google : null; -} - -function deleteGoogleExtraContentProperty(value: Record, property: string) { - const extraContent = value.extra_content; - const google = getGoogleExtraContent(value); - if (!isRecord(extraContent) || !google) { - return; - } - - delete google[property]; - if (Object.keys(google).length === 0) { - delete extraContent.google; - } - if (Object.keys(extraContent).length === 0) { - delete value.extra_content; - } -} - -/** - * Normalizes Google's OpenAI-compatible thought fields to OpenRouter reasoning - * details so clients can retain opaque Gemini signatures between tool calls. - */ -function rewriteGeminiThoughtToReasoningDetails(delta: unknown) { - if (!isRecord(delta)) { - return; - } - - const details: Array = []; - const google = getGoogleExtraContent(delta); - if (typeof delta.content === 'string' && google?.thought === true) { - details.push({ - type: ReasoningDetailType.Text, - text: delta.content, - index: 0, - format: ReasoningFormat.GoogleGeminiV1, - }); - delete delta.content; - deleteGoogleExtraContentProperty(delta, 'thought'); - } - - if (typeof google?.thought_signature === 'string') { - details.push({ - type: ReasoningDetailType.Encrypted, - data: google.thought_signature, - index: 0, - format: ReasoningFormat.GoogleGeminiV1, - }); - deleteGoogleExtraContentProperty(delta, 'thought_signature'); - } - - if (Array.isArray(delta.tool_calls)) { - for (const toolCall of delta.tool_calls) { - if (!isRecord(toolCall)) { - continue; - } - - const signature = getGoogleExtraContent(toolCall)?.thought_signature; - if (typeof signature !== 'string') { - continue; - } - - details.push({ - type: ReasoningDetailType.Encrypted, - data: signature, - id: typeof toolCall.id === 'string' ? toolCall.id : undefined, - index: 0, - format: ReasoningFormat.GoogleGeminiV1, - }); - deleteGoogleExtraContentProperty(toolCall, 'thought_signature'); - } - } - - if (details.length > 0) { - delta.reasoning_details = Array.isArray(delta.reasoning_details) - ? [...delta.reasoning_details, ...details] - : details; - } -} - -/** - * Converts the DeepSeek-style `reasoning_content` string into OpenRouter-style - * `reasoning_details`, matching the shape produced by OpenRouter and consumed - * by clients such as `@openrouter/ai-sdk-provider`. The flat string carries no - * format or signature, so every chunk becomes a `reasoning.text` detail at - * index 0; clients merge consecutive same-type deltas into a single block. - */ -function rewriteReasoningContentToReasoningDetails(delta: unknown) { - if ( - !isRecord(delta) || - typeof delta.reasoning_content !== 'string' || - typeof delta.reasoning_details !== 'undefined' - ) { - return; - } - - const detail = { - type: ReasoningDetailType.Text, - text: delta.reasoning_content, - index: 0, - format: ReasoningFormat.Unknown, - } satisfies ReasoningDetailText; - delta.reasoning_details = [detail]; - delete delta.reasoning_content; -} - export async function rewriteModelResponse_ChatCompletions({ response, removeCost, @@ -559,15 +434,8 @@ export async function rewriteModelResponse_ChatCompletions({ if (usage) { rewriteUsage(usage, removeCost); } - if (responseTransforms === ReasoningDetailsTransform.GeminiThought) { - for (const choice of json.choices ?? []) { - rewriteGeminiThoughtToReasoningDetails(choice.message); - } - } - if (responseTransforms === ReasoningDetailsTransform.ReasoningContent) { - for (const choice of json.choices ?? []) { - rewriteReasoningContentToReasoningDetails(choice.message); - } + for (const choice of json.choices ?? []) { + applyReasoningDetailsResponseTransform(responseTransforms, choice.message); } return NextResponse.json(json, { @@ -624,12 +492,7 @@ export async function rewriteModelResponse_ChatCompletions({ if (delta.role === null) { delete delta.role; } - if (responseTransforms === ReasoningDetailsTransform.GeminiThought) { - rewriteGeminiThoughtToReasoningDetails(delta); - } - if (responseTransforms === ReasoningDetailsTransform.ReasoningContent) { - rewriteReasoningContentToReasoningDetails(delta); - } + applyReasoningDetailsResponseTransform(responseTransforms, delta); } if (!json.choices) { diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index 83220afa76..e9e7c2aa30 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -2042,6 +2042,15 @@ export const CustomLlmMetadataSchema = z.object({ export type CustomLlmMetadata = z.infer; +export const ReasoningDetailsTransform = { + GeminiThought: 'gemini-thought', + ReasoningContent: 'reasoning-content', +} as const; + +export const ReasoningDetailsTransformSchema = z.enum(ReasoningDetailsTransform); + +export type ReasoningDetailsTransform = z.infer; + export const CustomLlmApiConfigSchema = z.object({ internal_id: z.string().min(1), base_url: z.url(), @@ -2050,7 +2059,7 @@ export const CustomLlmApiConfigSchema = z.object({ extra_headers: CustomLlmExtraHeadersSchema.optional(), extra_body: CustomLlmExtraBodySchema.optional(), remove_from_body: z.array(z.string()).optional(), - use_gemini_reasoning_transform: z.boolean().optional(), + reasoning_details_transform: ReasoningDetailsTransformSchema.optional(), }); export type CustomLlmApiConfig = z.infer; From 65464de21f9f3273853197c779027d607c24ae10 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 19 Aug 2026 21:42:39 +0200 Subject: [PATCH 3/3] fix(ai-gateway): retain message-level Gemini signatures --- .../apply-provider-specific-logic.test.ts | 47 +++++++++++++++++++ .../apply-provider-specific-logic.ts | 2 +- .../ai-gateway/reasoning-details-transform.ts | 2 + 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts index 904e65bb45..062eb4eb6d 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.test.ts @@ -192,6 +192,53 @@ describe('applyReasoningDetailsTransform', () => { expect(request.body.messages[0]).not.toHaveProperty('reasoning_details'); }); + it('keeps id-less Gemini signatures on the message when tool calls are present', () => { + const request: Extract = { + kind: 'chat_completions', + body: { + model: 'vendor/model', + messages: [ + { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }, + ], + reasoning_details: [ + { + type: 'reasoning.encrypted', + data: 'message-signature', + format: 'google-gemini-v1', + }, + { + type: 'reasoning.encrypted', + data: 'tool-signature', + id: 'call-1', + format: 'google-gemini-v1', + }, + ], + } as never, + ], + }, + }; + + applyReasoningDetailsTransform(makeProvider(ReasoningDetailsTransform.GeminiThought), request); + + expect(request.body.messages[0]).toMatchObject({ + extra_content: { google: { thought_signature: 'message-signature' } }, + tool_calls: [ + { + id: 'call-1', + extra_content: { google: { thought_signature: 'tool-signature' } }, + }, + ], + }); + }); + it('does not touch Messages requests', () => { const request = makeMessagesRequest('vendor/model'); diff --git a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts index 02d1746474..4cde8e57b5 100644 --- a/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts +++ b/apps/web/src/lib/ai-gateway/providers/apply-provider-specific-logic.ts @@ -105,7 +105,7 @@ function mapGeminiReasoningDetails(request: OpenRouterChatCompletionRequest) { const toolCall = typeof detail.id === 'string' ? toolCalls.find(candidate => candidate.id === detail.id) - : toolCalls[0]; + : undefined; setGeminiThoughtSignature(toolCall ?? message, detail.data); } } diff --git a/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts b/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts index d9f03c187b..5113ae967b 100644 --- a/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts +++ b/apps/web/src/lib/ai-gateway/reasoning-details-transform.ts @@ -43,6 +43,8 @@ function mapGeminiThoughtToReasoningDetails(value: unknown) { return; } + // OpenRouter commonly emits every block with index 0. Its AI SDK provider + // merges adjacent text by type and always keeps encrypted blocks discrete. const details: Array = []; const google = getGoogleExtraContent(value); if (typeof value.content === 'string' && google?.thought === true) {