fix(plugin-ai): sanitize malformed raw tool calls (#9900)

packages/plugins/@nocobase/plugin-ai/src/server/__tests__/ai-message-sanitizer.test.ts

@@ -0,0 +1,114 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { AIMessage } from '@langchain/core/messages';
+import { describe, expect, it, vi } from 'vitest';
+import { sanitizeAdditionalKwargsForToolCalls, sanitizeLangChainAIMessage } from '../ai-employees/tool-call-sanitizer';
+import { convertAIMessage } from '../ai-employees/utils';
+
+describe('AI message tool call sanitizer', () => {
+  const rawToolCall = {
+    id: 'call_bad',
+    type: 'function',
+    function: {
+      name: 'aiEmployeeWorkflowTaskOutput',
+      arguments: '{"result":{"reference_reply":":"bad json"}',
+    },
+  };
+
+  it('should drop malformed raw tool calls when converting LangChain AI messages to stored messages', () => {
+    const logger = { warn: vi.fn() };
+    const aiMessage = new AIMessage({
+      id: 'ai_bad',
+      content: '',
+      additional_kwargs: {
+        tool_calls: [rawToolCall],
+        reasoning_content: 'keep this reasoning',
+      },
+    });
+
+    const values = convertAIMessage({
+      aiEmployee: {
+        employee: { username: 'assistant' },
+        skillSettings: { tools: [] },
+        logger,
+      } as never,
+      providerName: 'deepseek',
+      model: 'deepseek-v4-flash',
+      aiMessage,
+    });
+
+    expect(values.metadata.additional_kwargs).toEqual({
+      reasoning_content: 'keep this reasoning',
+    });
+    expect(logger.warn).toHaveBeenCalledWith(
+      'Discard malformed raw tool calls from AI message',
+      expect.objectContaining({
+        phase: 'convertAIMessage',
+        messageId: 'ai_bad',
+        rawToolCallCount: 1,
+        parsedToolCallCount: 0,
+        rawToolCallIds: ['call_bad'],
+        rawToolCallNames: ['aiEmployeeWorkflowTaskOutput'],
+      }),
+    );
+    expect(JSON.stringify(logger.warn.mock.calls[0][1])).not.toContain('reference_reply');
+  });
+
+  it('should drop malformed raw tool calls when formatting stored assistant messages', () => {
+    const { additionalKwargs } = sanitizeAdditionalKwargsForToolCalls(
+      {
+        tool_calls: [rawToolCall],
+        reasoning_content: 'keep this reasoning',
+      },
+      null,
+    );
+
+    expect(additionalKwargs).toEqual({
+      reasoning_content: 'keep this reasoning',
+    });
+  });
+
+  it('should keep raw tool calls when parsed tool calls are present', () => {
+    const result = sanitizeAdditionalKwargsForToolCalls(
+      {
+        tool_calls: [rawToolCall],
+        reasoning_content: 'keep this reasoning',
+      },
+      [{ id: 'call_bad', name: 'aiEmployeeWorkflowTaskOutput', args: {}, type: 'tool_call' }],
+    );
+
+    expect(result.changed).toBe(false);
+    expect(result.additionalKwargs).toEqual({
+      tool_calls: [rawToolCall],
+      reasoning_content: 'keep this reasoning',
+    });
+  });
+
+  it('should sanitize LangChain AI messages in place without rebuilding the message instance', () => {
+    const aiMessage = new AIMessage({
+      id: 'ai_bad',
+      content: '',
+      additional_kwargs: {
+        tool_calls: [rawToolCall],
+        reasoning_content: 'keep this reasoning',
+      },
+    });
+    const messageWithRuntimeField = aiMessage as AIMessage & { runtimeField?: string };
+    messageWithRuntimeField.runtimeField = 'keep runtime field';
+
+    const sanitizedMessage = sanitizeLangChainAIMessage(aiMessage);
+
+    expect(sanitizedMessage).toBe(aiMessage);
+    expect(messageWithRuntimeField.runtimeField).toBe('keep runtime field');
+    expect(aiMessage.additional_kwargs).toEqual({
+      reasoning_content: 'keep this reasoning',
+    });
+  });
+});

packages/plugins/@nocobase/plugin-ai/src/server/__tests__/tool-call-sanitizer-middleware.test.ts

@@ -0,0 +1,226 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { AIMessage, BaseMessage, HumanMessage, ToolMessage } from '@langchain/core/messages';
+import { FakeListChatModel } from '@langchain/core/utils/testing';
+import { MemorySaver } from '@langchain/langgraph';
+import { convertMessagesToCompletionsMessageParams } from '@langchain/openai';
+import { createAgent, createMiddleware } from 'langchain';
+import { describe, expect, it, vi } from 'vitest';
+import { toolCallSanitizerMiddleware } from '../ai-employees/middleware';
+
+describe('toolCallSanitizerMiddleware', () => {
+  const rawToolCall = {
+    id: 'call_bad',
+    type: 'function',
+    function: {
+      name: 'aiEmployeeWorkflowTaskOutput',
+      arguments: '{"result":{"reference_reply":":"bad json"}',
+    },
+  };
+
+  type MessagePatchHook = (state: {
+    messages: BaseMessage[];
+  }) => void | { messages?: BaseMessage[] } | Promise<void | { messages?: BaseMessage[] }>;
+
+  const getMessagePatchHook = (hook: unknown): MessagePatchHook => {
+    if (typeof hook === 'function') {
+      return hook as MessagePatchHook;
+    }
+    if (hook && typeof hook === 'object' && 'hook' in hook) {
+      const nestedHook = (hook as { hook?: unknown }).hook;
+      if (typeof nestedHook === 'function') {
+        return nestedHook as MessagePatchHook;
+      }
+    }
+    throw new Error('Middleware hook is not callable');
+  };
+
+  const getPatchMessages = (result: Awaited<ReturnType<MessagePatchHook>>) => {
+    if (!result || !('messages' in result)) {
+      return [];
+    }
+    return result.messages ?? [];
+  };
+
+  it('should remove malformed raw tool calls before checkpoint persistence and next model request', async () => {
+    const model = new FakeListChatModel({
+      responses: [
+        new AIMessage({
+          id: 'ai_bad',
+          content: '',
+          additional_kwargs: {
+            tool_calls: [rawToolCall],
+            reasoning_content: 'keep this reasoning',
+          },
+        }),
+      ],
+    });
+    const logger = { warn: vi.fn() };
+    const agent = createAgent({
+      model,
+      tools: [],
+      middleware: [toolCallSanitizerMiddleware({ logger })],
+      checkpointer: new MemorySaver(),
+    });
+    const config = { configurable: { thread_id: 'malformed-tool-call' } };
+
+    await agent.invoke({ messages: [{ role: 'user', content: 'run' }] }, config);
+
+    const state = await agent.getState(config);
+    const lastMessage = state.values.messages.at(-1) as AIMessage;
+
+    expect(lastMessage.type).toBe('ai');
+    expect(lastMessage.tool_calls).toEqual([]);
+    expect(lastMessage.invalid_tool_calls).toHaveLength(1);
+    expect(lastMessage.additional_kwargs).toEqual({
+      reasoning_content: 'keep this reasoning',
+    });
+
+    const requestMessages = convertMessagesToCompletionsMessageParams({
+      messages: state.values.messages,
+    });
+    expect(requestMessages.at(-1)).toMatchObject({
+      role: 'assistant',
+      content: '',
+    });
+    expect(requestMessages.at(-1)).not.toHaveProperty('tool_calls');
+    expect(logger.warn).toHaveBeenCalledWith(
+      'Discard malformed raw tool calls from AI message',
+      expect.objectContaining({
+        phase: 'afterModel',
+        messageId: 'ai_bad',
+        rawToolCallCount: 1,
+        parsedToolCallCount: 0,
+        rawToolCallIds: ['call_bad'],
+        rawToolCallNames: ['aiEmployeeWorkflowTaskOutput'],
+      }),
+    );
+    expect(JSON.stringify(logger.warn.mock.calls[0][1])).not.toContain('reference_reply');
+  });
+
+  it('should remove malformed raw tool calls from existing state before model requests', async () => {
+    let capturedMessages: BaseMessage[] = [];
+    const model = new FakeListChatModel({
+      responses: ['done'],
+    });
+    const agent = createAgent({
+      model,
+      tools: [],
+      middleware: [
+        toolCallSanitizerMiddleware(),
+        createMiddleware({
+          name: 'CaptureMessagesMiddleware',
+          wrapModelCall: (request, handler) => {
+            capturedMessages = request.messages;
+            return handler(request);
+          },
+        }),
+      ],
+      checkpointer: new MemorySaver(),
+    });
+
+    await agent.invoke(
+      {
+        messages: [
+          new AIMessage({
+            id: 'ai_bad_history',
+            content: '',
+            additional_kwargs: {
+              tool_calls: [rawToolCall],
+              reasoning_content: 'keep this reasoning',
+            },
+          }),
+          { role: 'user', content: 'continue' },
+        ],
+      },
+      { configurable: { thread_id: 'malformed-tool-call-history' } },
+    );
+
+    const requestMessages = convertMessagesToCompletionsMessageParams({
+      messages: capturedMessages,
+    });
+    expect(requestMessages[0]).toMatchObject({
+      role: 'assistant',
+      content: '',
+    });
+    expect(requestMessages[0]).not.toHaveProperty('tool_calls');
+    expect(requestMessages).toHaveLength(2);
+    expect(requestMessages[1]).toMatchObject({
+      role: 'user',
+      content: 'continue',
+    });
+  });
+
+  it('should build a patch without replacing non-AI messages', async () => {
+    const badAIMessage = new AIMessage({
+      id: 'ai_bad_middle',
+      content: '',
+      additional_kwargs: {
+        tool_calls: [rawToolCall],
+        reasoning_content: 'keep this reasoning',
+      },
+    });
+    const humanMessage = new HumanMessage({ id: 'human_1', content: 'continue' });
+    const toolMessage = new ToolMessage({
+      id: 'tool_1',
+      content: 'tool result',
+      tool_call_id: 'call_existing',
+    });
+    const middleware = toolCallSanitizerMiddleware();
+    const beforeModel = getMessagePatchHook(middleware.beforeModel);
+    const messages = [humanMessage, badAIMessage, toolMessage];
+
+    const result = await beforeModel({
+      messages,
+    });
+    const patchMessages = getPatchMessages(result);
+
+    expect(messages).toEqual([humanMessage, badAIMessage, toolMessage]);
+    expect(patchMessages).toHaveLength(1);
+    expect(patchMessages[0]).not.toBe(badAIMessage);
+    expect(AIMessage.isInstance(patchMessages[0])).toBe(true);
+    expect(patchMessages[0].id).toBe(badAIMessage.id);
+    expect(badAIMessage.additional_kwargs).toEqual({
+      reasoning_content: 'keep this reasoning',
+    });
+  });
+
+  it('should sanitize malformed AI messages even when they are not the last message after model', async () => {
+    const badAIMessage = new AIMessage({
+      id: 'ai_bad_not_last',
+      content: '',
+      additional_kwargs: {
+        tool_calls: [rawToolCall],
+        reasoning_content: 'keep this reasoning',
+      },
+    });
+    const middleware = toolCallSanitizerMiddleware();
+    const afterModel = getMessagePatchHook(middleware.afterModel);
+    const messages = [
+      new HumanMessage({ id: 'human_1', content: 'start' }),
+      badAIMessage,
+      new HumanMessage({ id: 'human_2', content: 'not last' }),
+    ];
+
+    const result = await afterModel({
+      messages,
+    });
+    const patchMessages = getPatchMessages(result);
+
+    expect(messages[1]).toBe(badAIMessage);
+    expect(patchMessages).toHaveLength(1);
+    expect(patchMessages[0]).not.toBe(badAIMessage);
+    expect(AIMessage.isInstance(patchMessages[0])).toBe(true);
+    expect(patchMessages[0].id).toBe(badAIMessage.id);
+    expect(badAIMessage.additional_kwargs).toEqual({
+      reasoning_content: 'keep this reasoning',
+    });
+  });
+});

packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts

@@ -23,6 +23,7 @@ import type { AIEmployee as AIEmployeeType } from '../../collections/ai-employee
 import {
   conversationMiddleware,
   skillToolBindingMiddleware,
+  toolCallSanitizerMiddleware,
   toolCallStatusMiddleware,
   toolInteractionMiddleware,
   workflowHistoryMiddleware,
@@ -39,6 +40,7 @@ import { LLMResult } from '@langchain/core/outputs';
 import { Context } from '@nocobase/actions';
 import { listAccessibleAIEmployees, serializeEmployeeSummary } from '../../ai/tools/sub-agents/shared';
 import { LLMStreamCached } from '../manager/llm-stream-manager';
+import { sanitizeAdditionalKwargsForToolCalls } from './tool-call-sanitizer';
 
 export interface ModelRef {
   llmService: string;
@@ -1288,7 +1290,15 @@ If information is missing, clearly state it in the summary.</Important>`;
         role: 'assistant',
         content,
         tool_calls: msg.toolCalls,
-        additional_kwargs: msg.metadata?.additional_kwargs,
+        additional_kwargs: sanitizeAdditionalKwargsForToolCalls(msg.metadata?.additional_kwargs, msg.toolCalls, {
+          onDiscard: (info) => {
+            this.logger.warn('Discard malformed raw tool calls from AI message', {
+              phase: 'formatMessages',
+              messageId: msg.metadata?.id,
+              ...info,
+            });
+          },
+        }).additionalKwargs,
       });
     }
 
@@ -1533,6 +1543,7 @@ If information is missing, clearly state it in the summary.</Important>`;
       toolCallStatusMiddleware(this),
       ...(inWorkflow ? [workflowHistoryMiddleware(this, this.db)] : []),
       conversationMiddleware(this, { providerName, llmService, model, messageId, agentThread }),
+      toolCallSanitizerMiddleware({ logger: this.logger }),
     ];
   }
 

packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/index.ts

@@ -9,5 +9,6 @@
 
 export * from './conversation';
 export * from './skill-tools';
+export * from './tool-call-sanitizer';
 export * from './tools';
 export * from './workflow-history';

packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/tool-call-sanitizer.ts

@@ -0,0 +1,80 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { AIMessage, BaseMessage } from '@langchain/core/messages';
+import { createMiddleware } from 'langchain';
+import { sanitizeLangChainAIMessage } from '../tool-call-sanitizer';
+
+type ToolCallSanitizerLogger = {
+  warn: (message: string, meta?: Record<string, unknown>) => void;
+};
+
+type ToolCallSanitizerMiddlewareOptions = {
+  logger?: ToolCallSanitizerLogger;
+};
+
+const buildMessageUpdates = (
+  messages: readonly unknown[],
+  options: ToolCallSanitizerMiddlewareOptions,
+  phase: 'beforeModel' | 'afterModel',
+): BaseMessage[] => {
+  const updates: BaseMessage[] = [];
+
+  for (const message of messages) {
+    if (!AIMessage.isInstance(message) || !message.id) {
+      continue;
+    }
+    const sanitizedMessage = sanitizeLangChainAIMessage(message, {
+      onDiscard: (info) => {
+        options.logger?.warn('Discard malformed raw tool calls from AI message', {
+          phase,
+          messageId: message.id,
+          invalidToolCallCount: message.invalid_tool_calls?.length ?? 0,
+          ...info,
+        });
+      },
+    });
+    if (!sanitizedMessage) {
+      continue;
+    }
+    updates.push(
+      new AIMessage({
+        id: sanitizedMessage.id,
+        content: sanitizedMessage.content,
+        name: sanitizedMessage.name,
+        additional_kwargs: sanitizedMessage.additional_kwargs,
+        response_metadata: sanitizedMessage.response_metadata,
+        tool_calls: sanitizedMessage.tool_calls,
+        invalid_tool_calls: sanitizedMessage.invalid_tool_calls,
+        usage_metadata: sanitizedMessage.usage_metadata,
+      }),
+    );
+  }
+
+  return updates;
+};
+
+export const toolCallSanitizerMiddleware = (options: ToolCallSanitizerMiddlewareOptions = {}) =>
+  createMiddleware<never>({
+    name: 'ToolCallSanitizerMiddleware',
+    beforeModel: (state) => {
+      const updates = buildMessageUpdates(state.messages ?? [], options, 'beforeModel');
+      if (!updates.length) {
+        return;
+      }
+      return { messages: updates };
+    },
+    afterModel: (state) => {
+      const updates = buildMessageUpdates(state.messages ?? [], options, 'afterModel');
+      if (!updates.length) {
+        return;
+      }
+      return { messages: updates };
+    },
+  });

packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/tool-call-sanitizer.ts

@@ -0,0 +1,95 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { AIMessage } from '@langchain/core/messages';
+
+type AdditionalKwargs = Record<string, unknown>;
+
+type DiscardedToolCallsInfo = {
+  rawToolCallCount: number;
+  parsedToolCallCount: number;
+  rawToolCallIds: string[];
+  rawToolCallNames: string[];
+};
+
+type SanitizeAdditionalKwargsOptions = {
+  onDiscard?: (info: DiscardedToolCallsInfo) => void;
+};
+
+type SanitizeAdditionalKwargsResult = {
+  changed: boolean;
+  additionalKwargs?: AdditionalKwargs;
+};
+
+const hasRawToolCalls = (additionalKwargs?: AdditionalKwargs) =>
+  Array.isArray(additionalKwargs?.tool_calls) && additionalKwargs.tool_calls.length > 0;
+
+const hasParsedToolCalls = (toolCalls?: unknown[] | null) => Array.isArray(toolCalls) && toolCalls.length > 0;
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value);
+
+const collectRawToolCallValues = (toolCalls: unknown[], key: 'id' | 'name') =>
+  toolCalls
+    .map((toolCall) => {
+      if (!isRecord(toolCall)) {
+        return null;
+      }
+      if (key === 'name') {
+        const fn = toolCall.function;
+        if (isRecord(fn) && typeof fn.name === 'string') {
+          return fn.name;
+        }
+      }
+      const value = toolCall[key];
+      return typeof value === 'string' ? value : null;
+    })
+    .filter((value): value is string => Boolean(value));
+
+export const sanitizeAdditionalKwargsForToolCalls = (
+  additionalKwargs: AdditionalKwargs | undefined,
+  parsedToolCalls?: unknown[] | null,
+  options: SanitizeAdditionalKwargsOptions = {},
+): SanitizeAdditionalKwargsResult => {
+  if (!hasRawToolCalls(additionalKwargs) || hasParsedToolCalls(parsedToolCalls)) {
+    return {
+      changed: false,
+      additionalKwargs,
+    };
+  }
+
+  const rawToolCalls = additionalKwargs.tool_calls as unknown[];
+  options.onDiscard?.({
+    rawToolCallCount: rawToolCalls.length,
+    parsedToolCallCount: parsedToolCalls?.length ?? 0,
+    rawToolCallIds: collectRawToolCallValues(rawToolCalls, 'id'),
+    rawToolCallNames: collectRawToolCallValues(rawToolCalls, 'name'),
+  });
+
+  const sanitized = { ...additionalKwargs };
+  delete sanitized.tool_calls;
+
+  return {
+    changed: true,
+    additionalKwargs: Object.keys(sanitized).length > 0 ? sanitized : undefined,
+  };
+};
+
+export const sanitizeLangChainAIMessage = (message: AIMessage, options: SanitizeAdditionalKwargsOptions = {}) => {
+  const sanitized = sanitizeAdditionalKwargsForToolCalls(message.additional_kwargs, message.tool_calls, options);
+
+  if (!sanitized.changed) {
+    return null;
+  }
+
+  message.additional_kwargs = sanitized.additionalKwargs ?? {};
+  message.lc_kwargs.additional_kwargs = message.additional_kwargs;
+
+  return message;
+};

packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/utils.ts

@@ -10,6 +10,7 @@
 import { AIMessage, HumanMessage, ToolMessage } from 'langchain';
 import { AIMessageContent, AIMessageInput } from '../types';
 import { AIEmployee } from './ai-employee';
+import { sanitizeAdditionalKwargsForToolCalls } from './tool-call-sanitizer';
 
 export const convertAIMessage = ({
   aiEmployee,
@@ -84,8 +85,18 @@ export const convertAIMessage = ({
   if (aiMessage.response_metadata) {
     values.metadata.response_metadata = aiMessage.response_metadata;
   }
-  if (aiMessage.additional_kwargs) {
-    values.metadata.additional_kwargs = aiMessage.additional_kwargs;
+  const additionalKwargs = sanitizeAdditionalKwargsForToolCalls(aiMessage.additional_kwargs, toolCalls, {
+    onDiscard: (info) => {
+      aiEmployee.logger?.warn('Discard malformed raw tool calls from AI message', {
+        phase: 'convertAIMessage',
+        messageId: aiMessage.id,
+        invalidToolCallCount: aiMessage.invalid_tool_calls?.length ?? 0,
+        ...info,
+      });
+    },
+  }).additionalKwargs;
+  if (additionalKwargs) {
+    values.metadata.additional_kwargs = additionalKwargs;
   }
 
   return values;