fix(web): sanitize Gemini JSON ref tool results

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<string, unknown> {
-  return typeof value === 'object' && value !== null && !Array.isArray(value);
-}
-
-function renameJsonRefProperties(value: unknown): boolean {
-  if (Array.isArray(value)) {
-    return value.reduce<boolean>(
-      (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);
       }
     },
   };

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<GatewayRequest, { kind: 'chat_completions' }> {
   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<GatewayRequest, { kind: 'messages' }>['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<GatewayRequest, { kind: 'chat_completions' }> {
     return {
       kind: 'chat_completions',

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<string, unknown> {
   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);
 

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<string, unknown> {
+  return typeof value === 'object' && value !== null && !Array.isArray(value);
+}
+
+function renameJsonRefProperties(value: unknown): boolean {
+  if (Array.isArray(value)) {
+    return value.reduce<boolean>(
+      (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);
+        }
+      }
+    }
+  }
+}