diff --git a/api/server/middleware/buildEndpointOption.js b/api/server/middleware/buildEndpointOption.js index 1eaa1ef8d9..35bb36830c 100644 --- a/api/server/middleware/buildEndpointOption.js +++ b/api/server/middleware/buildEndpointOption.js @@ -132,7 +132,9 @@ async function buildEndpointOption(req, res, next) { req.body.endpointOption = await builder(endpoint, parsedBody, endpointType); if (req.body.files && !isAgents) { - req.body.endpointOption.attachments = updateFilesUsage(req.body.files); + req.body.endpointOption.attachments = updateFilesUsage(req.body.files, undefined, { + user: req.user.id, + }); } next(); diff --git a/api/server/middleware/buildEndpointOption.spec.js b/api/server/middleware/buildEndpointOption.spec.js index 9c353b498a..893a14471b 100644 --- a/api/server/middleware/buildEndpointOption.spec.js +++ b/api/server/middleware/buildEndpointOption.spec.js @@ -35,6 +35,7 @@ jest.mock('~/server/services/Endpoints/agents', () => ({ jest.mock('~/models', () => ({ updateFilesUsage: jest.fn(), })); +const { updateFilesUsage } = require('~/models'); const mockGetEndpointsConfig = jest.fn(); jest.mock('~/server/services/Config', () => ({ @@ -417,6 +418,29 @@ describe('buildEndpointOption - defaultParamsEndpoint parsing', () => { expect(parsedResult.max_tokens).toBe(4096); }); + it('should scope non-agent chat attachment usage updates to the authenticated user', async () => { + const attachments = Promise.resolve([]); + updateFilesUsage.mockReturnValueOnce(attachments); + mockGetEndpointsConfig.mockResolvedValue({}); + + const req = createReq( + { + endpoint: EModelEndpoint.assistants, + assistant_id: 'asst_123', + files: [{ file_id: 'forged-file-id' }], + }, + { modelSpecs: null }, + ); + req.user = { id: 'user-1' }; + + await buildEndpointOption(req, createRes(), jest.fn()); + + expect(updateFilesUsage).toHaveBeenCalledWith(req.body.files, undefined, { + user: 'user-1', + }); + expect(req.body.endpointOption.attachments).toBe(attachments); + }); + it('should not enter the enforce branch when modelSpecs.list is empty', async () => { mockGetEndpointsConfig.mockResolvedValue({}); diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index adb6689bc5..857638f59a 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -33,7 +33,8 @@ jest.mock('@librechat/agents', () => ({ })); import { Providers } from '@librechat/agents'; -import { EModelEndpoint, Tools } from 'librechat-data-provider'; +import { EModelEndpoint, EToolResources, Tools } from 'librechat-data-provider'; +import type { IMongoFile } from '@librechat/data-schemas'; import type { Agent } from 'librechat-data-provider'; import type { ServerRequest, InitializeResultBase, EndpointTokenConfig } from '~/types'; import type { InitializeAgentDbMethods } from '../initialize'; @@ -750,6 +751,44 @@ describe('initializeAgent — attachment scoping', () => { expect(result.requestAttachments).toEqual([requestFile]); expect(result.agentContextAttachments).toEqual([agentContextFile]); }); + + it('owner-scopes request file usage updates while preserving trusted tool files', async () => { + const requestFile = { file_id: 'request-file', filename: 'request.txt' } as IMongoFile; + const toolFile = { file_id: 'tool-file', filename: 'tool.txt' } as IMongoFile; + const { agent, req, res, loadTools, db } = createMocks(); + + agent.tools = [EToolResources.file_search]; + mockExtractLibreChatParams.mockReturnValueOnce({ + resendFiles: true, + maxContextTokens: undefined, + modelOptions: { model: agent.model }, + }); + (db.getConvoFiles as jest.Mock).mockResolvedValueOnce([toolFile.file_id]); + (db.getToolFilesByIds as jest.Mock).mockResolvedValueOnce([toolFile]); + (db.updateFilesUsage as jest.Mock) + .mockResolvedValueOnce([requestFile]) + .mockResolvedValueOnce([toolFile]); + + await initializeAgent( + { + req, + res, + agent, + loadTools, + requestFiles: [requestFile], + conversationId: 'conversation-1', + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([Providers.OPENAI]), + isInitialAgent: true, + }, + db, + ); + + expect(db.updateFilesUsage).toHaveBeenNthCalledWith(1, [requestFile], undefined, { + user: 'user-1', + }); + expect(db.updateFilesUsage).toHaveBeenNthCalledWith(2, [toolFile]); + }); }); describe('initializeAgent — maxContextTokens', () => { diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 122e5c3462..a6150fccb8 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -389,7 +389,11 @@ export interface InitializeAgentParams { */ export interface InitializeAgentDbMethods extends EndpointDbMethods { /** Update usage tracking for multiple files */ - updateFilesUsage: (files: Array<{ file_id: string }>, fileIds?: string[]) => Promise; + updateFilesUsage: ( + files: Array<{ file_id: string }>, + fileIds?: string[], + options?: { user?: string }, + ) => Promise; /** Get files from database */ getFiles: (filter: unknown, sort: unknown, select: unknown) => Promise; /** Filter files by agent access permissions (ownership or agent attachment) */ @@ -539,6 +543,7 @@ export async function initializeAgent( allowedProviders, isInitialAgent = false, } = params; + const requestFileOwnerId = req.user?.id; if (!db) { throw new Error('initializeAgent requires db methods to be passed'); @@ -640,10 +645,27 @@ export async function initializeAgent( const allToolFiles = toolFiles.concat(codeGeneratedFiles, userCodeFiles); if (requestFiles.length || allToolFiles.length) { - currentFiles = (await db.updateFilesUsage(requestFiles.concat(allToolFiles))) as IMongoFile[]; + const requestUsageFiles = + requestFiles.length && requestFileOwnerId + ? ((await db.updateFilesUsage(requestFiles, undefined, { + user: requestFileOwnerId, + })) as IMongoFile[]) + : []; + const requestUsageFileIds = new Set(requestUsageFiles.map((file) => file.file_id)); + const trustedToolFiles = allToolFiles.filter( + (file) => !requestUsageFileIds.has(file.file_id), + ); + const toolUsageFiles = trustedToolFiles.length + ? ((await db.updateFilesUsage(trustedToolFiles)) as IMongoFile[]) + : []; + currentFiles = requestUsageFiles.concat(toolUsageFiles); } } else if (requestFiles.length) { - currentFiles = (await db.updateFilesUsage(requestFiles)) as IMongoFile[]; + currentFiles = requestFileOwnerId + ? ((await db.updateFilesUsage(requestFiles, undefined, { + user: requestFileOwnerId, + })) as IMongoFile[]) + : []; } if (currentFiles && currentFiles.length) { diff --git a/packages/data-schemas/src/methods/file.spec.ts b/packages/data-schemas/src/methods/file.spec.ts index 08dd02326c..090783e263 100644 --- a/packages/data-schemas/src/methods/file.spec.ts +++ b/packages/data-schemas/src/methods/file.spec.ts @@ -2,10 +2,10 @@ import mongoose from 'mongoose'; import { v4 as uuidv4 } from 'uuid'; import { MongoMemoryServer } from 'mongodb-memory-server'; import { EToolResources, FileContext } from 'librechat-data-provider'; +import { _resetStrictCache } from '~/models/plugins/tenantIsolation'; +import { runAsSystem } from '~/config/tenantContext'; import { createFileMethods } from './file'; import { createModels } from '~/models'; -import { runAsSystem } from '~/config/tenantContext'; -import { _resetStrictCache } from '~/models/plugins/tenantIsolation'; let File: mongoose.Model; let fileMethods: ReturnType; @@ -660,6 +660,43 @@ describe('File Methods', () => { const updated2 = await fileMethods.updateFileUsage({ file_id: fileId, inc: 5 }); expect(updated2?.usage).toBe(6); }); + + it('should skip usage and TTL mutation when the owner filter does not match', async () => { + const fileId = uuidv4(); + const ownerId = new mongoose.Types.ObjectId(); + const otherUserId = new mongoose.Types.ObjectId(); + + await fileMethods.createFile({ + file_id: fileId, + temp_file_id: 'tmp-file', + user: ownerId, + filename: 'owned.txt', + filepath: '/uploads/owned.txt', + type: 'text/plain', + bytes: 100, + usage: 0, + }); + + const denied = await fileMethods.updateFileUsage({ + file_id: fileId, + user: otherUserId.toString(), + }); + const unchanged = await fileMethods.findFileById(fileId); + + expect(denied).toBeNull(); + expect(unchanged?.usage).toBe(0); + expect(unchanged?.temp_file_id).toBe('tmp-file'); + expect(unchanged?.expiresAt).toBeDefined(); + + const allowed = await fileMethods.updateFileUsage({ + file_id: fileId, + user: ownerId.toString(), + }); + + expect(allowed?.usage).toBe(1); + expect(allowed?.temp_file_id).toBeUndefined(); + expect(allowed?.expiresAt).toBeUndefined(); + }); }); describe('updateFilesUsage', () => { @@ -789,6 +826,48 @@ describe('File Methods', () => { expect(updated).toHaveLength(1); expect((updated[0] as { usage: number }).usage).toBe(1); }); + + it('should only return and mutate files owned by the scoped user', async () => { + const ownerId = new mongoose.Types.ObjectId(); + const otherUserId = new mongoose.Types.ObjectId(); + const ownerFileId = uuidv4(); + const otherFileId = uuidv4(); + + await fileMethods.createFile({ + file_id: ownerFileId, + user: ownerId, + filename: 'owner.txt', + filepath: '/uploads/owner.txt', + type: 'text/plain', + bytes: 100, + usage: 0, + }); + await fileMethods.createFile({ + file_id: otherFileId, + temp_file_id: 'tmp-other', + user: otherUserId, + filename: 'other.txt', + filepath: '/uploads/other.txt', + type: 'text/plain', + bytes: 100, + usage: 0, + text: 'private extracted text', + }); + + const updated = await fileMethods.updateFilesUsage( + [{ file_id: ownerFileId }, { file_id: otherFileId }], + undefined, + { user: ownerId.toString() }, + ); + const otherFile = await fileMethods.findFileById(otherFileId); + + expect(updated).toHaveLength(1); + expect(updated[0].file_id).toBe(ownerFileId); + expect(updated[0].usage).toBe(1); + expect(otherFile?.usage).toBe(0); + expect(otherFile?.temp_file_id).toBe('tmp-other'); + expect(otherFile?.expiresAt).toBeDefined(); + }); }); describe('deleteFile', () => { diff --git a/packages/data-schemas/src/methods/file.ts b/packages/data-schemas/src/methods/file.ts index a76cbc1315..29efc1bc49 100644 --- a/packages/data-schemas/src/methods/file.ts +++ b/packages/data-schemas/src/methods/file.ts @@ -1,8 +1,8 @@ -import logger from '../config/winston'; import { EToolResources, FileContext } from 'librechat-data-provider'; import type { FilterQuery, SortOrder, Model } from 'mongoose'; import type { IMongoFile } from '~/types/file'; import { tenantSafeBulkWrite } from '~/utils/tenantBulkWrite'; +import logger from '../config/winston'; /** Factory function that takes mongoose instance and returns the file methods */ export function createFileMethods(mongoose: typeof import('mongoose')) { @@ -305,14 +305,17 @@ export function createFileMethods(mongoose: typeof import('mongoose')) { async function updateFileUsage(data: { file_id: string; inc?: number; + user?: string; }): Promise { const File = mongoose.models.File as Model; - const { file_id, inc = 1 } = data; + const { file_id, inc = 1, user } = data; const updateOperation = { $inc: { usage: inc }, $unset: { expiresAt: '', temp_file_id: '' }, }; - return File.findOneAndUpdate({ file_id }, updateOperation, { + // Owner scoping is fail-closed: mismatches leave usage and TTL metadata unchanged. + const query: FilterQuery = user ? { file_id, user } : { file_id }; + return File.findOneAndUpdate(query, updateOperation, { new: true, }).lean(); } @@ -400,9 +403,12 @@ export function createFileMethods(mongoose: typeof import('mongoose')) { async function updateFilesUsage( files: Array<{ file_id: string }>, fileIds?: string[], + options?: { user?: string }, ): Promise { const promises: Promise[] = []; const seen = new Set(); + // Preserve the same owner scope for every deduped ID in this batch. + const user = options?.user; for (const file of files) { const { file_id } = file; @@ -410,7 +416,7 @@ export function createFileMethods(mongoose: typeof import('mongoose')) { continue; } seen.add(file_id); - promises.push(updateFileUsage({ file_id })); + promises.push(updateFileUsage({ file_id, user })); } if (!fileIds) { @@ -423,7 +429,7 @@ export function createFileMethods(mongoose: typeof import('mongoose')) { continue; } seen.add(file_id); - promises.push(updateFileUsage({ file_id })); + promises.push(updateFileUsage({ file_id, user })); } const results = await Promise.all(promises);