diff --git a/packages/twenty-front/src/modules/ai/components/AiChatAssistantMessageRenderer.tsx b/packages/twenty-front/src/modules/ai/components/AiChatAssistantMessageRenderer.tsx index 7b246f4d4b..f348b0ce43 100644 --- a/packages/twenty-front/src/modules/ai/components/AiChatAssistantMessageRenderer.tsx +++ b/packages/twenty-front/src/modules/ai/components/AiChatAssistantMessageRenderer.tsx @@ -9,7 +9,7 @@ import { ToolStepRenderer } from '@/ai/components/ToolStepRenderer'; import { groupContiguousThinkingStepParts } from '@/ai/utils/groupContiguousThinkingStepParts'; import { isCodeInterpreterToolPart } from '@/ai/utils/isCodeInterpreterToolPart'; import { styled } from '@linaria/react'; -import { isToolUIPart, type ToolUIPart } from 'ai'; +import { isToolUIPart } from 'ai'; import { type ExtendedUIMessagePart } from 'twenty-shared/ai'; import { useContext } from 'react'; import { ThemeContext, themeCssVariables } from 'twenty-ui/theme-constants'; @@ -75,13 +75,8 @@ const MessagePartRenderer = ({ /> ); default: - if (isToolUIPart(part) === true && part.type !== 'dynamic-tool') { - return ( - - ); + if (isToolUIPart(part)) { + return ; } return null; } diff --git a/packages/twenty-front/src/modules/ai/components/ThinkingStepsDisplay.tsx b/packages/twenty-front/src/modules/ai/components/ThinkingStepsDisplay.tsx index ad9ae1b4e9..d077c17f1c 100644 --- a/packages/twenty-front/src/modules/ai/components/ThinkingStepsDisplay.tsx +++ b/packages/twenty-front/src/modules/ai/components/ThinkingStepsDisplay.tsx @@ -1,7 +1,7 @@ import { styled } from '@linaria/react'; import { plural, t } from '@lingui/core/macro'; import { useState } from 'react'; -import { type ToolUIPart } from 'ai'; +import { type DynamicToolUIPart, getToolName, type ToolUIPart } from 'ai'; import { isDefined } from 'twenty-shared/utils'; import { IconChevronRight, @@ -258,12 +258,12 @@ const ThinkingToolStepRow = ({ rowIndex, }: { isActive: boolean; - part: ToolUIPart; + part: ToolUIPart | DynamicToolUIPart; rowIndex: number; }) => { const { copyToClipboard } = useCopyToClipboard(); const [isExpanded, setIsExpanded] = useState(false); - const rawToolName = part.type.split('-')[1]; + const rawToolName = getToolName(part); const { resolvedInput: toolInput, resolvedToolName } = resolveToolInput( part.input, rawToolName, diff --git a/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx b/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx index 01ea41444b..d1a0e989be 100644 --- a/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx +++ b/packages/twenty-front/src/modules/ai/components/ToolStepRenderer.tsx @@ -14,7 +14,7 @@ import { } from '@/ai/utils/getToolDisplayMessage'; import { getToolIcon } from '@/ai/utils/getToolIcon'; import { useLingui } from '@lingui/react/macro'; -import { type ToolUIPart } from 'ai'; +import { type DynamicToolUIPart, getToolName, type ToolUIPart } from 'ai'; import { isDefined } from 'twenty-shared/utils'; import { type JsonValue } from 'type-fest'; import { useCopyToClipboard } from '~/hooks/useCopyToClipboard'; @@ -131,7 +131,7 @@ export const ToolStepRenderer = ({ toolPart, isStreaming, }: { - toolPart: ToolUIPart; + toolPart: ToolUIPart | DynamicToolUIPart; isStreaming: boolean; }) => { const { theme } = useContext(ThemeContext); @@ -140,8 +140,8 @@ export const ToolStepRenderer = ({ const [isExpanded, setIsExpanded] = useState(false); const [activeTab, setActiveTab] = useState('output'); - const { input, output, type, errorText } = toolPart; - const rawToolName = type.split('-')[1]; + const { input, output, errorText } = toolPart; + const rawToolName = getToolName(toolPart); const { resolvedInput: toolInput, resolvedToolName: toolName } = resolveToolInput(input, rawToolName); diff --git a/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx b/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx index 00e4c36ad9..5408684f08 100644 --- a/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx +++ b/packages/twenty-front/src/modules/ai/components/__tests__/AiChatAssistantMessageRenderer.test.tsx @@ -233,4 +233,28 @@ describe('AiChatAssistantMessageRenderer', () => { ); expect(screen.getByTestId('code-execution-display')).toBeInTheDocument(); }); + + it('should group a dynamic-tool part (native web search) into ThinkingStepsDisplay', () => { + const messageParts = [ + { + type: 'dynamic-tool', + toolName: 'web_search', + toolCallId: 'dyn-1', + input: { query: 'crm software' }, + output: { result: { ok: true } }, + state: 'output-available', + providerExecuted: true, + }, + { + type: 'text', + text: 'Final answer', + }, + ] as ExtendedUIMessagePart[]; + + renderAssistantRenderer(messageParts); + + expect(screen.getByTestId('thinking-steps-display')).toHaveTextContent( + 'thinking-1-answer-started', + ); + }); }); diff --git a/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts b/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts index 14ab20ca07..9bd828d3d9 100644 --- a/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts +++ b/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts @@ -1,4 +1,4 @@ -import { type ReasoningUIPart, type ToolUIPart } from 'ai'; +import { type ReasoningUIPart } from 'ai'; import { type ExtendedFileUIPart, type ExtendedUIMessagePart, @@ -63,9 +63,13 @@ export const mapDBPartToUIMessagePart = ( }; default: { - if (part.type.includes('tool-') === true) { + const isStaticToolPart = part.type.startsWith('tool-'); + const isDynamicToolPart = part.type === 'dynamic-tool'; + + if (isStaticToolPart || isDynamicToolPart) { return { - type: part.type as `tool-${string}`, + type: part.type as `tool-${string}` | 'dynamic-tool', + ...(isDynamicToolPart && { toolName: part.toolName ?? '' }), toolCallId: part.toolCallId!, input: part.toolInput ?? {}, output: part.toolOutput, @@ -74,7 +78,10 @@ export const mapDBPartToUIMessagePart = ( ...(part.providerExecuted != null && { providerExecuted: part.providerExecuted, }), - } as ToolUIPart; + ...(part.providerMetadata != null && { + callProviderMetadata: part.providerMetadata, + }), + } as ExtendedUIMessagePart; } } throw new Error(`Unsupported part type: ${part.type}`); diff --git a/packages/twenty-front/src/modules/ai/utils/thinkingStepPart.ts b/packages/twenty-front/src/modules/ai/utils/thinkingStepPart.ts index cff13a070b..8ab2a9bca8 100644 --- a/packages/twenty-front/src/modules/ai/utils/thinkingStepPart.ts +++ b/packages/twenty-front/src/modules/ai/utils/thinkingStepPart.ts @@ -1,3 +1,7 @@ -import { type ReasoningUIPart, type ToolUIPart } from 'ai'; +import { + type DynamicToolUIPart, + type ReasoningUIPart, + type ToolUIPart, +} from 'ai'; -export type ThinkingStepPart = ReasoningUIPart | ToolUIPart; +export type ThinkingStepPart = ReasoningUIPart | ToolUIPart | DynamicToolUIPart; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/__tests__/map-message-parts.dynamic-tool.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/__tests__/map-message-parts.dynamic-tool.spec.ts new file mode 100644 index 0000000000..e795e7abd3 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/__tests__/map-message-parts.dynamic-tool.spec.ts @@ -0,0 +1,137 @@ +import { type ExtendedUIMessagePart } from 'twenty-shared/ai'; + +import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity'; +import { mapDBPartToUIMessagePart } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart'; +import { mapUIMessagePartsToDBParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts'; + +const dynamicToolPart = ( + overrides: Record = {}, +): ExtendedUIMessagePart => + ({ + type: 'dynamic-tool', + toolName: 'unknown_remote_tool', + toolCallId: 'call_dyn_1', + state: 'output-available', + input: { query: 'hello' }, + output: { ok: true }, + ...overrides, + }) as unknown as ExtendedUIMessagePart; + +const staticToolPart = ( + overrides: Record = {}, +): ExtendedUIMessagePart => + ({ + type: 'tool-execute_tool', + toolCallId: 'call_static_1', + state: 'output-available', + input: { name: 'foo' }, + output: { ok: true }, + ...overrides, + }) as unknown as ExtendedUIMessagePart; + +describe('AgentMessagePart mappers — dynamic-tool support', () => { + it('persists a dynamic-tool part without throwing', () => { + expect(() => + mapUIMessagePartsToDBParts( + [dynamicToolPart()], + 'message-1', + 'workspace-1', + ), + ).not.toThrow(); + }); + + it('stores the tool name on the row for dynamic-tool parts', () => { + const [row] = mapUIMessagePartsToDBParts( + [dynamicToolPart()], + 'message-1', + 'workspace-1', + ); + + expect(row).toMatchObject({ + type: 'dynamic-tool', + toolName: 'unknown_remote_tool', + toolCallId: 'call_dyn_1', + toolInput: { query: 'hello' }, + toolOutput: { ok: true }, + }); + }); + + it('stores the tool name on the row for static tool parts', () => { + const [row] = mapUIMessagePartsToDBParts( + [staticToolPart()], + 'message-1', + 'workspace-1', + ); + + expect(row).toMatchObject({ + type: 'tool-execute_tool', + toolName: 'execute_tool', + toolCallId: 'call_static_1', + }); + }); + + it('round-trips a dynamic-tool part through DB and back', () => { + const original = dynamicToolPart(); + + const [row] = mapUIMessagePartsToDBParts( + [original], + 'message-1', + 'workspace-1', + ); + const reloaded = mapDBPartToUIMessagePart(row as AgentMessagePartEntity); + + expect(reloaded).toEqual({ + type: 'dynamic-tool', + toolName: 'unknown_remote_tool', + toolCallId: 'call_dyn_1', + input: { query: 'hello' }, + output: { ok: true }, + errorText: '', + state: 'output-available', + }); + }); + + it('round-trips callProviderMetadata for provider-executed tools', () => { + const original = dynamicToolPart({ + providerExecuted: true, + callProviderMetadata: { anthropic: { encryptedContent: 'abc123' } }, + }); + + const [row] = mapUIMessagePartsToDBParts( + [original], + 'message-1', + 'workspace-1', + ); + + expect(row).toMatchObject({ + providerExecuted: true, + providerMetadata: { anthropic: { encryptedContent: 'abc123' } }, + }); + + const reloaded = mapDBPartToUIMessagePart(row as AgentMessagePartEntity); + + expect(reloaded).toMatchObject({ + providerExecuted: true, + callProviderMetadata: { anthropic: { encryptedContent: 'abc123' } }, + }); + }); + + it('round-trips a static tool part through DB and back', () => { + const original = staticToolPart(); + + const [row] = mapUIMessagePartsToDBParts( + [original], + 'message-1', + 'workspace-1', + ); + const reloaded = mapDBPartToUIMessagePart(row as AgentMessagePartEntity); + + expect(reloaded).toMatchObject({ + type: 'tool-execute_tool', + toolCallId: 'call_static_1', + input: { name: 'foo' }, + output: { ok: true }, + }); + expect(reloaded).not.toHaveProperty('toolName'); + }); +}); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart.ts index 9523ba989a..5c943ea248 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart.ts @@ -56,9 +56,15 @@ export const mapDBPartToUIMessagePart = ( case 'data-routing-status': return null; default: { - if (part.type.includes('tool-') && part.toolCallId) { + const isStaticToolPart = + part.type.startsWith('tool-') && part.toolCallId !== null; + const isDynamicToolPart = + part.type === 'dynamic-tool' && part.toolCallId !== null; + + if (isStaticToolPart || isDynamicToolPart) { return { type: part.type, + ...(isDynamicToolPart && { toolName: part.toolName ?? '' }), toolCallId: part.toolCallId, input: part.toolInput ?? {}, output: part.toolOutput, @@ -67,6 +73,9 @@ export const mapDBPartToUIMessagePart = ( ...(part.providerExecuted != null && { providerExecuted: part.providerExecuted, }), + ...(part.providerMetadata != null && { + callProviderMetadata: part.providerMetadata, + }), } as ExtendedUIMessagePart; } diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts index 73e1f51a9d..3435d35fd7 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts @@ -1,4 +1,4 @@ -import { type ToolUIPart } from 'ai'; +import { getToolName, isToolUIPart } from 'ai'; import { isExtendedFileUIPart, type ExtendedUIMessagePart, @@ -6,10 +6,6 @@ import { import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity'; -const isToolPart = (part: ExtendedUIMessagePart): part is ToolUIPart => { - return part.type.includes('tool-') && 'toolCallId' in part; -}; - export const mapUIMessagePartsToDBParts = ( uiMessageParts: ExtendedUIMessagePart[], messageId: string, @@ -80,23 +76,25 @@ export const mapUIMessagePartsToDBParts = ( case 'data-thread-title': // Thread title is a transient notification for the client return null; - default: - { - if (isToolPart(part)) { - const { toolCallId, input, output, errorText, state } = part; - - return { - ...basePart, - toolCallId: toolCallId, - toolInput: input, - toolOutput: output, - errorMessage: errorText, - state, - providerExecuted: part.providerExecuted ?? null, - }; - } + default: { + if (isToolUIPart(part)) { + return { + ...basePart, + toolName: getToolName(part), + toolCallId: part.toolCallId, + toolInput: part.input, + toolOutput: part.output, + errorMessage: part.errorText, + state: part.state, + providerExecuted: part.providerExecuted ?? null, + providerMetadata: part.callProviderMetadata ?? null, + }; } - throw new Error(`Unsupported part type: ${part.type}`); + + throw new Error( + `Unsupported part type: ${(part as { type: string }).type}`, + ); + } } }) .filter((part): part is Partial => part !== null);