@librechat/agents 3.1.36 → 3.1.38

package/src/tools/ToolNode.ts

@@ -161,6 +161,9 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
        * Each file uses its own session_id (supporting multi-session file tracking).
        * Both session_id and _injected_files are injected directly to invokeParams
        * (not inside args) so they bypass Zod schema validation and reach config.toolCall.
+       *
+       * session_id is always injected when available (even without tracked files)
+       * so the CodeExecutor can fall back to the /files endpoint for session continuity.
        */
       if (
         call.name === Constants.EXECUTE_CODE ||
@@ -169,23 +172,20 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         const codeSession = this.sessions?.get(Constants.EXECUTE_CODE) as
           | t.CodeSessionContext
           | undefined;
-        if (codeSession?.files != null && codeSession.files.length > 0) {
-          /**
-           * Convert tracked files to CodeEnvFile format for the API.
-           * Each file uses its own session_id (set when file was created).
-           * This supports files from multiple parallel/sequential executions.
-           */
-          const fileRefs: t.CodeEnvFile[] = codeSession.files.map((file) => ({
-            session_id: file.session_id ?? codeSession.session_id,
-            id: file.id,
-            name: file.name,
-          }));
-          /** Inject latest session_id and files - bypasses Zod, reaches config.toolCall */
+        if (codeSession?.session_id != null && codeSession.session_id !== '') {
           invokeParams = {
             ...invokeParams,
             session_id: codeSession.session_id,
-            _injected_files: fileRefs,
           };
+
+          if (codeSession.files != null && codeSession.files.length > 0) {
+            const fileRefs: t.CodeEnvFile[] = codeSession.files.map((file) => ({
+              session_id: file.session_id ?? codeSession.session_id,
+              id: file.id,
+              name: file.name,
+            }));
+            invokeParams._injected_files = fileRefs;
+          }
         }
       }
 
@@ -256,6 +256,100 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     }
   }
 
+  /**
+   * Builds code session context for injection into event-driven tool calls.
+   * Mirrors the session injection logic in runTool() for direct execution.
+   */
+  private getCodeSessionContext(): t.ToolCallRequest['codeSessionContext'] {
+    if (!this.sessions) {
+      return undefined;
+    }
+
+    const codeSession = this.sessions.get(Constants.EXECUTE_CODE) as
+      | t.CodeSessionContext
+      | undefined;
+    if (!codeSession) {
+      return undefined;
+    }
+
+    const context: NonNullable<t.ToolCallRequest['codeSessionContext']> = {
+      session_id: codeSession.session_id,
+    };
+
+    if (codeSession.files && codeSession.files.length > 0) {
+      context.files = codeSession.files.map((file) => ({
+        session_id: file.session_id ?? codeSession.session_id,
+        id: file.id,
+        name: file.name,
+      }));
+    }
+
+    return context;
+  }
+
+  /**
+   * Extracts code execution session context from tool results and stores in Graph.sessions.
+   * Mirrors the session storage logic in Graph.handleToolCallCompleted() for direct execution.
+   */
+  private storeCodeSessionFromResults(
+    results: t.ToolExecuteResult[],
+    requests: t.ToolCallRequest[]
+  ): void {
+    if (!this.sessions) {
+      return;
+    }
+
+    for (let i = 0; i < results.length; i++) {
+      const result = results[i];
+      if (result.status !== 'success' || result.artifact == null) {
+        continue;
+      }
+
+      const request = requests.find((r) => r.id === result.toolCallId);
+      if (
+        request?.name !== Constants.EXECUTE_CODE &&
+        request?.name !== Constants.PROGRAMMATIC_TOOL_CALLING
+      ) {
+        continue;
+      }
+
+      const artifact = result.artifact as t.CodeExecutionArtifact | undefined;
+      if (artifact?.session_id == null || artifact.session_id === '') {
+        continue;
+      }
+
+      const newFiles = artifact.files ?? [];
+      const existingSession = this.sessions.get(Constants.EXECUTE_CODE) as
+        | t.CodeSessionContext
+        | undefined;
+      const existingFiles = existingSession?.files ?? [];
+
+      if (newFiles.length > 0) {
+        const filesWithSession: t.FileRefs = newFiles.map((file) => ({
+          ...file,
+          session_id: artifact.session_id,
+        }));
+
+        const newFileNames = new Set(filesWithSession.map((f) => f.name));
+        const filteredExisting = existingFiles.filter(
+          (f) => !newFileNames.has(f.name)
+        );
+
+        this.sessions.set(Constants.EXECUTE_CODE, {
+          session_id: artifact.session_id,
+          files: [...filteredExisting, ...filesWithSession],
+          lastUpdated: Date.now(),
+        });
+      } else {
+        this.sessions.set(Constants.EXECUTE_CODE, {
+          session_id: artifact.session_id,
+          files: existingFiles,
+          lastUpdated: Date.now(),
+        });
+      }
+    }
+  }
+
   /**
    * Dispatches tool calls to the host via ON_TOOL_EXECUTE event and returns raw ToolMessages.
    * Core logic for event-driven execution, separated from output shaping.
@@ -267,13 +361,23 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     const requests: t.ToolCallRequest[] = toolCalls.map((call) => {
       const turn = this.toolUsageCount.get(call.name) ?? 0;
       this.toolUsageCount.set(call.name, turn + 1);
-      return {
+
+      const request: t.ToolCallRequest = {
         id: call.id!,
         name: call.name,
         args: call.args as Record<string, unknown>,
         stepId: this.toolCallStepIds?.get(call.id!),
         turn,
       };
+
+      if (
+        call.name === Constants.EXECUTE_CODE ||
+        call.name === Constants.PROGRAMMATIC_TOOL_CALLING
+      ) {
+        request.codeSessionContext = this.getCodeSessionContext();
+      }
+
+      return request;
     });
 
     const results = await new Promise<t.ToolExecuteResult[]>(
@@ -294,6 +398,8 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       }
     );
 
+    this.storeCodeSessionFromResults(results, requests);
+
     return results.map((result) => {
       const request = requests.find((r) => r.id === result.toolCallId);
       const toolName = request?.name ?? 'unknown';

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