Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions apps/web/src/app/api/openrouter/[...path]/route.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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 },
Expand Down
149 changes: 114 additions & 35 deletions apps/web/src/lib/ai-gateway/experiments/build-direct-provider.test.ts
Original file line number Diff line number Diff line change
@@ -1,19 +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<GatewayRequest, { kind: 'chat_completions' }>;

async function transformRequest(
request: GatewayRequest,
options: {
sanitize_ref_fields?: boolean;
use_gemini_reasoning_transform?: boolean;
extra_body?: Record<string, unknown>;
remove_from_body?: string[];
} = {}
options: Partial<Omit<CustomLlmApiConfig, 'internal_id' | 'base_url'>> = {}
) {
const provider = buildDirectProvider('custom', ['chat_completions'], {
internal_id: 'upstream-model',
Expand All @@ -33,6 +30,7 @@ async function transformRequest(
organization_id: null,
session_id: null,
});
applyReasoningDetailsTransform(provider, request);
}

function makeRequest(): ChatCompletionRequest {
Expand All @@ -46,27 +44,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;
}
Expand All @@ -77,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);
});
Expand All @@ -93,13 +102,10 @@ 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).toEqual({
mapGeminiThoughtContent: true,
mapReasoningContentToDetails: false,
});
expect(provider.responseTransforms).toBe(ReasoningDetailsTransform.GeminiThought);
});

it('sets response transforms to null when the transform is not enabled', () => {
Expand All @@ -114,11 +120,11 @@ 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, {
use_gemini_reasoning_transform: true,
reasoning_details_transform: ReasoningDetailsTransform.GeminiThought,
});

expect(request.body.model).toBe('upstream-model');
Expand Down Expand Up @@ -146,12 +152,77 @@ 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, {
reasoning_details_transform: ReasoningDetailsTransform.GeminiThought,
});

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<string, unknown>;
assistant.reasoning_details = [
{
type: 'reasoning.encrypted',
data: 'anthropic-signature',
format: 'anthropic-claude-v1',
},
];

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');
});

it('continues mapping legacy tool-call signatures', async () => {
const request = makeRequest();
const assistant = request.body.messages[0] as unknown as Record<string, unknown>;
delete assistant.reasoning_details;
Object.assign(assistant.tool_calls as object[], [{ thoughtSignature: 'legacy-signature' }]);

await transformRequest(request, {
reasoning_details_transform: ReasoningDetailsTransform.GeminiThought,
});

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';

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'],
});
Expand All @@ -176,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 } },
});

Expand All @@ -188,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: {
Expand All @@ -200,16 +273,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' },
]);
});
});
Expand Down
83 changes: 5 additions & 78 deletions apps/web/src/lib/ai-gateway/experiments/build-direct-provider.ts
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import { addCacheBreakpoints } from '@/lib/ai-gateway/providers/openrouter/request-helpers';
import type { CustomLlmApiConfig } from '@kilocode/db';
import type {
GatewayChatApiKind,
Provider,
TransformRequestContext,
import {
type GatewayChatApiKind,
type Provider,
type TransformRequestContext,
} from '@/lib/ai-gateway/providers/types';

/**
Expand All @@ -21,68 +21,6 @@ function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === 'object' && value !== null && !Array.isArray(value);
}

function mapGeminiThoughtSignatures(context: TransformRequestContext) {
if (context.request.kind !== 'chat_completions') {
return;
}

for (const message of context.request.body.messages) {
if (!isRecord(message)) {
continue;
}

delete message.thoughtSignature;

if (!Array.isArray(message.tool_calls)) {
continue;
}

for (const toolCall of message.tool_calls) {
if (!isRecord(toolCall)) {
continue;
}

const signature = toolCall.thoughtSignature;
delete toolCall.thoughtSignature;
if (typeof signature !== 'string') {
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,
},
};
}
}
}

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,
},
};
}

mapGeminiThoughtSignatures(context);
}

function renameJsonRefProperties(value: unknown): boolean {
if (Array.isArray(value)) {
return value.reduce<boolean>(
Expand Down Expand Up @@ -160,16 +98,8 @@ export function buildDirectProvider(
apiUrlOverrides: {},
apiKey: upstream.api_key,
supportedChatApis,
responseTransforms: upstream.use_gemini_reasoning_transform
? { mapGeminiThoughtContent: true, mapReasoningContentToDetails: false }
: 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<string, unknown>;
for (const key of upstream.remove_from_body) {
Expand All @@ -184,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);
}
Expand Down
Loading