From 180046a3c57b5fda653ba5591bcbc83046469c64 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 11 Sep 2025 14:34:07 -0400 Subject: [PATCH] =?UTF-8?q?=E2=9C=82=EF=B8=8F=20refactor:=20Artifacts=20an?= =?UTF-8?q?d=20Tool=20Callbacks=20to=20Pass=20UI=20Resources=20(#9581)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ✂️ refactor: use artifacts and callbacks to pass UI resources * chore: imports * refactor: Update UIResource type imports and definitions across components and tests * refactor: Update ToolCallInfo test data structure and enhance TAttachment type definition --------- Co-authored-by: Samuel Path --- .../agents/__tests__/callbacks.spec.js | 342 ++++++++++++++++ api/server/controllers/agents/callbacks.js | 24 ++ client/src/common/types.ts | 7 - .../Chat/Messages/Content/ToolCall.tsx | 1 + .../Chat/Messages/Content/ToolCallInfo.tsx | 30 +- .../Messages/Content/UIResourceCarousel.tsx | 4 +- .../Content/__tests__/ToolCall.test.tsx | 382 ++++++++++++++++++ .../Content/__tests__/ToolCallInfo.test.tsx | 254 +++++++----- .../__tests__/UIResourceCarousel.test.tsx | 6 +- .../api/src/mcp/__tests__/parsers.test.ts | 138 +++++-- packages/api/src/mcp/parsers.ts | 24 +- packages/api/src/mcp/types/index.ts | 41 +- packages/data-provider/src/schemas.ts | 17 +- .../data-provider/src/types/assistants.ts | 1 + 14 files changed, 1072 insertions(+), 199 deletions(-) create mode 100644 api/server/controllers/agents/__tests__/callbacks.spec.js create mode 100644 client/src/components/Chat/Messages/Content/__tests__/ToolCall.test.tsx diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js new file mode 100644 index 000000000..25f00bab8 --- /dev/null +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -0,0 +1,342 @@ +const { Tools } = require('librechat-data-provider'); + +// Mock all dependencies before requiring the module +jest.mock('nanoid', () => ({ + nanoid: jest.fn(() => 'mock-id'), +})); + +jest.mock('@librechat/api', () => ({ + sendEvent: jest.fn(), +})); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { + error: jest.fn(), + }, +})); + +jest.mock('@librechat/agents', () => ({ + EnvVar: { CODE_API_KEY: 'CODE_API_KEY' }, + Providers: { GOOGLE: 'google' }, + GraphEvents: {}, + getMessageId: jest.fn(), + ToolEndHandler: jest.fn(), + handleToolCalls: jest.fn(), + ChatModelStreamHandler: jest.fn(), +})); + +jest.mock('~/server/services/Files/Citations', () => ({ + processFileCitations: jest.fn(), +})); + +jest.mock('~/server/services/Files/Code/process', () => ({ + processCodeOutput: jest.fn(), +})); + +jest.mock('~/server/services/Tools/credentials', () => ({ + loadAuthValues: jest.fn(), +})); + +jest.mock('~/server/services/Files/process', () => ({ + saveBase64Image: jest.fn(), +})); + +describe('createToolEndCallback', () => { + let req, res, artifactPromises, createToolEndCallback; + let logger; + + beforeEach(() => { + jest.clearAllMocks(); + + // Get the mocked logger + logger = require('@librechat/data-schemas').logger; + + // Now require the module after all mocks are set up + const callbacks = require('../callbacks'); + createToolEndCallback = callbacks.createToolEndCallback; + + req = { + user: { id: 'user123' }, + }; + res = { + headersSent: false, + write: jest.fn(), + }; + artifactPromises = []; + }); + + describe('ui_resources artifact handling', () => { + it('should process ui_resources artifact and return attachment when headers not sent', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const output = { + tool_call_id: 'tool123', + artifact: { + [Tools.ui_resources]: { + data: { + 0: { type: 'button', label: 'Click me' }, + 1: { type: 'input', placeholder: 'Enter text' }, + }, + }, + }, + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + + // Wait for all promises to resolve + const results = await Promise.all(artifactPromises); + + // When headers are not sent, it returns attachment without writing + expect(res.write).not.toHaveBeenCalled(); + + const attachment = results[0]; + expect(attachment).toEqual({ + type: Tools.ui_resources, + messageId: 'run456', + toolCallId: 'tool123', + conversationId: 'thread789', + [Tools.ui_resources]: { + 0: { type: 'button', label: 'Click me' }, + 1: { type: 'input', placeholder: 'Enter text' }, + }, + }); + }); + + it('should write to response when headers are already sent', async () => { + res.headersSent = true; + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const output = { + tool_call_id: 'tool123', + artifact: { + [Tools.ui_resources]: { + data: { + 0: { type: 'carousel', items: [] }, + }, + }, + }, + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + const results = await Promise.all(artifactPromises); + + expect(res.write).toHaveBeenCalled(); + expect(results[0]).toEqual({ + type: Tools.ui_resources, + messageId: 'run456', + toolCallId: 'tool123', + conversationId: 'thread789', + [Tools.ui_resources]: { + 0: { type: 'carousel', items: [] }, + }, + }); + }); + + it('should handle errors when processing ui_resources', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + // Mock res.write to throw an error + res.headersSent = true; + res.write.mockImplementation(() => { + throw new Error('Write failed'); + }); + + const output = { + tool_call_id: 'tool123', + artifact: { + [Tools.ui_resources]: { + data: { + 0: { type: 'test' }, + }, + }, + }, + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + const results = await Promise.all(artifactPromises); + + expect(logger.error).toHaveBeenCalledWith( + 'Error processing artifact content:', + expect.any(Error), + ); + expect(results[0]).toBeNull(); + }); + + it('should handle multiple artifacts including ui_resources', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const output = { + tool_call_id: 'tool123', + artifact: { + [Tools.ui_resources]: { + data: { + 0: { type: 'chart', data: [] }, + }, + }, + [Tools.web_search]: { + results: ['result1', 'result2'], + }, + }, + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + const results = await Promise.all(artifactPromises); + + // Both ui_resources and web_search should be processed + expect(artifactPromises).toHaveLength(2); + expect(results).toHaveLength(2); + + // Check ui_resources attachment + const uiResourceAttachment = results.find((r) => r?.type === Tools.ui_resources); + expect(uiResourceAttachment).toBeTruthy(); + expect(uiResourceAttachment[Tools.ui_resources]).toEqual({ + 0: { type: 'chart', data: [] }, + }); + + // Check web_search attachment + const webSearchAttachment = results.find((r) => r?.type === Tools.web_search); + expect(webSearchAttachment).toBeTruthy(); + expect(webSearchAttachment[Tools.web_search]).toEqual({ + results: ['result1', 'result2'], + }); + }); + + it('should not process artifacts when output has no artifacts', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const output = { + tool_call_id: 'tool123', + content: 'Some regular content', + // No artifact property + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + + expect(artifactPromises).toHaveLength(0); + expect(res.write).not.toHaveBeenCalled(); + }); + }); + + describe('edge cases', () => { + it('should handle empty ui_resources data object', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const output = { + tool_call_id: 'tool123', + artifact: { + [Tools.ui_resources]: { + data: {}, + }, + }, + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + const results = await Promise.all(artifactPromises); + + expect(results[0]).toEqual({ + type: Tools.ui_resources, + messageId: 'run456', + toolCallId: 'tool123', + conversationId: 'thread789', + [Tools.ui_resources]: {}, + }); + }); + + it('should handle ui_resources with complex nested data', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const complexData = { + 0: { + type: 'form', + fields: [ + { name: 'field1', type: 'text', required: true }, + { name: 'field2', type: 'select', options: ['a', 'b', 'c'] }, + ], + nested: { + deep: { + value: 123, + array: [1, 2, 3], + }, + }, + }, + }; + + const output = { + tool_call_id: 'tool123', + artifact: { + [Tools.ui_resources]: { + data: complexData, + }, + }, + }; + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output }, metadata); + const results = await Promise.all(artifactPromises); + + expect(results[0][Tools.ui_resources]).toEqual(complexData); + }); + + it('should handle when output is undefined', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback({ output: undefined }, metadata); + + expect(artifactPromises).toHaveLength(0); + expect(res.write).not.toHaveBeenCalled(); + }); + + it('should handle when data parameter is undefined', async () => { + const toolEndCallback = createToolEndCallback({ req, res, artifactPromises }); + + const metadata = { + run_id: 'run456', + thread_id: 'thread789', + }; + + await toolEndCallback(undefined, metadata); + + expect(artifactPromises).toHaveLength(0); + expect(res.write).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 6138964ba..d700f0a9c 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -265,6 +265,30 @@ function createToolEndCallback({ req, res, artifactPromises }) { ); } + // TODO: a lot of duplicated code in createToolEndCallback + // we should refactor this to use a helper function in a follow-up PR + if (output.artifact[Tools.ui_resources]) { + artifactPromises.push( + (async () => { + const attachment = { + type: Tools.ui_resources, + messageId: metadata.run_id, + toolCallId: output.tool_call_id, + conversationId: metadata.thread_id, + [Tools.ui_resources]: output.artifact[Tools.ui_resources].data, + }; + if (!res.headersSent) { + return attachment; + } + res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`); + return attachment; + })().catch((error) => { + logger.error('Error processing artifact content:', error); + return null; + }), + ); + } + if (output.artifact[Tools.web_search]) { artifactPromises.push( (async () => { diff --git a/client/src/common/types.ts b/client/src/common/types.ts index b03a80f07..5ea6ad6bb 100644 --- a/client/src/common/types.ts +++ b/client/src/common/types.ts @@ -642,10 +642,3 @@ declare global { google_tag_manager?: unknown; } } - -export type UIResource = { - uri: string; - mimeType: string; - text: string; - [key: string]: unknown; -}; diff --git a/client/src/components/Chat/Messages/Content/ToolCall.tsx b/client/src/components/Chat/Messages/Content/ToolCall.tsx index 7943dd6f1..4af8dd1a2 100644 --- a/client/src/components/Chat/Messages/Content/ToolCall.tsx +++ b/client/src/components/Chat/Messages/Content/ToolCall.tsx @@ -211,6 +211,7 @@ export default function ToolCall({ domain={authDomain || (domain ?? '')} function_name={function_name} pendingAuth={authDomain.length > 0 && !cancelled && progress < 1} + attachments={attachments} /> )} diff --git a/client/src/components/Chat/Messages/Content/ToolCallInfo.tsx b/client/src/components/Chat/Messages/Content/ToolCallInfo.tsx index 1cc472f83..7d9a941d1 100644 --- a/client/src/components/Chat/Messages/Content/ToolCallInfo.tsx +++ b/client/src/components/Chat/Messages/Content/ToolCallInfo.tsx @@ -1,8 +1,9 @@ import React from 'react'; import { useLocalize } from '~/hooks'; +import { Tools } from 'librechat-data-provider'; import { UIResourceRenderer } from '@mcp-ui/client'; import UIResourceCarousel from './UIResourceCarousel'; -import type { UIResource } from '~/common'; +import type { TAttachment, UIResource } from 'librechat-data-provider'; function OptimizedCodeBlock({ text, maxHeight = 320 }: { text: string; maxHeight?: number }) { return ( @@ -27,12 +28,14 @@ export default function ToolCallInfo({ domain, function_name, pendingAuth, + attachments, }: { input: string; function_name: string; output?: string | null; domain?: string; pendingAuth?: boolean; + attachments?: TAttachment[]; }) { const localize = useLocalize(); const formatText = (text: string) => { @@ -54,25 +57,12 @@ export default function ToolCallInfo({ : localize('com_assistants_attempt_info'); } - // Extract ui_resources from the output to display them in the UI - let uiResources: UIResource[] = []; - if (output?.includes('ui_resources')) { - try { - const parsedOutput = JSON.parse(output); - const uiResourcesItem = parsedOutput.find( - (contentItem) => contentItem.metadata?.type === 'ui_resources', - ); - if (uiResourcesItem?.metadata?.data) { - uiResources = uiResourcesItem.metadata.data; - output = JSON.stringify( - parsedOutput.filter((contentItem) => contentItem.metadata?.type !== 'ui_resources'), - ); - } - } catch (error) { - // If JSON parsing fails, keep original output - console.error('Failed to parse output:', error); - } - } + const uiResources: UIResource[] = + attachments + ?.filter((attachment) => attachment.type === Tools.ui_resources) + .flatMap((attachment) => { + return attachment[Tools.ui_resources] as UIResource[]; + }) ?? []; return (
diff --git a/client/src/components/Chat/Messages/Content/UIResourceCarousel.tsx b/client/src/components/Chat/Messages/Content/UIResourceCarousel.tsx index 4fdb7a38a..dfffa0893 100644 --- a/client/src/components/Chat/Messages/Content/UIResourceCarousel.tsx +++ b/client/src/components/Chat/Messages/Content/UIResourceCarousel.tsx @@ -1,6 +1,6 @@ -import { UIResourceRenderer } from '@mcp-ui/client'; -import type { UIResource } from '~/common'; import React, { useState } from 'react'; +import { UIResourceRenderer } from '@mcp-ui/client'; +import type { UIResource } from 'librechat-data-provider'; interface UIResourceCarouselProps { uiResources: UIResource[]; diff --git a/client/src/components/Chat/Messages/Content/__tests__/ToolCall.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/ToolCall.test.tsx new file mode 100644 index 000000000..d1f4f8921 --- /dev/null +++ b/client/src/components/Chat/Messages/Content/__tests__/ToolCall.test.tsx @@ -0,0 +1,382 @@ +import React from 'react'; +import { render, screen, fireEvent } from '@testing-library/react'; +import { RecoilRoot } from 'recoil'; +import { Tools } from 'librechat-data-provider'; +import ToolCall from '../ToolCall'; + +// Mock dependencies +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string, values?: any) => { + const translations: Record = { + com_assistants_function_use: `Used ${values?.[0]}`, + com_assistants_completed_function: `Completed ${values?.[0]}`, + com_assistants_completed_action: `Completed action on ${values?.[0]}`, + com_assistants_running_var: `Running ${values?.[0]}`, + com_assistants_running_action: 'Running action', + com_ui_sign_in_to_domain: `Sign in to ${values?.[0]}`, + com_ui_cancelled: 'Cancelled', + com_ui_requires_auth: 'Requires authentication', + com_assistants_allow_sites_you_trust: 'Only allow sites you trust', + }; + return translations[key] || key; + }, + useProgress: (initialProgress: number) => (initialProgress >= 1 ? 1 : initialProgress), +})); + +jest.mock('~/components/Chat/Messages/Content/MessageContent', () => ({ + __esModule: true, + default: ({ content }: { content: string }) =>
{content}
, +})); + +jest.mock('../ToolCallInfo', () => ({ + __esModule: true, + default: ({ attachments, ...props }: any) => ( +
+ {JSON.stringify(props)} +
+ ), +})); + +jest.mock('../ProgressText', () => ({ + __esModule: true, + default: ({ onClick, inProgressText, finishedText, _error, _hasInput, _isExpanded }: any) => ( +
{finishedText || inProgressText}
+ ), +})); + +jest.mock('../Parts', () => ({ + AttachmentGroup: ({ attachments }: any) => ( +
{JSON.stringify(attachments)}
+ ), +})); + +jest.mock('~/components/ui', () => ({ + Button: ({ children, onClick, ...props }: any) => ( + + ), +})); + +jest.mock('lucide-react', () => ({ + ChevronDown: () => {'ChevronDown'}, + ChevronUp: () => {'ChevronUp'}, + TriangleAlert: () => {'TriangleAlert'}, +})); + +jest.mock('~/utils', () => ({ + logger: { + error: jest.fn(), + }, + cn: (...classes: any[]) => classes.filter(Boolean).join(' '), +})); + +describe('ToolCall', () => { + const mockProps = { + args: '{"test": "input"}', + name: 'testFunction', + output: 'Test output', + initialProgress: 1, + isSubmitting: false, + }; + + const renderWithRecoil = (component: React.ReactElement) => { + return render({component}); + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('attachments prop passing', () => { + it('should pass attachments to ToolCallInfo when provided', () => { + const attachments = [ + { + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: { + '0': { type: 'button', label: 'Click me' }, + }, + }, + ]; + + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + expect(toolCallInfo).toBeInTheDocument(); + + const attachmentsData = toolCallInfo.getAttribute('data-attachments'); + expect(attachmentsData).toBe(JSON.stringify(attachments)); + }); + + it('should pass empty array when no attachments', () => { + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const attachmentsData = toolCallInfo.getAttribute('data-attachments'); + expect(attachmentsData).toBeNull(); // JSON.stringify(undefined) returns undefined, so attribute is not set + }); + + it('should pass multiple attachments of different types', () => { + const attachments = [ + { + type: Tools.ui_resources, + messageId: 'msg1', + toolCallId: 'tool1', + conversationId: 'conv1', + [Tools.ui_resources]: { + '0': { type: 'form', fields: [] }, + }, + }, + { + type: Tools.web_search, + messageId: 'msg2', + toolCallId: 'tool2', + conversationId: 'conv2', + [Tools.web_search]: { + results: ['result1', 'result2'], + }, + }, + ]; + + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const attachmentsData = toolCallInfo.getAttribute('data-attachments'); + expect(JSON.parse(attachmentsData!)).toEqual(attachments); + }); + }); + + describe('attachment group rendering', () => { + it('should render AttachmentGroup when attachments are provided', () => { + const attachments = [ + { + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: { + '0': { type: 'chart', data: [] }, + }, + }, + ]; + + renderWithRecoil(); + + const attachmentGroup = screen.getByTestId('attachment-group'); + expect(attachmentGroup).toBeInTheDocument(); + expect(attachmentGroup.textContent).toBe(JSON.stringify(attachments)); + }); + + it('should not render AttachmentGroup when no attachments', () => { + renderWithRecoil(); + + expect(screen.queryByTestId('attachment-group')).not.toBeInTheDocument(); + }); + + it('should not render AttachmentGroup when attachments is empty array', () => { + renderWithRecoil(); + + expect(screen.queryByTestId('attachment-group')).not.toBeInTheDocument(); + }); + }); + + describe('tool call info visibility', () => { + it('should toggle tool call info when clicking header', () => { + renderWithRecoil(); + + // Initially closed + expect(screen.queryByTestId('tool-call-info')).not.toBeInTheDocument(); + + // Click to open + fireEvent.click(screen.getByText('Completed testFunction')); + expect(screen.getByTestId('tool-call-info')).toBeInTheDocument(); + + // Click to close + fireEvent.click(screen.getByText('Completed testFunction')); + expect(screen.queryByTestId('tool-call-info')).not.toBeInTheDocument(); + }); + + it('should pass all required props to ToolCallInfo', () => { + const attachments = [ + { + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: { + '0': { type: 'button', label: 'Test' }, + }, + }, + ]; + + // Use a name with domain separator (_action_) and domain separator (---) + const propsWithDomain = { + ...mockProps, + name: 'testFunction_action_test---domain---com', // domain will be extracted and --- replaced with dots + attachments, + }; + + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed action on test.domain.com')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const props = JSON.parse(toolCallInfo.textContent!); + + expect(props.input).toBe('{"test": "input"}'); + expect(props.output).toBe('Test output'); + expect(props.function_name).toBe('testFunction'); + // Domain is extracted from name and --- are replaced with dots + expect(props.domain).toBe('test.domain.com'); + expect(props.pendingAuth).toBe(false); + }); + }); + + describe('authentication flow', () => { + it('should show sign-in button when auth URL is provided', () => { + const originalOpen = window.open; + window.open = jest.fn(); + + renderWithRecoil( + , + ); + + const signInButton = screen.getByText('Sign in to auth.example.com'); + expect(signInButton).toBeInTheDocument(); + + fireEvent.click(signInButton); + expect(window.open).toHaveBeenCalledWith( + 'https://auth.example.com', + '_blank', + 'noopener,noreferrer', + ); + + window.open = originalOpen; + }); + + it('should pass pendingAuth as true when auth is pending', () => { + renderWithRecoil( + , + ); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const props = JSON.parse(toolCallInfo.textContent!); + expect(props.pendingAuth).toBe(true); + }); + + it('should not show auth section when cancelled', () => { + renderWithRecoil( + , + ); + + expect(screen.queryByText('Sign in to example.com')).not.toBeInTheDocument(); + }); + + it('should not show auth section when progress is complete', () => { + renderWithRecoil( + , + ); + + expect(screen.queryByText('Sign in to example.com')).not.toBeInTheDocument(); + }); + }); + + describe('edge cases', () => { + it('should handle undefined args', () => { + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const props = JSON.parse(toolCallInfo.textContent!); + expect(props.input).toBe(''); + }); + + it('should handle null output', () => { + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const props = JSON.parse(toolCallInfo.textContent!); + expect(props.output).toBeNull(); + }); + + it('should handle missing domain', () => { + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const props = JSON.parse(toolCallInfo.textContent!); + expect(props.domain).toBe(''); + }); + + it('should handle complex nested attachments', () => { + const complexAttachments = [ + { + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: { + '0': { + type: 'nested', + data: { + deep: { + value: 'test', + array: [1, 2, 3], + object: { key: 'value' }, + }, + }, + }, + }, + }, + ]; + + renderWithRecoil(); + + fireEvent.click(screen.getByText('Completed testFunction')); + + const toolCallInfo = screen.getByTestId('tool-call-info'); + const attachmentsData = toolCallInfo.getAttribute('data-attachments'); + expect(JSON.parse(attachmentsData!)).toEqual(complexAttachments); + + const attachmentGroup = screen.getByTestId('attachment-group'); + expect(JSON.parse(attachmentGroup.textContent!)).toEqual(complexAttachments); + }); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/__tests__/ToolCallInfo.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/ToolCallInfo.test.tsx index d612fb8e9..3e74504fc 100644 --- a/client/src/components/Chat/Messages/Content/__tests__/ToolCallInfo.test.tsx +++ b/client/src/components/Chat/Messages/Content/__tests__/ToolCallInfo.test.tsx @@ -1,8 +1,10 @@ import React from 'react'; -import { render, screen } from '@testing-library/react'; -import ToolCallInfo from '../ToolCallInfo'; +import { Tools } from 'librechat-data-provider'; import { UIResourceRenderer } from '@mcp-ui/client'; -import UIResourceCarousel from '../UIResourceCarousel'; +import { render, screen } from '@testing-library/react'; +import type { TAttachment } from 'librechat-data-provider'; +import UIResourceCarousel from '~/components/Chat/Messages/Content/UIResourceCarousel'; +import ToolCallInfo from '~/components/Chat/Messages/Content/ToolCallInfo'; // Mock the dependencies jest.mock('~/hooks', () => ({ @@ -46,24 +48,25 @@ describe('ToolCallInfo', () => { jest.clearAllMocks(); }); - describe('ui_resources extraction', () => { - it('should extract single ui_resource from output', () => { + describe('ui_resources from attachments', () => { + it('should render single ui_resource from attachments', () => { const uiResource = { type: 'text', data: 'Test resource', }; - const output = JSON.stringify([ - { type: 'text', text: 'Regular output' }, + const attachments: TAttachment[] = [ { - metadata: { - type: 'ui_resources', - data: [uiResource], - }, + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: [uiResource], }, - ]); + ]; - render(); + // Need output for ui_resources to render + render(); // Should render UIResourceRenderer for single resource expect(UIResourceRenderer).toHaveBeenCalledWith( @@ -81,29 +84,33 @@ describe('ToolCallInfo', () => { expect(UIResourceCarousel).not.toHaveBeenCalled(); }); - it('should extract multiple ui_resources from output', () => { - const uiResources = [ - { type: 'text', data: 'Resource 1' }, - { type: 'text', data: 'Resource 2' }, - { type: 'text', data: 'Resource 3' }, + it('should render carousel for multiple ui_resources from attachments', () => { + // To test multiple resources, we can use a single attachment with multiple resources + const attachments: TAttachment[] = [ + { + type: Tools.ui_resources, + messageId: 'msg1', + toolCallId: 'tool1', + conversationId: 'conv1', + [Tools.ui_resources]: [ + { type: 'text', data: 'Resource 1' }, + { type: 'text', data: 'Resource 2' }, + { type: 'text', data: 'Resource 3' }, + ], + }, ]; - const output = JSON.stringify([ - { type: 'text', text: 'Regular output' }, - { - metadata: { - type: 'ui_resources', - data: uiResources, - }, - }, - ]); - - render(); + // Need output for ui_resources to render + render(); // Should render carousel for multiple resources expect(UIResourceCarousel).toHaveBeenCalledWith( expect.objectContaining({ - uiResources, + uiResources: [ + { type: 'text', data: 'Resource 1' }, + { type: 'text', data: 'Resource 2' }, + { type: 'text', data: 'Resource 3' }, + ], }), expect.any(Object), ); @@ -112,34 +119,38 @@ describe('ToolCallInfo', () => { expect(UIResourceRenderer).not.toHaveBeenCalled(); }); - it('should filter out ui_resources from displayed output', () => { - const regularContent = [ - { type: 'text', text: 'Regular output 1' }, - { type: 'text', text: 'Regular output 2' }, + it('should handle attachments with normal output', () => { + const attachments: TAttachment[] = [ + { + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: [{ type: 'text', data: 'UI Resource' }], + }, ]; const output = JSON.stringify([ - ...regularContent, - { - metadata: { - type: 'ui_resources', - data: [{ type: 'text', data: 'UI Resource' }], - }, - }, + { type: 'text', text: 'Regular output 1' }, + { type: 'text', text: 'Regular output 2' }, ]); - const { container } = render(); + const { container } = render( + , + ); - // Check that the displayed output doesn't contain ui_resources + // Check that the output is displayed normally const codeBlocks = container.querySelectorAll('code'); const outputCode = codeBlocks[1]?.textContent; // Second code block is the output expect(outputCode).toContain('Regular output 1'); expect(outputCode).toContain('Regular output 2'); - expect(outputCode).not.toContain('ui_resources'); + + // UI resources should be rendered via attachments + expect(UIResourceRenderer).toHaveBeenCalled(); }); - it('should handle output without ui_resources', () => { + it('should handle no attachments', () => { const output = JSON.stringify([{ type: 'text', text: 'Regular output' }]); render(); @@ -148,66 +159,56 @@ describe('ToolCallInfo', () => { expect(UIResourceCarousel).not.toHaveBeenCalled(); }); - it('should handle malformed ui_resources gracefully', () => { - const output = JSON.stringify([ - { - metadata: 'ui_resources', // metadata should be an object, not a string - text: 'some text content', - }, - ]); + it('should handle empty attachments array', () => { + const attachments: TAttachment[] = []; - // Component should not throw error and should render without UI resources - const { container } = render(); + render(); - // Should render the component without crashing - expect(container).toBeTruthy(); - - // UIResourceCarousel should not be called since the metadata structure is invalid + expect(UIResourceRenderer).not.toHaveBeenCalled(); expect(UIResourceCarousel).not.toHaveBeenCalled(); }); - it('should handle ui_resources as plain text without breaking', () => { - const outputWithTextOnly = - 'This output contains ui_resources as plain text but not as a proper structure'; + it('should handle attachments with non-ui_resources type', () => { + const attachments: TAttachment[] = [ + { + type: Tools.web_search as any, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.web_search]: { + organic: [], + }, + }, + ]; - render(); + render(); - // Should render normally without errors - expect(screen.getByText(`Used ${mockProps.function_name}`)).toBeInTheDocument(); - expect(screen.getByText('Result')).toBeInTheDocument(); - - // The output text should be displayed in a code block - const codeBlocks = screen.getAllByText((content, element) => { - return element?.tagName === 'CODE' && content.includes(outputWithTextOnly); - }); - expect(codeBlocks.length).toBeGreaterThan(0); - - // Should not render UI resources components + // Should not render UI resources components for non-ui_resources attachments expect(UIResourceRenderer).not.toHaveBeenCalled(); expect(UIResourceCarousel).not.toHaveBeenCalled(); }); }); describe('rendering logic', () => { - it('should render UI Resources heading when ui_resources exist', () => { - const output = JSON.stringify([ + it('should render UI Resources heading when ui_resources exist in attachments', () => { + const attachments: TAttachment[] = [ { - metadata: { - type: 'ui_resources', - data: [{ type: 'text', data: 'Test' }], - }, + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: [{ type: 'text', data: 'Test' }], }, - ]); + ]; - render(); + // Need output for ui_resources section to render + render(); expect(screen.getByText('UI Resources')).toBeInTheDocument(); }); - it('should not render UI Resources heading when no ui_resources', () => { - const output = JSON.stringify([{ type: 'text', text: 'Regular output' }]); - - render(); + it('should not render UI Resources heading when no ui_resources in attachments', () => { + render(); expect(screen.queryByText('UI Resources')).not.toBeInTheDocument(); }); @@ -218,16 +219,18 @@ describe('ToolCallInfo', () => { data: { fields: [{ name: 'test', type: 'text' }] }, }; - const output = JSON.stringify([ + const attachments: TAttachment[] = [ { - metadata: { - type: 'ui_resources', - data: [uiResource], - }, + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: [uiResource], }, - ]); + ]; - render(); + // Need output for ui_resources to render + render(); expect(UIResourceRenderer).toHaveBeenCalledWith( expect.objectContaining({ @@ -244,16 +247,18 @@ describe('ToolCallInfo', () => { it('should console.log when UIAction is triggered', async () => { const consoleSpy = jest.spyOn(console, 'log').mockImplementation(); - const output = JSON.stringify([ + const attachments: TAttachment[] = [ { - metadata: { - type: 'ui_resources', - data: [{ type: 'text', data: 'Test' }], - }, + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: [{ type: 'text', data: 'Test' }], }, - ]); + ]; - render(); + // Need output for ui_resources to render + render(); const mockUIResourceRenderer = UIResourceRenderer as jest.MockedFunction< typeof UIResourceRenderer @@ -270,4 +275,55 @@ describe('ToolCallInfo', () => { consoleSpy.mockRestore(); }); }); + + describe('backward compatibility', () => { + it('should handle output with ui_resources for backward compatibility', () => { + const output = JSON.stringify([ + { type: 'text', text: 'Regular output' }, + { + metadata: { + type: 'ui_resources', + data: [{ type: 'text', data: 'UI Resource' }], + }, + }, + ]); + + render(); + + // Since we now use attachments, ui_resources in output should be ignored + expect(UIResourceRenderer).not.toHaveBeenCalled(); + expect(UIResourceCarousel).not.toHaveBeenCalled(); + }); + + it('should prioritize attachments over output ui_resources', () => { + const attachments: TAttachment[] = [ + { + type: Tools.ui_resources, + messageId: 'msg123', + toolCallId: 'tool456', + conversationId: 'conv789', + [Tools.ui_resources]: [{ type: 'attachment', data: 'From attachments' }], + }, + ]; + + const output = JSON.stringify([ + { + metadata: { + type: 'ui_resources', + data: [{ type: 'output', data: 'From output' }], + }, + }, + ]); + + render(); + + // Should use attachments, not output + expect(UIResourceRenderer).toHaveBeenCalledWith( + expect.objectContaining({ + resource: { type: 'attachment', data: 'From attachments' }, + }), + expect.any(Object), + ); + }); + }); }); diff --git a/client/src/components/Chat/Messages/Content/__tests__/UIResourceCarousel.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/UIResourceCarousel.test.tsx index faeddd1e0..2878c6444 100644 --- a/client/src/components/Chat/Messages/Content/__tests__/UIResourceCarousel.test.tsx +++ b/client/src/components/Chat/Messages/Content/__tests__/UIResourceCarousel.test.tsx @@ -1,8 +1,8 @@ import React from 'react'; -import { render, screen, fireEvent, waitFor } from '@testing-library/react'; import '@testing-library/jest-dom'; -import UIResourceCarousel from '../UIResourceCarousel'; -import type { UIResource } from '~/common'; +import { render, screen, fireEvent, waitFor } from '@testing-library/react'; +import type { UIResource } from 'librechat-data-provider'; +import UIResourceCarousel from '~/components/Chat/Messages/Content/UIResourceCarousel'; // Mock the UIResourceRenderer component jest.mock('@mcp-ui/client', () => ({ diff --git a/packages/api/src/mcp/__tests__/parsers.test.ts b/packages/api/src/mcp/__tests__/parsers.test.ts index de3d4cd70..a26373aef 100644 --- a/packages/api/src/mcp/__tests__/parsers.test.ts +++ b/packages/api/src/mcp/__tests__/parsers.test.ts @@ -161,7 +161,7 @@ describe('formatToolContent', () => { }); describe('resource handling', () => { - it('should handle UI resources', () => { + it('should handle UI resources in artifacts', () => { const result: t.MCPToolCallResponse = { content: [ { @@ -181,22 +181,27 @@ describe('formatToolContent', () => { expect(content).toEqual([ { type: 'text', - text: '', - metadata: { - type: 'ui_resources', - data: [ - { - uri: 'ui://carousel', - mimeType: 'application/json', - text: '{"items": []}', - name: 'carousel', - description: 'A carousel component', - }, - ], - }, + text: + 'Resource Text: {"items": []}\n' + + 'Resource URI: ui://carousel\n' + + 'Resource: carousel\n' + + 'Resource Description: A carousel component\n' + + 'Resource MIME Type: application/json', }, ]); - expect(artifacts).toBeUndefined(); + expect(artifacts).toEqual({ + ui_resources: { + data: [ + { + uri: 'ui://carousel', + mimeType: 'application/json', + text: '{"items": []}', + name: 'carousel', + description: 'A carousel component', + }, + ], + }, + }); }); it('should handle regular resources', () => { @@ -281,24 +286,75 @@ describe('formatToolContent', () => { expect(content).toEqual([ { type: 'text', - text: 'Some text\n\n' + 'Resource URI: file://data.csv\n' + 'Resource: Data file', + text: + 'Some text\n\n' + + 'Resource Text: {"label": "Click me"}\n' + + 'Resource URI: ui://button\n' + + 'Resource MIME Type: application/json\n\n' + + 'Resource URI: file://data.csv\n' + + 'Resource: Data file', + }, + ]); + expect(artifacts).toEqual({ + ui_resources: { + data: [ + { + uri: 'ui://button', + mimeType: 'application/json', + text: '{"label": "Click me"}', + }, + ], + }, + }); + }); + + it('should handle both images and UI resources in artifacts', () => { + const result: t.MCPToolCallResponse = { + content: [ + { type: 'text', text: 'Content with multimedia' }, + { type: 'image', data: 'base64imagedata', mimeType: 'image/png' }, + { + type: 'resource', + resource: { + uri: 'ui://graph', + mimeType: 'application/json', + text: '{"type": "line"}', + }, + }, + ], + }; + + const [content, artifacts] = formatToolContent(result, 'openai'); + expect(content).toEqual([ + { + type: 'text', + text: 'Content with multimedia', }, { type: 'text', - text: '', - metadata: { - type: 'ui_resources', - data: [ - { - uri: 'ui://button', - mimeType: 'application/json', - text: '{"label": "Click me"}', - }, - ], - }, + text: + 'Resource Text: {"type": "line"}\n' + + 'Resource URI: ui://graph\n' + + 'Resource MIME Type: application/json', }, ]); - expect(artifacts).toBeUndefined(); + expect(artifacts).toEqual({ + content: [ + { + type: 'image_url', + image_url: { url: 'data:image/png;base64,base64imagedata' }, + }, + ], + ui_resources: { + data: [ + { + uri: 'ui://graph', + mimeType: 'application/json', + text: '{"type": "line"}', + }, + ], + }, + }); }); }); @@ -358,25 +414,14 @@ describe('formatToolContent', () => { type: 'text', text: 'Middle section\n\n' + + 'Resource Text: {"type": "bar"}\n' + + 'Resource URI: ui://chart\n' + + 'Resource MIME Type: application/json\n\n' + 'Resource URI: https://api.example.com/data\n' + 'Resource: API Data\n' + 'Resource Description: External data source', }, { type: 'text', text: 'Conclusion' }, - { - type: 'text', - text: '', - metadata: { - type: 'ui_resources', - data: [ - { - uri: 'ui://chart', - mimeType: 'application/json', - text: '{"type": "bar"}', - }, - ], - }, - }, ]); expect(artifacts).toEqual({ content: [ @@ -389,6 +434,15 @@ describe('formatToolContent', () => { image_url: { url: 'https://example.com/image2.jpg' }, }, ], + ui_resources: { + data: [ + { + uri: 'ui://chart', + mimeType: 'application/json', + text: '{"type": "bar"}', + }, + ], + }, }); }); diff --git a/packages/api/src/mcp/parsers.ts b/packages/api/src/mcp/parsers.ts index 77af29bee..97b1efc4c 100644 --- a/packages/api/src/mcp/parsers.ts +++ b/packages/api/src/mcp/parsers.ts @@ -1,4 +1,7 @@ +import { Tools } from 'librechat-data-provider'; +import type { UIResource } from 'librechat-data-provider'; import type * as t from './types'; + const RECOGNIZED_PROVIDERS = new Set([ 'google', 'anthropic', @@ -111,7 +114,7 @@ export function formatToolContent( const formattedContent: t.FormattedContent[] = []; const imageUrls: t.FormattedContent[] = []; let currentTextBlock = ''; - const uiResources: t.UIResource[] = []; + const uiResources: UIResource[] = []; type ContentHandler = undefined | ((item: t.ToolContentPart) => void); @@ -144,8 +147,7 @@ export function formatToolContent( resource: (item) => { if (item.resource.uri.startsWith('ui://')) { - uiResources.push(item.resource as t.UIResource); - return; + uiResources.push(item.resource as UIResource); } const resourceText = []; @@ -182,18 +184,14 @@ export function formatToolContent( formattedContent.push({ type: 'text', text: currentTextBlock }); } - if (uiResources.length) { - formattedContent.push({ - type: 'text', - metadata: { - type: 'ui_resources', - data: uiResources, - }, - text: '', - }); + let artifacts: t.Artifacts = undefined; + if (imageUrls.length || uiResources.length) { + artifacts = { + ...(imageUrls.length && { content: imageUrls }), + ...(uiResources.length && { [Tools.ui_resources]: { data: uiResources } }), + }; } - const artifacts = imageUrls.length ? { content: imageUrls } : undefined; if (CONTENT_ARRAY_PROVIDERS.has(provider)) { return [formattedContent, artifacts]; } diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index 95086a367..b774a2fe1 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -6,8 +6,9 @@ import { StdioOptionsSchema, WebSocketOptionsSchema, StreamableHTTPOptionsSchema, + Tools, } from 'librechat-data-provider'; -import type { TPlugin, TUser } from 'librechat-data-provider'; +import type { UIResource, TPlugin, TUser } from 'librechat-data-provider'; import type * as t from '@modelcontextprotocol/sdk/types.js'; import type { TokenMethods } from '@librechat/data-schemas'; import type { FlowStateManager } from '~/flow/manager'; @@ -86,7 +87,7 @@ export type FormattedContent = metadata?: { type: string; data: UIResource[]; - } + }; text?: string; } | { @@ -111,18 +112,36 @@ export type FormattedContent = }; }; -export type FormattedContentResult = [ - string | FormattedContent[], - undefined | { content: FormattedContent[] }, -]; - -export type UIResource = { - uri: string; - mimeType: string; - text: string; +export type FileSearchSource = { + fileId: string; + relevance: number; + fileName?: string; + metadata?: { + storageType?: string; + [key: string]: unknown; + }; [key: string]: unknown; }; +export type Artifacts = + | { + content?: FormattedContent[]; + [Tools.ui_resources]?: { + data: UIResource[]; + }; + [Tools.file_search]?: { + sources: FileSearchSource[]; + fileCitations?: boolean; + }; + [Tools.web_search]?: import('librechat-data-provider').SearchResultData; + files?: Array<{ id: string; name: string }>; + session_id?: string; + file_ids?: string[]; + } + | undefined; + +export type FormattedContentResult = [string | FormattedContent[], undefined | Artifacts]; + export type ImageFormatter = (item: ImageContent) => FormattedContent; export type FormattedToolResponse = [ diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index dc4baebc0..8ae503ef2 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -552,20 +552,33 @@ export type MemoryArtifact = { type: 'update' | 'delete' | 'error'; }; +export type UIResource = { + type?: string; + data?: unknown; + uri?: string; + mimeType?: string; + text?: string; + [key: string]: unknown; +}; + export type TAttachmentMetadata = { type?: Tools; messageId: string; toolCallId: string; + [Tools.memory]?: MemoryArtifact; + [Tools.ui_resources]?: UIResource[]; [Tools.web_search]?: SearchResultData; [Tools.file_search]?: SearchResultData; - [Tools.memory]?: MemoryArtifact; }; export type TAttachment = | (TFile & TAttachmentMetadata) | (Pick & { expiresAt: number; - } & TAttachmentMetadata); + } & TAttachmentMetadata) + | (Partial> & + Pick & + TAttachmentMetadata); export type TMessage = z.input & { children?: TMessage[]; diff --git a/packages/data-provider/src/types/assistants.ts b/packages/data-provider/src/types/assistants.ts index df9009728..03e287cbc 100644 --- a/packages/data-provider/src/types/assistants.ts +++ b/packages/data-provider/src/types/assistants.ts @@ -23,6 +23,7 @@ export enum Tools { retrieval = 'retrieval', function = 'function', memory = 'memory', + ui_resources = 'ui_resources', } export enum EToolResources {