fix(ai): gemini not working in ask ai (#21898)

Upstream issue: https://github.com/vercel/ai/issues/14369 Gemini 400s whenever a tool result contains JSON Schema `$ref`/`$defs` (it reads `$ref` as a function declaration name and finds no match). We hit this because `learn_tools` returns tool input schemas, and our recursive filter schema emits `$ref`/`$defs`. Other providers accept it fine, so this only blocks Gemini. Adds a Google-only `wrapLanguageModel` middleware that serializes ref-bearing tool results to text before they reach Gemini, so the pointers travel as a string instead of structured keys. The model still reads the full schema (same as the MCP path). Guarded so normal tool results pass through untouched.  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/21898?utm_source=github" target="_blank" rel="noopener noreferrer" data-no-image-dialog="true"><picture><source media="(prefers-color-scheme: dark)" srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img alt="Review in cubic" src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a>  --------- Co-authored-by: Charles Bochet <charles@twenty.com>
2026-06-22 12:46:59 +05:30
parent 3ad5259106
commit 1b7dc0367e
4 changed files with 232 additions and 4 deletions
@@ -0,0 +1,12 @@
+import { type LanguageModelMiddleware } from 'ai';
+
+import { sanitizeToolResultRefs } from 'src/engine/metadata-modules/ai/ai-models/utils/sanitize-tool-result-refs.util';
+
+// Gemini 400s on JSON Schema `$ref`/`$defs` in tool results, so serialize them to text (Google only): https://github.com/vercel/ai/issues/14369
+export const sanitizeGeminiToolResultRefsMiddleware: LanguageModelMiddleware = {
+  specificationVersion: 'v3',
+  transformParams: async ({ params }) => ({
+    ...params,
+    prompt: sanitizeToolResultRefs(params.prompt),
+  }),
+};
@@ -9,7 +9,11 @@ import { createOpenAI, type OpenAIProvider } from '@ai-sdk/openai';
 import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
 import { createXai, type XaiProvider } from '@ai-sdk/xai';
 import { fromNodeProviderChain } from '@aws-sdk/credential-providers';
-import { type LanguageModel } from 'ai';
+import {
+  wrapLanguageModel,
+  type LanguageModel,
+  type LanguageModelMiddleware,
+} from 'ai';
 import { type AiSdkPackage } from 'twenty-shared/ai';

 import {
@@ -22,6 +26,7 @@ import {
  AI_SDK_OPENAI_COMPATIBLE,
  AI_SDK_XAI,
 } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-sdk-package.const';
+import { sanitizeGeminiToolResultRefsMiddleware } from 'src/engine/metadata-modules/ai/ai-models/middleware/sanitize-gemini-tool-result-refs.middleware';
 import { type AiProviderConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-provider-config.type';

 export type AiSdkProviderInstance = {
@@ -92,7 +97,9 @@ export class SdkProviderFactoryService {
      case AI_SDK_ANTHROPIC:
        return this.buildStandardProvider(config, createAnthropic);
      case AI_SDK_GOOGLE:
-        return this.buildStandardProvider(config, createGoogleGenerativeAI);
+        return this.buildStandardProvider(config, createGoogleGenerativeAI, {
+          middleware: sanitizeGeminiToolResultRefsMiddleware,
+        });
      case AI_SDK_MISTRAL:
        return this.buildStandardProvider(config, createMistral);
      case AI_SDK_XAI:
@@ -111,6 +118,7 @@ export class SdkProviderFactoryService {
  private buildStandardProvider(
    config: AiProviderConfig,
    factory: (opts: { apiKey?: string; baseURL?: string }) => CallableFunction,
+    options?: { middleware?: LanguageModelMiddleware },
  ): AiSdkProviderInstance {
    const provider = factory({
      ...(config.apiKey && { apiKey: config.apiKey }),
@@ -118,8 +126,13 @@ export class SdkProviderFactoryService {
    });

    return {
-      createModel: (modelId: string) =>
-        (provider as CallableFunction)(modelId) as LanguageModel,
+      createModel: (modelId: string) => {
+        const model = (provider as CallableFunction)(modelId);
+
+        return options?.middleware
+          ? wrapLanguageModel({ model, middleware: options.middleware })
+          : model;
+      },
      rawProvider: provider,
      sdkPackage: config.npm,
    };
@@ -0,0 +1,148 @@
+import {
+  type LanguageModelV3Prompt,
+  type LanguageModelV3ToolResultPart,
+} from '@ai-sdk/provider';
+
+import { sanitizeToolResultRefs } from 'src/engine/metadata-modules/ai/ai-models/utils/sanitize-tool-result-refs.util';
+
+const SCHEMA_WITH_REFS = {
+  tools: [
+    {
+      name: 'find_many_companies',
+      inputSchema: {
+        type: 'object',
+        properties: { filter: { $ref: '#/$defs/__schema0' } },
+        $defs: { __schema0: { type: 'object' } },
+      },
+    },
+  ],
+};
+
+const getFirstToolResultPart = (
+  prompt: LanguageModelV3Prompt,
+): LanguageModelV3ToolResultPart => {
+  const toolMessage = prompt.find((message) => message.role === 'tool');
+
+  if (toolMessage?.role !== 'tool') {
+    throw new Error('Expected a tool message in the prompt');
+  }
+
+  const toolResultPart = toolMessage.content.find(
+    (part) => part.type === 'tool-result',
+  );
+
+  if (toolResultPart?.type !== 'tool-result') {
+    throw new Error('Expected a tool-result part in the tool message');
+  }
+
+  return toolResultPart;
+};
+
+describe('sanitizeToolResultRefs', () => {
+  it('should serialize tool-result json output containing $ref/$defs to text', () => {
+    const prompt: LanguageModelV3Prompt = [
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'learn_tools',
+            output: { type: 'json', value: SCHEMA_WITH_REFS },
+          },
+        ],
+      },
+    ];
+
+    const part = getFirstToolResultPart(sanitizeToolResultRefs(prompt));
+
+    if (part.output.type !== 'text') {
+      throw new Error('Expected the output to be serialized to text');
+    }
+
+    expect(JSON.parse(part.output.value)).toEqual(SCHEMA_WITH_REFS);
+  });
+
+  it('should preserve output-level providerOptions when serializing to text', () => {
+    const providerOptions = { google: { cacheControl: { type: 'ephemeral' } } };
+    const prompt: LanguageModelV3Prompt = [
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'learn_tools',
+            output: { type: 'json', value: SCHEMA_WITH_REFS, providerOptions },
+          },
+        ],
+      },
+    ];
+
+    const part = getFirstToolResultPart(sanitizeToolResultRefs(prompt));
+
+    if (part.output.type !== 'text') {
+      throw new Error('Expected the output to be serialized to text');
+    }
+
+    expect(part.output.providerOptions).toEqual(providerOptions);
+  });
+
+  it('should convert error-json output containing refs to error-text', () => {
+    const prompt: LanguageModelV3Prompt = [
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'learn_tools',
+            output: { type: 'error-json', value: SCHEMA_WITH_REFS },
+          },
+        ],
+      },
+    ];
+
+    const part = getFirstToolResultPart(sanitizeToolResultRefs(prompt));
+
+    if (part.output.type !== 'error-text') {
+      throw new Error('Expected the output to be serialized to error-text');
+    }
+
+    expect(JSON.parse(part.output.value)).toEqual(SCHEMA_WITH_REFS);
+  });
+
+  it('should leave tool results without refs untouched', () => {
+    const value = { success: true, result: { id: '1', name: 'Acme' } };
+    const prompt: LanguageModelV3Prompt = [
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'find_one_company',
+            output: { type: 'json', value },
+          },
+        ],
+      },
+    ];
+
+    const part = getFirstToolResultPart(sanitizeToolResultRefs(prompt));
+
+    expect(part.output.type).toBe('json');
+    expect(part.output).toEqual({ type: 'json', value });
+  });
+
+  it('should not touch non-tool messages', () => {
+    const prompt: LanguageModelV3Prompt = [
+      { role: 'system', content: 'You are helpful.' },
+      {
+        role: 'user',
+        content: [{ type: 'text', text: 'hello' }],
+      },
+    ];
+
+    expect(sanitizeToolResultRefs(prompt)).toEqual(prompt);
+  });
+});
@@ -0,0 +1,55 @@
+import {
+  type LanguageModelV3Prompt,
+  type LanguageModelV3ToolResultPart,
+} from '@ai-sdk/provider';
+import { isArray, isObject, isString } from '@sniptt/guards';
+
+const containsJsonSchemaDefsRef = (value: unknown): boolean => {
+  if (isArray(value)) {
+    return value.some(containsJsonSchemaDefsRef);
+  }
+
+  if (!isObject(value)) {
+    return false;
+  }
+
+  return (
+    ('$ref' in value && isString(value.$ref)) ||
+    '$defs' in value ||
+    Object.values(value).some(containsJsonSchemaDefsRef)
+  );
+};
+
+const sanitizeToolResultPart = (
+  part: LanguageModelV3ToolResultPart,
+): LanguageModelV3ToolResultPart => {
+  if (
+    (part.output.type !== 'json' && part.output.type !== 'error-json') ||
+    !containsJsonSchemaDefsRef(part.output.value)
+  ) {
+    return part;
+  }
+
+  const value = JSON.stringify(part.output.value);
+
+  const providerOptions = part.output.providerOptions
+    ? { providerOptions: part.output.providerOptions }
+    : {};
+
+  return {
+    ...part,
+    output:
+      part.output.type === 'error-json'
+        ? { type: 'error-text', value, ...providerOptions }
+        : { type: 'text', value, ...providerOptions },
+  };
+};
+
+export const sanitizeToolResultRefs = (
+  prompt: LanguageModelV3Prompt,
+): LanguageModelV3Prompt =>
+  prompt.map((message) =>
+    message.role === 'tool'
+      ? { ...message, content: message.content.map(sanitizeToolResultPart) }
+      : message,
+  );