🗂️ feat: Better Persistence for Code Execution Files Between Sessions (#11362)

* refactor: process code output files for re-use (WIP)

* feat: file attachment handling with additional metadata for downloads

* refactor: Update directory path logic for local file saving based on basePath

* refactor: file attachment handling to support TFile type and improve data merging logic

* feat: thread filtering of code-generated files

- Introduced parentMessageId parameter in addedConvo and initialize functions to enhance thread management.
- Updated related methods to utilize parentMessageId for retrieving messages and filtering code-generated files by conversation threads.
- Enhanced type definitions to include parentMessageId in relevant interfaces for better clarity and usage.

* chore: imports/params ordering

* feat: update file model to use messageId for filtering and processing

- Changed references from 'message' to 'messageId' in file-related methods for consistency.
- Added messageId field to the file schema and updated related types.
- Enhanced file processing logic to accommodate the new messageId structure.

* feat: enhance file retrieval methods to support user-uploaded execute_code files

- Added a new method `getUserCodeFiles` to retrieve user-uploaded execute_code files, excluding code-generated files.
- Updated existing file retrieval methods to improve filtering logic and handle edge cases.
- Enhanced thread data extraction to collect both message IDs and file IDs efficiently.
- Integrated `getUserCodeFiles` into relevant endpoints for better file management in conversations.

* chore: update @librechat/agents package version to 3.0.78 in package-lock.json and related package.json files

* refactor: file processing and retrieval logic

- Added a fallback mechanism for download URLs when files exceed size limits or cannot be processed locally.
- Implemented a deduplication strategy for code-generated files based on conversationId and filename to optimize storage.
- Updated file retrieval methods to ensure proper filtering by messageIds, preventing orphaned files from being included.
- Introduced comprehensive tests for new thread data extraction functionality, covering edge cases and performance considerations.

* fix: improve file retrieval tests and handling of optional properties

- Updated tests to safely access optional properties using non-null assertions.
- Modified test descriptions for clarity regarding the exclusion of execute_code files.
- Ensured that the retrieval logic correctly reflects the expected outcomes for file queries.

* test: add comprehensive unit tests for processCodeOutput functionality

- Introduced a new test suite for the processCodeOutput function, covering various scenarios including file retrieval, creation, and processing for both image and non-image files.
- Implemented mocks for dependencies such as axios, logger, and file models to isolate tests and ensure reliable outcomes.
- Validated behavior for existing files, new file creation, and error handling, including size limits and fallback mechanisms.
- Enhanced test coverage for metadata handling and usage increment logic, ensuring robust verification of file processing outcomes.

* test: enhance file size limit enforcement in processCodeOutput tests

- Introduced a configurable file size limit for tests to improve flexibility and coverage.
- Mocked the `librechat-data-provider` to allow dynamic adjustment of file size limits during tests.
- Updated the file size limit enforcement test to validate behavior when files exceed specified limits, ensuring proper fallback to download URLs.
- Reset file size limit after tests to maintain isolation for subsequent test cases.

packages/api/src/agents/initialize.ts

@@ -1,5 +1,6 @@
 import { Providers } from '@librechat/agents';
 import {
+  Constants,
   ErrorTypes,
   EModelEndpoint,
   EToolResources,
@@ -20,7 +21,12 @@ import type { GenericTool, LCToolRegistry, ToolMap } from '@librechat/agents';
 import type { Response as ServerResponse } from 'express';
 import type { IMongoFile } from '@librechat/data-schemas';
 import type { InitializeResultBase, ServerRequest, EndpointDbMethods } from '~/types';
-import { getModelMaxTokens, extractLibreChatParams, optionalChainWithEmptyCheck } from '~/utils';
+import {
+  optionalChainWithEmptyCheck,
+  extractLibreChatParams,
+  getModelMaxTokens,
+  getThreadData,
+} from '~/utils';
 import { filterFilesByEndpointConfig } from '~/files';
 import { generateArtifactsPrompt } from '~/prompts';
 import { getProviderConfig } from '~/endpoints';
@@ -58,6 +64,8 @@ export interface InitializeAgentParams {
   agent: Agent;
   /** Conversation ID (optional) */
   conversationId?: string | null;
+  /** Parent message ID for determining the current thread (optional) */
+  parentMessageId?: string | null;
   /** Request files */
   requestFiles?: IMongoFile[];
   /** Function to load agent tools */
@@ -95,10 +103,23 @@ export interface InitializeAgentDbMethods extends EndpointDbMethods {
   updateFilesUsage: (files: Array<{ file_id: string }>, fileIds?: string[]) => Promise<unknown[]>;
   /** Get files from database */
   getFiles: (filter: unknown, sort: unknown, select: unknown, opts?: unknown) => Promise<unknown[]>;
-  /** Get tool files by IDs */
+  /** Get tool files by IDs (user-uploaded files only, code files handled separately) */
   getToolFilesByIds: (fileIds: string[], toolSet: Set<EToolResources>) => Promise<unknown[]>;
   /** Get conversation file IDs */
   getConvoFiles: (conversationId: string) => Promise<string[] | null>;
+  /** Get code-generated files by conversation ID and optional message IDs */
+  getCodeGeneratedFiles?: (conversationId: string, messageIds?: string[]) => Promise<unknown[]>;
+  /** Get user-uploaded execute_code files by file IDs (from message.files in thread) */
+  getUserCodeFiles?: (fileIds: string[]) => Promise<unknown[]>;
+  /** Get messages for a conversation (supports select for field projection) */
+  getMessages?: (
+    filter: { conversationId: string },
+    select?: string,
+  ) => Promise<Array<{
+    messageId: string;
+    parentMessageId?: string;
+    files?: Array<{ file_id: string }>;
+  }> | null>;
 }
 
 /**
@@ -125,6 +146,7 @@ export async function initializeAgent(
     requestFiles = [],
     conversationId,
     endpointOption,
+    parentMessageId,
     allowedProviders,
     isInitialAgent = false,
   } = params;
@@ -174,9 +196,51 @@ export async function initializeAgent(
         toolResourceSet.add(EToolResources[tool as keyof typeof EToolResources]);
       }
     }
+
     const toolFiles = (await db.getToolFilesByIds(fileIds, toolResourceSet)) as IMongoFile[];
-    if (requestFiles.length || toolFiles.length) {
-      currentFiles = (await db.updateFilesUsage(requestFiles.concat(toolFiles))) as IMongoFile[];
+
+    /**
+     * Retrieve execute_code files filtered to the current thread.
+     * This includes both code-generated files and user-uploaded execute_code files.
+     */
+    let codeGeneratedFiles: IMongoFile[] = [];
+    let userCodeFiles: IMongoFile[] = [];
+
+    if (toolResourceSet.has(EToolResources.execute_code)) {
+      let threadMessageIds: string[] | undefined;
+      let threadFileIds: string[] | undefined;
+
+      if (parentMessageId && parentMessageId !== Constants.NO_PARENT && db.getMessages) {
+        /** Only select fields needed for thread traversal */
+        const messages = await db.getMessages(
+          { conversationId },
+          'messageId parentMessageId files',
+        );
+        if (messages && messages.length > 0) {
+          /** Single O(n) pass: build Map, traverse thread, collect both IDs */
+          const threadData = getThreadData(messages, parentMessageId);
+          threadMessageIds = threadData.messageIds;
+          threadFileIds = threadData.fileIds;
+        }
+      }
+
+      /** Code-generated files (context: execute_code) filtered by messageId */
+      if (db.getCodeGeneratedFiles) {
+        codeGeneratedFiles = (await db.getCodeGeneratedFiles(
+          conversationId,
+          threadMessageIds,
+        )) as IMongoFile[];
+      }
+
+      /** User-uploaded execute_code files (context: agents/message_attachment) from thread messages */
+      if (db.getUserCodeFiles && threadFileIds && threadFileIds.length > 0) {
+        userCodeFiles = (await db.getUserCodeFiles(threadFileIds)) as IMongoFile[];
+      }
+    }
+
+    const allToolFiles = toolFiles.concat(codeGeneratedFiles, userCodeFiles);
+    if (requestFiles.length || allToolFiles.length) {
+      currentFiles = (await db.updateFilesUsage(requestFiles.concat(allToolFiles))) as IMongoFile[];
     }
   } else if (requestFiles.length) {
     currentFiles = (await db.updateFilesUsage(requestFiles)) as IMongoFile[];