mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-01-20 09:16:13 +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
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
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue