From 1b7dc0367e3d9e0780b55965f4a448adc077fd98 Mon Sep 17 00:00:00 2001 From: Abdul Rahman <81605929+abdulrahmancodes@users.noreply.github.com> Date: Mon, 22 Jun 2026 12:46:59 +0530 Subject: [PATCH] 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. Review in cubic --------- Co-authored-by: Charles Bochet --- ...tize-gemini-tool-result-refs.middleware.ts | 12 ++ .../services/sdk-provider-factory.service.ts | 21 ++- .../sanitize-tool-result-refs.util.spec.ts | 148 ++++++++++++++++++ .../utils/sanitize-tool-result-refs.util.ts | 55 +++++++ 4 files changed, 232 insertions(+), 4 deletions(-) create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-models/middleware/sanitize-gemini-tool-result-refs.middleware.ts create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/__tests__/sanitize-tool-result-refs.util.spec.ts create mode 100644 packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/sanitize-tool-result-refs.util.ts diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/middleware/sanitize-gemini-tool-result-refs.middleware.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/middleware/sanitize-gemini-tool-result-refs.middleware.ts new file mode 100644 index 0000000000..acaec10604 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/middleware/sanitize-gemini-tool-result-refs.middleware.ts @@ -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), + }), +}; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/services/sdk-provider-factory.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/services/sdk-provider-factory.service.ts index 448f880fbf..85e6ceeb0e 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/services/sdk-provider-factory.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/services/sdk-provider-factory.service.ts @@ -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, }; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/__tests__/sanitize-tool-result-refs.util.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/__tests__/sanitize-tool-result-refs.util.spec.ts new file mode 100644 index 0000000000..34ffdbd5ca --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/__tests__/sanitize-tool-result-refs.util.spec.ts @@ -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); + }); +}); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/sanitize-tool-result-refs.util.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/sanitize-tool-result-refs.util.ts new file mode 100644 index 0000000000..126a70a3e5 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-models/utils/sanitize-tool-result-refs.util.ts @@ -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, + );