mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-01-20 17:26:12 +01:00
🗂️ 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.
This commit is contained in:
Danny Avila
parent
fe32cbedf9
commit
cc32895d13
22 changed files with 1364 additions and 83 deletions
|
|
@ -31,6 +31,7 @@ setGetAgent(getAgent);
|
|||
* @param {Function} params.loadTools - Function to load agent tools
|
||||
* @param {Array} params.requestFiles - Request files
|
||||
* @param {string} params.conversationId - The conversation ID
|
||||
* @param {string} [params.parentMessageId] - The parent message ID for thread filtering
|
||||
* @param {Set} params.allowedProviders - Set of allowed providers
|
||||
* @param {Map} params.agentConfigs - Map of agent configs to add to
|
||||
* @param {string} params.primaryAgentId - The primary agent ID
|
||||
|
|
@ -46,6 +47,7 @@ const processAddedConvo = async ({
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
allowedProviders,
|
||||
agentConfigs,
|
||||
primaryAgentId,
|
||||
|
|
@ -91,6 +93,7 @@ const processAddedConvo = async ({
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
agent: addedAgent,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
|
|
@ -99,9 +102,12 @@ const processAddedConvo = async ({
|
|||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -3,10 +3,10 @@ const { createContentAggregator } = require('@librechat/agents');
|
|||
const {
|
||||
initializeAgent,
|
||||
validateAgentModel,
|
||||
getCustomEndpointConfig,
|
||||
createSequentialChainEdges,
|
||||
createEdgeCollector,
|
||||
filterOrphanedEdges,
|
||||
getCustomEndpointConfig,
|
||||
createSequentialChainEdges,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
EModelEndpoint,
|
||||
|
|
@ -129,6 +129,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
const requestFiles = req.body.files ?? [];
|
||||
/** @type {string} */
|
||||
const conversationId = req.body.conversationId;
|
||||
/** @type {string | undefined} */
|
||||
const parentMessageId = req.body.parentMessageId;
|
||||
|
||||
const primaryConfig = await initializeAgent(
|
||||
{
|
||||
|
|
@ -137,6 +139,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
agent: primaryAgent,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
|
|
@ -146,9 +149,12 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -188,6 +194,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
loadTools,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
endpointOption,
|
||||
allowedProviders,
|
||||
},
|
||||
|
|
@ -195,9 +202,12 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
getConvoFiles,
|
||||
getFiles: db.getFiles,
|
||||
getUserKey: db.getUserKey,
|
||||
getMessages: db.getMessages,
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
getUserKeyValues: db.getUserKeyValues,
|
||||
getUserCodeFiles: db.getUserCodeFiles,
|
||||
getToolFilesByIds: db.getToolFilesByIds,
|
||||
getCodeGeneratedFiles: db.getCodeGeneratedFiles,
|
||||
},
|
||||
);
|
||||
if (userMCPAuthMap != null) {
|
||||
|
|
@ -252,17 +262,18 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
const { userMCPAuthMap: updatedMCPAuthMap } = await processAddedConvo({
|
||||
req,
|
||||
res,
|
||||
endpointOption,
|
||||
modelsConfig,
|
||||
logViolation,
|
||||
loadTools,
|
||||
logViolation,
|
||||
modelsConfig,
|
||||
requestFiles,
|
||||
conversationId,
|
||||
allowedProviders,
|
||||
agentConfigs,
|
||||
primaryAgentId: primaryConfig.id,
|
||||
primaryAgent,
|
||||
endpointOption,
|
||||
userMCPAuthMap,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
allowedProviders,
|
||||
primaryAgentId: primaryConfig.id,
|
||||
});
|
||||
|
||||
if (updatedMCPAuthMap) {
|
||||
|
|
|
|||
|
|
@ -6,27 +6,112 @@ const { getCodeBaseURL } = require('@librechat/agents');
|
|||
const { logAxiosError, getBasePath } = require('@librechat/api');
|
||||
const {
|
||||
Tools,
|
||||
megabyte,
|
||||
fileConfig,
|
||||
FileContext,
|
||||
FileSources,
|
||||
imageExtRegex,
|
||||
inferMimeType,
|
||||
EToolResources,
|
||||
EModelEndpoint,
|
||||
mergeFileConfig,
|
||||
getEndpointFileConfig,
|
||||
} = require('librechat-data-provider');
|
||||
const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions');
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { convertImage } = require('~/server/services/Files/images/convert');
|
||||
const { createFile, getFiles, updateFile } = require('~/models');
|
||||
const { determineFileType } = require('~/server/utils');
|
||||
|
||||
/**
|
||||
* Process OpenAI image files, convert to target format, save and return file metadata.
|
||||
* Creates a fallback download URL response when file cannot be processed locally.
|
||||
* Used when: file exceeds size limit, storage strategy unavailable, or download error occurs.
|
||||
* @param {Object} params - The parameters.
|
||||
* @param {string} params.name - The filename.
|
||||
* @param {string} params.session_id - The code execution session ID.
|
||||
* @param {string} params.id - The file ID from the code environment.
|
||||
* @param {string} params.conversationId - The current conversation ID.
|
||||
* @param {string} params.toolCallId - The tool call ID that generated the file.
|
||||
* @param {string} params.messageId - The current message ID.
|
||||
* @param {number} params.expiresAt - Expiration timestamp (24 hours from creation).
|
||||
* @returns {Object} Fallback response with download URL.
|
||||
*/
|
||||
const createDownloadFallback = ({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
expiresAt,
|
||||
session_id,
|
||||
toolCallId,
|
||||
conversationId,
|
||||
}) => {
|
||||
const basePath = getBasePath();
|
||||
return {
|
||||
filename: name,
|
||||
filepath: `${basePath}/api/files/code/download/${session_id}/${id}`,
|
||||
expiresAt,
|
||||
conversationId,
|
||||
toolCallId,
|
||||
messageId,
|
||||
};
|
||||
};
|
||||
|
||||
/**
|
||||
* Find an existing code-generated file by filename in the conversation.
|
||||
* Used to update existing files instead of creating duplicates.
|
||||
*
|
||||
* ## Deduplication Strategy
|
||||
*
|
||||
* Files are deduplicated by `(conversationId, filename)` - NOT including `messageId`.
|
||||
* This is an intentional design decision to handle iterative code development patterns:
|
||||
*
|
||||
* **Rationale:**
|
||||
* - When users iteratively refine code (e.g., "regenerate that chart with red bars"),
|
||||
* the same logical file (e.g., "chart.png") is produced multiple times
|
||||
* - Without deduplication, each iteration would create a new file, leading to storage bloat
|
||||
* - The latest version is what matters for re-upload to the code environment
|
||||
*
|
||||
* **Implications:**
|
||||
* - Different messages producing files with the same name will update the same file record
|
||||
* - The `messageId` field tracks which message last updated the file
|
||||
* - The `usage` counter tracks how many times the file has been generated
|
||||
*
|
||||
* **Future Considerations:**
|
||||
* - If file versioning is needed, consider adding a `versions` array or separate version collection
|
||||
* - The current approach prioritizes storage efficiency over history preservation
|
||||
*
|
||||
* @param {string} filename - The filename to search for.
|
||||
* @param {string} conversationId - The conversation ID.
|
||||
* @returns {Promise<MongoFile | null>} The existing file or null.
|
||||
*/
|
||||
const findExistingCodeFile = async (filename, conversationId) => {
|
||||
if (!filename || !conversationId) {
|
||||
return null;
|
||||
}
|
||||
const files = await getFiles(
|
||||
{
|
||||
filename,
|
||||
conversationId,
|
||||
context: FileContext.execute_code,
|
||||
},
|
||||
{ createdAt: -1 },
|
||||
{ text: 0 },
|
||||
);
|
||||
return files?.[0] ?? null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Process code execution output files - downloads and saves both images and non-image files.
|
||||
* All files are saved to local storage with fileIdentifier metadata for code env re-upload.
|
||||
* @param {ServerRequest} params.req - The Express request object.
|
||||
* @param {string} params.id - The file ID.
|
||||
* @param {string} params.id - The file ID from the code environment.
|
||||
* @param {string} params.name - The filename.
|
||||
* @param {string} params.apiKey - The code execution API key.
|
||||
* @param {string} params.toolCallId - The tool call ID that generated the file.
|
||||
* @param {string} params.session_id - The code execution session ID.
|
||||
* @param {string} params.conversationId - The current conversation ID.
|
||||
* @param {string} params.messageId - The current message ID.
|
||||
* @returns {Promise<MongoFile & { messageId: string, toolCallId: string } | { filename: string; filepath: string; expiresAt: number; conversationId: string; toolCallId: string; messageId: string } | undefined>} The file metadata or undefined if an error occurs.
|
||||
* @returns {Promise<MongoFile & { messageId: string, toolCallId: string } | undefined>} The file metadata or undefined if an error occurs.
|
||||
*/
|
||||
const processCodeOutput = async ({
|
||||
req,
|
||||
|
|
@ -41,19 +126,15 @@ const processCodeOutput = async ({
|
|||
const appConfig = req.config;
|
||||
const currentDate = new Date();
|
||||
const baseURL = getCodeBaseURL();
|
||||
const basePath = getBasePath();
|
||||
const fileExt = path.extname(name);
|
||||
if (!fileExt || !imageExtRegex.test(name)) {
|
||||
return {
|
||||
filename: name,
|
||||
filepath: `${basePath}/api/files/code/download/${session_id}/${id}`,
|
||||
/** Note: expires 24 hours after creation */
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
conversationId,
|
||||
toolCallId,
|
||||
messageId,
|
||||
};
|
||||
}
|
||||
const fileExt = path.extname(name).toLowerCase();
|
||||
const isImage = fileExt && imageExtRegex.test(name);
|
||||
|
||||
const mergedFileConfig = mergeFileConfig(appConfig.fileConfig);
|
||||
const endpointFileConfig = getEndpointFileConfig({
|
||||
fileConfig: mergedFileConfig,
|
||||
endpoint: EModelEndpoint.agents,
|
||||
});
|
||||
const fileSizeLimit = endpointFileConfig.fileSizeLimit ?? mergedFileConfig.serverFileSizeLimit;
|
||||
|
||||
try {
|
||||
const formattedDate = currentDate.toISOString();
|
||||
|
|
@ -70,29 +151,135 @@ const processCodeOutput = async ({
|
|||
|
||||
const buffer = Buffer.from(response.data, 'binary');
|
||||
|
||||
const file_id = v4();
|
||||
const _file = await convertImage(req, buffer, 'high', `${file_id}${fileExt}`);
|
||||
// Enforce file size limit
|
||||
if (buffer.length > fileSizeLimit) {
|
||||
logger.warn(
|
||||
`[processCodeOutput] File "${name}" (${(buffer.length / megabyte).toFixed(2)} MB) exceeds size limit of ${(fileSizeLimit / megabyte).toFixed(2)} MB, falling back to download URL`,
|
||||
);
|
||||
return createDownloadFallback({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
toolCallId,
|
||||
session_id,
|
||||
conversationId,
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
});
|
||||
}
|
||||
|
||||
const fileIdentifier = `${session_id}/${id}`;
|
||||
|
||||
/**
|
||||
* Check for existing file with same filename in this conversation.
|
||||
* If found, we'll update it instead of creating a duplicate.
|
||||
*/
|
||||
const existingFile = await findExistingCodeFile(name, conversationId);
|
||||
const file_id = existingFile?.file_id ?? v4();
|
||||
const isUpdate = !!existingFile;
|
||||
|
||||
if (isUpdate) {
|
||||
logger.debug(
|
||||
`[processCodeOutput] Updating existing file "${name}" (${file_id}) instead of creating duplicate`,
|
||||
);
|
||||
}
|
||||
|
||||
if (isImage) {
|
||||
const _file = await convertImage(req, buffer, 'high', `${file_id}${fileExt}`);
|
||||
const file = {
|
||||
..._file,
|
||||
file_id,
|
||||
messageId,
|
||||
usage: isUpdate ? (existingFile.usage ?? 0) + 1 : 1,
|
||||
filename: name,
|
||||
conversationId,
|
||||
user: req.user.id,
|
||||
type: `image/${appConfig.imageOutputType}`,
|
||||
createdAt: isUpdate ? existingFile.createdAt : formattedDate,
|
||||
updatedAt: formattedDate,
|
||||
source: appConfig.fileStrategy,
|
||||
context: FileContext.execute_code,
|
||||
metadata: { fileIdentifier },
|
||||
};
|
||||
createFile(file, true);
|
||||
return Object.assign(file, { messageId, toolCallId });
|
||||
}
|
||||
|
||||
// For non-image files, save to configured storage strategy
|
||||
const { saveBuffer } = getStrategyFunctions(appConfig.fileStrategy);
|
||||
if (!saveBuffer) {
|
||||
logger.warn(
|
||||
`[processCodeOutput] saveBuffer not available for strategy ${appConfig.fileStrategy}, falling back to download URL`,
|
||||
);
|
||||
return createDownloadFallback({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
toolCallId,
|
||||
session_id,
|
||||
conversationId,
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
});
|
||||
}
|
||||
|
||||
// Determine MIME type from buffer or extension
|
||||
const detectedType = await determineFileType(buffer, true);
|
||||
const mimeType = detectedType?.mime || inferMimeType(name, '') || 'application/octet-stream';
|
||||
|
||||
/** Check MIME type support - for code-generated files, we're lenient but log unsupported types */
|
||||
const isSupportedMimeType = fileConfig.checkType(
|
||||
mimeType,
|
||||
endpointFileConfig.supportedMimeTypes,
|
||||
);
|
||||
if (!isSupportedMimeType) {
|
||||
logger.warn(
|
||||
`[processCodeOutput] File "${name}" has unsupported MIME type "${mimeType}", proceeding with storage but may not be usable as tool resource`,
|
||||
);
|
||||
}
|
||||
|
||||
const fileName = `${file_id}__${name}`;
|
||||
const filepath = await saveBuffer({
|
||||
userId: req.user.id,
|
||||
buffer,
|
||||
fileName,
|
||||
basePath: 'uploads',
|
||||
});
|
||||
|
||||
const file = {
|
||||
..._file,
|
||||
file_id,
|
||||
usage: 1,
|
||||
filepath,
|
||||
messageId,
|
||||
object: 'file',
|
||||
filename: name,
|
||||
type: mimeType,
|
||||
conversationId,
|
||||
user: req.user.id,
|
||||
type: `image/${appConfig.imageOutputType}`,
|
||||
createdAt: formattedDate,
|
||||
bytes: buffer.length,
|
||||
updatedAt: formattedDate,
|
||||
metadata: { fileIdentifier },
|
||||
source: appConfig.fileStrategy,
|
||||
context: FileContext.execute_code,
|
||||
usage: isUpdate ? (existingFile.usage ?? 0) + 1 : 1,
|
||||
createdAt: isUpdate ? existingFile.createdAt : formattedDate,
|
||||
};
|
||||
|
||||
createFile(file, true);
|
||||
/** Note: `messageId` & `toolCallId` are not part of file DB schema; message object records associated file ID */
|
||||
return Object.assign(file, { messageId, toolCallId });
|
||||
} catch (error) {
|
||||
logAxiosError({
|
||||
message: 'Error downloading code environment file',
|
||||
message: 'Error downloading/processing code environment file',
|
||||
error,
|
||||
});
|
||||
|
||||
// Fallback for download errors - return download URL so user can still manually download
|
||||
return createDownloadFallback({
|
||||
id,
|
||||
name,
|
||||
messageId,
|
||||
toolCallId,
|
||||
session_id,
|
||||
conversationId,
|
||||
expiresAt: currentDate.getTime() + 86400000,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -204,9 +391,16 @@ const primeFiles = async (options, apiKey) => {
|
|||
if (!toolContext) {
|
||||
toolContext = `- Note: The following files are available in the "${Tools.execute_code}" tool environment:`;
|
||||
}
|
||||
toolContext += `\n\t- /mnt/data/${file.filename}${
|
||||
agentResourceIds.has(file.file_id) ? '' : ' (just attached by user)'
|
||||
}`;
|
||||
|
||||
let fileSuffix = '';
|
||||
if (!agentResourceIds.has(file.file_id)) {
|
||||
fileSuffix =
|
||||
file.context === FileContext.execute_code
|
||||
? ' (from previous code execution)'
|
||||
: ' (attached by user)';
|
||||
}
|
||||
|
||||
toolContext += `\n\t- /mnt/data/${file.filename}${fileSuffix}`;
|
||||
files.push({
|
||||
id,
|
||||
session_id,
|
||||
|
|
|
|||
418
api/server/services/Files/Code/process.spec.js
Normal file
418
api/server/services/Files/Code/process.spec.js
Normal file
|
|
@ -0,0 +1,418 @@
|
|||
// Configurable file size limit for tests - use a getter so it can be changed per test
|
||||
const fileSizeLimitConfig = { value: 20 * 1024 * 1024 }; // Default 20MB
|
||||
|
||||
// Mock librechat-data-provider with configurable file size limit
|
||||
jest.mock('librechat-data-provider', () => {
|
||||
const actual = jest.requireActual('librechat-data-provider');
|
||||
return {
|
||||
...actual,
|
||||
mergeFileConfig: jest.fn((config) => {
|
||||
const merged = actual.mergeFileConfig(config);
|
||||
// Override the serverFileSizeLimit with our test value
|
||||
return {
|
||||
...merged,
|
||||
get serverFileSizeLimit() {
|
||||
return fileSizeLimitConfig.value;
|
||||
},
|
||||
};
|
||||
}),
|
||||
getEndpointFileConfig: jest.fn((options) => {
|
||||
const config = actual.getEndpointFileConfig(options);
|
||||
// Override fileSizeLimit with our test value
|
||||
return {
|
||||
...config,
|
||||
get fileSizeLimit() {
|
||||
return fileSizeLimitConfig.value;
|
||||
},
|
||||
};
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
const { FileContext } = require('librechat-data-provider');
|
||||
|
||||
// Mock uuid
|
||||
jest.mock('uuid', () => ({
|
||||
v4: jest.fn(() => 'mock-uuid-1234'),
|
||||
}));
|
||||
|
||||
// Mock axios
|
||||
jest.mock('axios');
|
||||
const axios = require('axios');
|
||||
|
||||
// Mock logger
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
warn: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock getCodeBaseURL
|
||||
jest.mock('@librechat/agents', () => ({
|
||||
getCodeBaseURL: jest.fn(() => 'https://code-api.example.com'),
|
||||
}));
|
||||
|
||||
// Mock logAxiosError and getBasePath
|
||||
jest.mock('@librechat/api', () => ({
|
||||
logAxiosError: jest.fn(),
|
||||
getBasePath: jest.fn(() => ''),
|
||||
}));
|
||||
|
||||
// Mock models
|
||||
jest.mock('~/models', () => ({
|
||||
createFile: jest.fn(),
|
||||
getFiles: jest.fn(),
|
||||
updateFile: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock permissions (must be before process.js import)
|
||||
jest.mock('~/server/services/Files/permissions', () => ({
|
||||
filterFilesByAgentAccess: jest.fn((options) => Promise.resolve(options.files)),
|
||||
}));
|
||||
|
||||
// Mock strategy functions
|
||||
jest.mock('~/server/services/Files/strategies', () => ({
|
||||
getStrategyFunctions: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock convertImage
|
||||
jest.mock('~/server/services/Files/images/convert', () => ({
|
||||
convertImage: jest.fn(),
|
||||
}));
|
||||
|
||||
// Mock determineFileType
|
||||
jest.mock('~/server/utils', () => ({
|
||||
determineFileType: jest.fn(),
|
||||
}));
|
||||
|
||||
const { createFile, getFiles } = require('~/models');
|
||||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { convertImage } = require('~/server/services/Files/images/convert');
|
||||
const { determineFileType } = require('~/server/utils');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
||||
// Import after mocks
|
||||
const { processCodeOutput } = require('./process');
|
||||
|
||||
describe('Code Process', () => {
|
||||
const mockReq = {
|
||||
user: { id: 'user-123' },
|
||||
config: {
|
||||
fileConfig: {},
|
||||
fileStrategy: 'local',
|
||||
imageOutputType: 'webp',
|
||||
},
|
||||
};
|
||||
|
||||
const baseParams = {
|
||||
req: mockReq,
|
||||
id: 'file-id-123',
|
||||
name: 'test-file.txt',
|
||||
apiKey: 'test-api-key',
|
||||
toolCallId: 'tool-call-123',
|
||||
conversationId: 'conv-123',
|
||||
messageId: 'msg-123',
|
||||
session_id: 'session-123',
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
// Default mock implementations
|
||||
getFiles.mockResolvedValue(null);
|
||||
createFile.mockResolvedValue({});
|
||||
getStrategyFunctions.mockReturnValue({
|
||||
saveBuffer: jest.fn().mockResolvedValue('/uploads/mock-file-path.txt'),
|
||||
});
|
||||
determineFileType.mockResolvedValue({ mime: 'text/plain' });
|
||||
});
|
||||
|
||||
describe('findExistingCodeFile (via processCodeOutput)', () => {
|
||||
it('should find existing file by filename and conversationId', async () => {
|
||||
const existingFile = {
|
||||
file_id: 'existing-file-id',
|
||||
filename: 'test-file.txt',
|
||||
usage: 2,
|
||||
createdAt: '2024-01-01T00:00:00.000Z',
|
||||
};
|
||||
getFiles.mockResolvedValue([existingFile]);
|
||||
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
// Verify getFiles was called with correct deduplication query
|
||||
expect(getFiles).toHaveBeenCalledWith(
|
||||
{
|
||||
filename: 'test-file.txt',
|
||||
conversationId: 'conv-123',
|
||||
context: FileContext.execute_code,
|
||||
},
|
||||
{ createdAt: -1 },
|
||||
{ text: 0 },
|
||||
);
|
||||
|
||||
// Verify the existing file_id was reused
|
||||
expect(result.file_id).toBe('existing-file-id');
|
||||
// Verify usage was incremented
|
||||
expect(result.usage).toBe(3);
|
||||
// Verify original createdAt was preserved
|
||||
expect(result.createdAt).toBe('2024-01-01T00:00:00.000Z');
|
||||
});
|
||||
|
||||
it('should create new file when no existing file found', async () => {
|
||||
getFiles.mockResolvedValue(null);
|
||||
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
// Should use the mocked uuid
|
||||
expect(result.file_id).toBe('mock-uuid-1234');
|
||||
// Should have usage of 1 for new file
|
||||
expect(result.usage).toBe(1);
|
||||
});
|
||||
|
||||
it('should return null for invalid inputs (empty filename)', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
// The function handles this internally - with empty name
|
||||
// findExistingCodeFile returns null early for empty filename (guard clause)
|
||||
const result = await processCodeOutput({ ...baseParams, name: '' });
|
||||
|
||||
// getFiles should NOT be called due to early return in findExistingCodeFile
|
||||
expect(getFiles).not.toHaveBeenCalled();
|
||||
// A new file_id should be generated since no existing file was found
|
||||
expect(result.file_id).toBe('mock-uuid-1234');
|
||||
});
|
||||
});
|
||||
|
||||
describe('processCodeOutput', () => {
|
||||
describe('image file processing', () => {
|
||||
it('should process image files using convertImage', async () => {
|
||||
const imageParams = { ...baseParams, name: 'chart.png' };
|
||||
const imageBuffer = Buffer.alloc(500);
|
||||
axios.mockResolvedValue({ data: imageBuffer });
|
||||
|
||||
const convertedFile = {
|
||||
filepath: '/uploads/converted-image.webp',
|
||||
bytes: 400,
|
||||
};
|
||||
convertImage.mockResolvedValue(convertedFile);
|
||||
getFiles.mockResolvedValue(null);
|
||||
|
||||
const result = await processCodeOutput(imageParams);
|
||||
|
||||
expect(convertImage).toHaveBeenCalledWith(
|
||||
mockReq,
|
||||
imageBuffer,
|
||||
'high',
|
||||
'mock-uuid-1234.png',
|
||||
);
|
||||
expect(result.type).toBe('image/webp');
|
||||
expect(result.context).toBe(FileContext.execute_code);
|
||||
expect(result.filename).toBe('chart.png');
|
||||
});
|
||||
|
||||
it('should update existing image file and increment usage', async () => {
|
||||
const imageParams = { ...baseParams, name: 'chart.png' };
|
||||
const existingFile = {
|
||||
file_id: 'existing-img-id',
|
||||
usage: 1,
|
||||
createdAt: '2024-01-01T00:00:00.000Z',
|
||||
};
|
||||
getFiles.mockResolvedValue([existingFile]);
|
||||
|
||||
const imageBuffer = Buffer.alloc(500);
|
||||
axios.mockResolvedValue({ data: imageBuffer });
|
||||
convertImage.mockResolvedValue({ filepath: '/uploads/img.webp' });
|
||||
|
||||
const result = await processCodeOutput(imageParams);
|
||||
|
||||
expect(result.file_id).toBe('existing-img-id');
|
||||
expect(result.usage).toBe(2);
|
||||
expect(logger.debug).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Updating existing file'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('non-image file processing', () => {
|
||||
it('should process non-image files using saveBuffer', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const mockSaveBuffer = jest.fn().mockResolvedValue('/uploads/saved-file.txt');
|
||||
getStrategyFunctions.mockReturnValue({ saveBuffer: mockSaveBuffer });
|
||||
determineFileType.mockResolvedValue({ mime: 'text/plain' });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(mockSaveBuffer).toHaveBeenCalledWith({
|
||||
userId: 'user-123',
|
||||
buffer: smallBuffer,
|
||||
fileName: 'mock-uuid-1234__test-file.txt',
|
||||
basePath: 'uploads',
|
||||
});
|
||||
expect(result.type).toBe('text/plain');
|
||||
expect(result.filepath).toBe('/uploads/saved-file.txt');
|
||||
expect(result.bytes).toBe(100);
|
||||
});
|
||||
|
||||
it('should detect MIME type from buffer', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
determineFileType.mockResolvedValue({ mime: 'application/pdf' });
|
||||
|
||||
const result = await processCodeOutput({ ...baseParams, name: 'document.pdf' });
|
||||
|
||||
expect(determineFileType).toHaveBeenCalledWith(smallBuffer, true);
|
||||
expect(result.type).toBe('application/pdf');
|
||||
});
|
||||
|
||||
it('should fallback to application/octet-stream for unknown types', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
determineFileType.mockResolvedValue(null);
|
||||
|
||||
const result = await processCodeOutput({ ...baseParams, name: 'unknown.xyz' });
|
||||
|
||||
expect(result.type).toBe('application/octet-stream');
|
||||
});
|
||||
});
|
||||
|
||||
describe('file size limit enforcement', () => {
|
||||
it('should fallback to download URL when file exceeds size limit', async () => {
|
||||
// Set a small file size limit for this test
|
||||
fileSizeLimitConfig.value = 1000; // 1KB limit
|
||||
|
||||
const largeBuffer = Buffer.alloc(5000); // 5KB - exceeds 1KB limit
|
||||
axios.mockResolvedValue({ data: largeBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('exceeds size limit'));
|
||||
expect(result.filepath).toContain('/api/files/code/download/session-123/file-id-123');
|
||||
expect(result.expiresAt).toBeDefined();
|
||||
// Should not call createFile for oversized files (fallback path)
|
||||
expect(createFile).not.toHaveBeenCalled();
|
||||
|
||||
// Reset to default for other tests
|
||||
fileSizeLimitConfig.value = 20 * 1024 * 1024;
|
||||
});
|
||||
});
|
||||
|
||||
describe('fallback behavior', () => {
|
||||
it('should fallback to download URL when saveBuffer is not available', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
getStrategyFunctions.mockReturnValue({ saveBuffer: null });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(logger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('saveBuffer not available'),
|
||||
);
|
||||
expect(result.filepath).toContain('/api/files/code/download/');
|
||||
expect(result.filename).toBe('test-file.txt');
|
||||
});
|
||||
|
||||
it('should fallback to download URL on axios error', async () => {
|
||||
axios.mockRejectedValue(new Error('Network error'));
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.filepath).toContain('/api/files/code/download/session-123/file-id-123');
|
||||
expect(result.conversationId).toBe('conv-123');
|
||||
expect(result.messageId).toBe('msg-123');
|
||||
expect(result.toolCallId).toBe('tool-call-123');
|
||||
});
|
||||
});
|
||||
|
||||
describe('usage counter increment', () => {
|
||||
it('should set usage to 1 for new files', async () => {
|
||||
getFiles.mockResolvedValue(null);
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.usage).toBe(1);
|
||||
});
|
||||
|
||||
it('should increment usage for existing files', async () => {
|
||||
const existingFile = { file_id: 'existing-id', usage: 5, createdAt: '2024-01-01' };
|
||||
getFiles.mockResolvedValue([existingFile]);
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.usage).toBe(6);
|
||||
});
|
||||
|
||||
it('should handle existing file with undefined usage', async () => {
|
||||
const existingFile = { file_id: 'existing-id', createdAt: '2024-01-01' };
|
||||
getFiles.mockResolvedValue([existingFile]);
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
// (undefined ?? 0) + 1 = 1
|
||||
expect(result.usage).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('metadata and file properties', () => {
|
||||
it('should include fileIdentifier in metadata', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.metadata).toEqual({
|
||||
fileIdentifier: 'session-123/file-id-123',
|
||||
});
|
||||
});
|
||||
|
||||
it('should set correct context for code-generated files', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.context).toBe(FileContext.execute_code);
|
||||
});
|
||||
|
||||
it('should include toolCallId and messageId in result', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
const result = await processCodeOutput(baseParams);
|
||||
|
||||
expect(result.toolCallId).toBe('tool-call-123');
|
||||
expect(result.messageId).toBe('msg-123');
|
||||
});
|
||||
|
||||
it('should call createFile with upsert enabled', async () => {
|
||||
const smallBuffer = Buffer.alloc(100);
|
||||
axios.mockResolvedValue({ data: smallBuffer });
|
||||
|
||||
await processCodeOutput(baseParams);
|
||||
|
||||
expect(createFile).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
file_id: 'mock-uuid-1234',
|
||||
context: FileContext.execute_code,
|
||||
}),
|
||||
true, // upsert flag
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -67,7 +67,12 @@ async function saveLocalBuffer({ userId, buffer, fileName, basePath = 'images' }
|
|||
try {
|
||||
const { publicPath, uploads } = paths;
|
||||
|
||||
const directoryPath = path.join(basePath === 'images' ? publicPath : uploads, basePath, userId);
|
||||
/**
|
||||
* For 'images': save to publicPath/images/userId (images are served statically)
|
||||
* For 'uploads': save to uploads/userId (files downloaded via API)
|
||||
* */
|
||||
const directoryPath =
|
||||
basePath === 'images' ? path.join(publicPath, basePath, userId) : path.join(uploads, userId);
|
||||
|
||||
if (!fs.existsSync(directoryPath)) {
|
||||
fs.mkdirSync(directoryPath, { recursive: true });
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue