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 1768e2a5d2..35f2bbbd09 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,10 +1,7 @@ import { addCacheBreakpoints } from '@/lib/ai-gateway/providers/openrouter/request-helpers'; import type { CustomLlmApiConfig } from '@kilocode/db'; -import { - type GatewayChatApiKind, - type Provider, - type TransformRequestContext, -} from '@/lib/ai-gateway/providers/types'; +import { type GatewayChatApiKind, type Provider } from '@/lib/ai-gateway/providers/types'; +import { sanitizeJsonRefToolResults } from '@/lib/ai-gateway/providers/sanitize-json-ref-tool-results'; /** * Plain in-memory shape: a `CustomLlmApiConfig` merged with the decrypted @@ -17,65 +14,6 @@ import { */ export type ResolvedExperimentUpstream = CustomLlmApiConfig & { api_key: string }; -function isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null && !Array.isArray(value); -} - -function renameJsonRefProperties(value: unknown): boolean { - if (Array.isArray(value)) { - return value.reduce( - (changed, item) => renameJsonRefProperties(item) || changed, - false - ); - } - - if (!isRecord(value)) { - return false; - } - - let changed = false; - for (const [key, nestedValue] of Object.entries(value)) { - changed = renameJsonRefProperties(nestedValue) || changed; - if (key === '$ref') { - delete value.$ref; - value._ref = nestedValue; - changed = true; - } - } - return changed; -} - -function sanitizeJsonRefContent(content: string): string { - try { - const result: unknown = JSON.parse(content); - return renameJsonRefProperties(result) ? JSON.stringify(result) : content; - } catch { - return content; - } -} - -function sanitizeJsonRefToolResults(context: TransformRequestContext) { - if (context.request.kind !== 'chat_completions') { - return; - } - - for (const message of context.request.body.messages) { - if (message.role !== 'tool') { - continue; - } - - if (typeof message.content === 'string') { - message.content = sanitizeJsonRefContent(message.content); - } else { - for (const part of message.content) { - if (part.type === 'text') { - part.text = sanitizeJsonRefContent(part.text); - } - } - } - } -} - /** * Builds a `Provider` that points directly at a partner-issued upstream. * @@ -117,7 +55,7 @@ export function buildDirectProvider( addCacheBreakpoints(context.request); } if (upstream.sanitize_ref_fields) { - sanitizeJsonRefToolResults(context); + sanitizeJsonRefToolResults(context.request); } }, }; 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 8688ecb38a..d54401b746 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 @@ -4,6 +4,7 @@ import { applyAnthropicThinkingDefault, applyGatewayModelsFallback, applyPreferredProvider, + applyProviderSpecificLogic, applyReasoningDetailsTransform, removeUnsupportedRequestServiceTier, } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; @@ -19,8 +20,12 @@ import { gpt_5_6_sol_discounted_model, gpt_6_astra_flex_model, } from '@/lib/ai-gateway/providers/openai-exclusive'; +import { EmptyFraudDetectionHeaders } from '@/lib/utils'; -function makeRequest(model: string, models?: string[]): GatewayRequest { +function makeRequest( + model: string, + models?: string[] +): Extract { return { kind: 'chat_completions', body: { @@ -31,6 +36,19 @@ function makeRequest(model: string, models?: string[]): GatewayRequest { }; } +function makeProvider(responseTransforms: Provider['responseTransforms']): Provider { + return { + id: 'perplexity', + apiUrl: 'https://example.com/v1', + apiUrlOverrides: {}, + apiKey: 'test-key', + apiKeyHeader: null, + supportedChatApis: ['chat_completions'], + responseTransforms, + async transformRequest() {}, + }; +} + type MessagesThinking = Extract['body']['thinking']; function makeMessagesRequest( @@ -140,20 +158,57 @@ describe('removeUnsupportedRequestServiceTier', () => { }); }); -describe('applyReasoningDetailsTransform', () => { - function makeProvider(responseTransforms: Provider['responseTransforms']): Provider { - return { - id: 'perplexity', - apiUrl: 'https://example.com/v1', - apiUrlOverrides: {}, - apiKey: 'test-key', - apiKeyHeader: null, - supportedChatApis: ['chat_completions'], - responseTransforms, - async transformRequest() {}, - }; +describe('applyProviderSpecificLogic JSON ref field sanitization', () => { + async function applyToToolResult(model: string, content: string) { + const request = makeRequest(model); + request.body.messages = [ + { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }, + ], + }, + { role: 'tool', tool_call_id: 'call-1', content }, + ]; + + await applyProviderSpecificLogic( + makeProvider(null), + model, + request, + {}, + null, + EmptyFraudDetectionHeaders, + 'user-1', + null, + null, + null + ); + + return request.body.messages.find(message => message.role === 'tool')?.content; } + it('sanitizes JSON ref fields for Gemini models', async () => { + const content = await applyToToolResult( + 'google/gemini-3.1-pro-preview:free', + '{"$ref":"#/$defs/result"}' + ); + + expect(content).toBe('{"_ref":"#/$defs/result"}'); + }); + + it('preserves JSON ref fields for non-Gemini models', async () => { + const content = await applyToToolResult('vendor/model:free', '{"$ref":"#/$defs/result"}'); + + expect(content).toBe('{"$ref":"#/$defs/result"}'); + }); +}); + +describe('applyReasoningDetailsTransform', () => { function makeReasoningRequest(): Extract { return { kind: 'chat_completions', 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 dfa5b49a2f..305d244d2b 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 @@ -53,6 +53,8 @@ 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'; import { getCustomPricing } from '@/lib/ai-gateway/custom-pricing'; +import { isGeminiModel } from '@/lib/ai-gateway/providers/google'; +import { sanitizeJsonRefToolResults } from '@/lib/ai-gateway/providers/sanitize-json-ref-tool-results'; function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); @@ -295,6 +297,10 @@ export async function applyProviderSpecificLogic( sanitizeBinaryToolResults(requestToMutate); + if (isGeminiModel(requestedModel)) { + sanitizeJsonRefToolResults(requestToMutate); + } + if (requestToMutate.kind === 'chat_completions') { scrubOpenCodeSpecificProperties(requestToMutate.body); diff --git a/apps/web/src/lib/ai-gateway/providers/sanitize-json-ref-tool-results.ts b/apps/web/src/lib/ai-gateway/providers/sanitize-json-ref-tool-results.ts new file mode 100644 index 0000000000..b49f34641a --- /dev/null +++ b/apps/web/src/lib/ai-gateway/providers/sanitize-json-ref-tool-results.ts @@ -0,0 +1,60 @@ +import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function renameJsonRefProperties(value: unknown): boolean { + if (Array.isArray(value)) { + return value.reduce( + (changed, item) => renameJsonRefProperties(item) || changed, + false + ); + } + + if (!isRecord(value)) { + return false; + } + + let changed = false; + for (const [key, nestedValue] of Object.entries(value)) { + changed = renameJsonRefProperties(nestedValue) || changed; + if (key === '$ref') { + delete value.$ref; + value._ref = nestedValue; + changed = true; + } + } + return changed; +} + +function sanitizeJsonRefContent(content: string): string { + try { + const result: unknown = JSON.parse(content); + return renameJsonRefProperties(result) ? JSON.stringify(result) : content; + } catch { + return content; + } +} + +export function sanitizeJsonRefToolResults(request: GatewayRequest) { + if (request.kind !== 'chat_completions') { + return; + } + + for (const message of request.body.messages) { + if (message.role !== 'tool') { + continue; + } + + if (typeof message.content === 'string') { + message.content = sanitizeJsonRefContent(message.content); + } else { + for (const part of message.content) { + if (part.type === 'text') { + part.text = sanitizeJsonRefContent(part.text); + } + } + } + } +}