✂️ refactor: use artifacts and callbacks to pass UI resources

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();
+    });
+  });
+});

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 () => {

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}
               />
             )}
           </div>