@librechat/agents 3.1.37 → 3.1.39

package/src/tools/__tests__/ToolNode.session.test.ts

@@ -0,0 +1,465 @@
+import { z } from 'zod';
+import { tool } from '@langchain/core/tools';
+import { AIMessage } from '@langchain/core/messages';
+import { describe, it, expect } from '@jest/globals';
+import type { StructuredToolInterface } from '@langchain/core/tools';
+import type * as t from '@/types';
+import { ToolNode } from '../ToolNode';
+import { Constants } from '@/common';
+
+/**
+ * Creates a mock execute_code tool that captures the toolCall config it receives.
+ * Returns a content_and_artifact response with configurable session/files.
+ */
+function createMockCodeTool(options: {
+  capturedConfigs: Record<string, unknown>[];
+  artifact?: t.CodeExecutionArtifact;
+}): StructuredToolInterface {
+  const { capturedConfigs, artifact } = options;
+  const defaultArtifact: t.CodeExecutionArtifact = {
+    session_id: 'new-session-123',
+    files: [],
+  };
+
+  return tool(
+    async (_input, config) => {
+      capturedConfigs.push({ ...(config.toolCall ?? {}) });
+      return ['stdout:\nhello world\n', artifact ?? defaultArtifact];
+    },
+    {
+      name: Constants.EXECUTE_CODE,
+      description: 'Execute code in a sandbox',
+      schema: z.object({ lang: z.string(), code: z.string() }),
+      responseFormat: Constants.CONTENT_AND_ARTIFACT,
+    }
+  ) as unknown as StructuredToolInterface;
+}
+
+function createAIMessageWithCodeCall(callId: string): AIMessage {
+  return new AIMessage({
+    content: '',
+    tool_calls: [
+      {
+        id: callId,
+        name: Constants.EXECUTE_CODE,
+        args: { lang: 'python', code: 'print("hello")' },
+      },
+    ],
+  });
+}
+
+describe('ToolNode code execution session management', () => {
+  describe('session injection via runTool (direct execution)', () => {
+    it('injects session_id and _injected_files when session has files', async () => {
+      const capturedConfigs: Record<string, unknown>[] = [];
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'prev-session-abc',
+        files: [
+          { id: 'file1', name: 'data.csv', session_id: 'prev-session-abc' },
+          { id: 'file2', name: 'chart.png', session_id: 'prev-session-abc' },
+        ],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs });
+      const toolNode = new ToolNode({ tools: [mockTool], sessions });
+
+      const aiMsg = createAIMessageWithCodeCall('call_1');
+      await toolNode.invoke({ messages: [aiMsg] });
+
+      expect(capturedConfigs).toHaveLength(1);
+      expect(capturedConfigs[0].session_id).toBe('prev-session-abc');
+      expect(capturedConfigs[0]._injected_files).toEqual([
+        { session_id: 'prev-session-abc', id: 'file1', name: 'data.csv' },
+        { session_id: 'prev-session-abc', id: 'file2', name: 'chart.png' },
+      ]);
+    });
+
+    it('injects session_id even when session has no tracked files', async () => {
+      const capturedConfigs: Record<string, unknown>[] = [];
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'prev-session-no-files',
+        files: [],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs });
+      const toolNode = new ToolNode({ tools: [mockTool], sessions });
+
+      const aiMsg = createAIMessageWithCodeCall('call_2');
+      await toolNode.invoke({ messages: [aiMsg] });
+
+      expect(capturedConfigs).toHaveLength(1);
+      expect(capturedConfigs[0].session_id).toBe('prev-session-no-files');
+      expect(capturedConfigs[0]._injected_files).toBeUndefined();
+    });
+
+    it('does not inject session context when no session exists', async () => {
+      const capturedConfigs: Record<string, unknown>[] = [];
+      const sessions: t.ToolSessionMap = new Map();
+
+      const mockTool = createMockCodeTool({ capturedConfigs });
+      const toolNode = new ToolNode({ tools: [mockTool], sessions });
+
+      const aiMsg = createAIMessageWithCodeCall('call_3');
+      await toolNode.invoke({ messages: [aiMsg] });
+
+      expect(capturedConfigs).toHaveLength(1);
+      expect(capturedConfigs[0].session_id).toBeUndefined();
+      expect(capturedConfigs[0]._injected_files).toBeUndefined();
+    });
+
+    it('preserves per-file session_id for multi-session files', async () => {
+      const capturedConfigs: Record<string, unknown>[] = [];
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'session-B',
+        files: [
+          { id: 'f1', name: 'old.csv', session_id: 'session-A' },
+          { id: 'f2', name: 'new.png', session_id: 'session-B' },
+        ],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs });
+      const toolNode = new ToolNode({ tools: [mockTool], sessions });
+
+      const aiMsg = createAIMessageWithCodeCall('call_4');
+      await toolNode.invoke({ messages: [aiMsg] });
+
+      const files = capturedConfigs[0]._injected_files as t.CodeEnvFile[];
+      expect(files[0].session_id).toBe('session-A');
+      expect(files[1].session_id).toBe('session-B');
+    });
+  });
+
+  describe('getCodeSessionContext (via dispatchToolEvents request building)', () => {
+    it('builds session context with files for event-driven requests', () => {
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'evt-session',
+        files: [{ id: 'ef1', name: 'out.parquet', session_id: 'evt-session' }],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const context = (
+        toolNode as unknown as { getCodeSessionContext: () => unknown }
+      ).getCodeSessionContext();
+
+      expect(context).toEqual({
+        session_id: 'evt-session',
+        files: [{ session_id: 'evt-session', id: 'ef1', name: 'out.parquet' }],
+      });
+    });
+
+    it('builds session context without files when session has no tracked files', () => {
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'evt-session-empty',
+        files: [],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const context = (
+        toolNode as unknown as { getCodeSessionContext: () => unknown }
+      ).getCodeSessionContext();
+
+      expect(context).toEqual({ session_id: 'evt-session-empty' });
+    });
+
+    it('returns undefined when no session exists', () => {
+      const sessions: t.ToolSessionMap = new Map();
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const context = (
+        toolNode as unknown as { getCodeSessionContext: () => unknown }
+      ).getCodeSessionContext();
+
+      expect(context).toBeUndefined();
+    });
+  });
+
+  describe('storeCodeSessionFromResults (session storage from artifacts)', () => {
+    it('stores session with files from code execution results', () => {
+      const sessions: t.ToolSessionMap = new Map();
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const storeMethod = (
+        toolNode as unknown as {
+          storeCodeSessionFromResults: (
+            results: t.ToolExecuteResult[],
+            requests: t.ToolCallRequest[]
+          ) => void;
+        }
+      ).storeCodeSessionFromResults.bind(toolNode);
+
+      storeMethod(
+        [
+          {
+            toolCallId: 'tc1',
+            content: 'output',
+            artifact: {
+              session_id: 'new-sess',
+              files: [{ id: 'f1', name: 'result.csv' }],
+            },
+            status: 'success',
+          },
+        ],
+        [{ id: 'tc1', name: Constants.EXECUTE_CODE, args: {} }]
+      );
+
+      const stored = sessions.get(
+        Constants.EXECUTE_CODE
+      ) as t.CodeSessionContext;
+      expect(stored).toBeDefined();
+      expect(stored.session_id).toBe('new-sess');
+      expect(stored.files).toHaveLength(1);
+      expect(stored.files![0]).toEqual(
+        expect.objectContaining({
+          id: 'f1',
+          name: 'result.csv',
+          session_id: 'new-sess',
+        })
+      );
+    });
+
+    it('stores session_id even when Code API returns no files', () => {
+      const sessions: t.ToolSessionMap = new Map();
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const storeMethod = (
+        toolNode as unknown as {
+          storeCodeSessionFromResults: (
+            results: t.ToolExecuteResult[],
+            requests: t.ToolCallRequest[]
+          ) => void;
+        }
+      ).storeCodeSessionFromResults.bind(toolNode);
+
+      storeMethod(
+        [
+          {
+            toolCallId: 'tc2',
+            content: 'stdout:\nSaved parquet\n',
+            artifact: { session_id: 'parquet-session', files: [] },
+            status: 'success',
+          },
+        ],
+        [{ id: 'tc2', name: Constants.EXECUTE_CODE, args: {} }]
+      );
+
+      const stored = sessions.get(
+        Constants.EXECUTE_CODE
+      ) as t.CodeSessionContext;
+      expect(stored).toBeDefined();
+      expect(stored.session_id).toBe('parquet-session');
+      expect(stored.files).toEqual([]);
+    });
+
+    it('merges new files with existing session, replacing same-name files', () => {
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'old-sess',
+        files: [
+          { id: 'f1', name: 'data.csv', session_id: 'old-sess' },
+          { id: 'f2', name: 'chart.png', session_id: 'old-sess' },
+        ],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const storeMethod = (
+        toolNode as unknown as {
+          storeCodeSessionFromResults: (
+            results: t.ToolExecuteResult[],
+            requests: t.ToolCallRequest[]
+          ) => void;
+        }
+      ).storeCodeSessionFromResults.bind(toolNode);
+
+      storeMethod(
+        [
+          {
+            toolCallId: 'tc3',
+            content: 'output',
+            artifact: {
+              session_id: 'new-sess',
+              files: [{ id: 'f3', name: 'chart.png' }],
+            },
+            status: 'success',
+          },
+        ],
+        [{ id: 'tc3', name: Constants.EXECUTE_CODE, args: {} }]
+      );
+
+      const stored = sessions.get(
+        Constants.EXECUTE_CODE
+      ) as t.CodeSessionContext;
+      expect(stored.session_id).toBe('new-sess');
+      expect(stored.files).toHaveLength(2);
+
+      const csvFile = stored.files!.find((f) => f.name === 'data.csv');
+      expect(csvFile!.session_id).toBe('old-sess');
+
+      const chartFile = stored.files!.find((f) => f.name === 'chart.png');
+      expect(chartFile!.id).toBe('f3');
+      expect(chartFile!.session_id).toBe('new-sess');
+    });
+
+    it('preserves existing files when new execution has no files', () => {
+      const sessions: t.ToolSessionMap = new Map();
+      sessions.set(Constants.EXECUTE_CODE, {
+        session_id: 'old-sess',
+        files: [{ id: 'f1', name: 'data.csv', session_id: 'old-sess' }],
+        lastUpdated: Date.now(),
+      } satisfies t.CodeSessionContext);
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const storeMethod = (
+        toolNode as unknown as {
+          storeCodeSessionFromResults: (
+            results: t.ToolExecuteResult[],
+            requests: t.ToolCallRequest[]
+          ) => void;
+        }
+      ).storeCodeSessionFromResults.bind(toolNode);
+
+      storeMethod(
+        [
+          {
+            toolCallId: 'tc4',
+            content: 'stdout:\nno files generated\n',
+            artifact: { session_id: 'new-sess', files: [] },
+            status: 'success',
+          },
+        ],
+        [{ id: 'tc4', name: Constants.EXECUTE_CODE, args: {} }]
+      );
+
+      const stored = sessions.get(
+        Constants.EXECUTE_CODE
+      ) as t.CodeSessionContext;
+      expect(stored.session_id).toBe('new-sess');
+      expect(stored.files).toHaveLength(1);
+      expect(stored.files![0].name).toBe('data.csv');
+    });
+
+    it('ignores non-code-execution tool results', () => {
+      const sessions: t.ToolSessionMap = new Map();
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const storeMethod = (
+        toolNode as unknown as {
+          storeCodeSessionFromResults: (
+            results: t.ToolExecuteResult[],
+            requests: t.ToolCallRequest[]
+          ) => void;
+        }
+      ).storeCodeSessionFromResults.bind(toolNode);
+
+      storeMethod(
+        [
+          {
+            toolCallId: 'tc5',
+            content: 'search results',
+            artifact: { session_id: 'should-not-store' },
+            status: 'success',
+          },
+        ],
+        [{ id: 'tc5', name: 'web_search', args: {} }]
+      );
+
+      expect(sessions.has(Constants.EXECUTE_CODE)).toBe(false);
+    });
+
+    it('ignores error results', () => {
+      const sessions: t.ToolSessionMap = new Map();
+
+      const mockTool = createMockCodeTool({ capturedConfigs: [] });
+      const toolNode = new ToolNode({
+        tools: [mockTool],
+        sessions,
+        eventDrivenMode: true,
+      });
+
+      const storeMethod = (
+        toolNode as unknown as {
+          storeCodeSessionFromResults: (
+            results: t.ToolExecuteResult[],
+            requests: t.ToolCallRequest[]
+          ) => void;
+        }
+      ).storeCodeSessionFromResults.bind(toolNode);
+
+      storeMethod(
+        [
+          {
+            toolCallId: 'tc6',
+            content: '',
+            artifact: {
+              session_id: 'error-session',
+              files: [{ id: 'f1', name: 'x' }],
+            },
+            status: 'error',
+            errorMessage: 'execution failed',
+          },
+        ],
+        [{ id: 'tc6', name: Constants.EXECUTE_CODE, args: {} }]
+      );
+
+      expect(sessions.has(Constants.EXECUTE_CODE)).toBe(false);
+    });
+  });
+});