diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 2f0b576dbb..2601fb3be0 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -676,6 +676,7 @@ class AgentClient extends BaseClient { getFormattedMemories: db.getFormattedMemories, }, res: this.options.res, + user: createSafeUser(this.options.req.user), }); this.processMemory = processMemory; diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts new file mode 100644 index 0000000000..1b5242f78d --- /dev/null +++ b/packages/api/src/agents/memory.spec.ts @@ -0,0 +1,298 @@ +import { Types } from 'mongoose'; +import type { Response } from 'express'; +import { Run } from '@librechat/agents'; +import type { IUser } from '@librechat/data-schemas'; +import { createSafeUser } from '~/utils/env'; +import { processMemory } from './memory'; + +jest.mock('~/stream/GenerationJobManager'); +jest.mock('~/utils', () => ({ + Tokenizer: { + getTokenCount: jest.fn(() => 10), + }, +})); + +jest.mock('@librechat/agents', () => ({ + Run: { + create: jest.fn(() => ({ + processStream: jest.fn(() => Promise.resolve('success')), + })), + }, + Providers: { + OPENAI: 'openai', + }, + GraphEvents: { + TOOL_END: 'tool_end', + }, +})); + +function createTestUser(overrides: Partial = {}): IUser { + return { + _id: new Types.ObjectId(), + id: new Types.ObjectId().toString(), + username: 'testuser', + email: 'test@example.com', + name: 'Test User', + avatar: 'https://example.com/avatar.png', + provider: 'email', + role: 'user', + createdAt: new Date('2021-01-01'), + updatedAt: new Date('2021-01-01'), + emailVerified: true, + ...overrides, + } as IUser; +} + +describe('Memory Agent Header Resolution', () => { + let testUser: IUser; + let mockRes: Response; + let mockMemoryMethods: { + setMemory: jest.Mock; + deleteMemory: jest.Mock; + getFormattedMemories: jest.Mock; + }; + + beforeEach(() => { + process.env.CUSTOM_API_KEY = 'sk-custom-test-key'; + process.env.TEST_CUSTOM_API_KEY = 'sk-custom-test-key'; + + testUser = createTestUser({ + id: 'user-123', + email: 'test@example.com', + }); + + mockRes = { + write: jest.fn(), + end: jest.fn(), + headersSent: false, + } as unknown as Response; + + mockMemoryMethods = { + setMemory: jest.fn(), + deleteMemory: jest.fn(), + getFormattedMemories: jest.fn(() => + Promise.resolve({ + withKeys: 'formatted memories', + withoutKeys: 'memories without keys', + totalTokens: 100, + }), + ), + }; + + jest.clearAllMocks(); + }); + + afterEach(() => { + delete process.env.CUSTOM_API_KEY; + delete process.env.TEST_CUSTOM_API_KEY; + }); + + it('should resolve environment variables in custom endpoint headers', async () => { + const llmConfig = { + provider: 'custom', + model: 'gpt-4o-mini', + configuration: { + defaultHeaders: { + 'x-custom-api-key': '${CUSTOM_API_KEY}', + 'api-key': '${TEST_CUSTOM_API_KEY}', + }, + }, + }; + + await processMemory({ + res: mockRes, + userId: 'user-123', + setMemory: mockMemoryMethods.setMemory, + deleteMemory: mockMemoryMethods.deleteMemory, + messages: [], + memory: 'test memory', + messageId: 'msg-123', + conversationId: 'conv-123', + validKeys: ['preferences'], + instructions: 'test instructions', + llmConfig, + user: testUser, + }); + + expect(Run.create as jest.Mock).toHaveBeenCalled(); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0]; + expect(runConfig.graphConfig.llmConfig.configuration.defaultHeaders).toEqual({ + 'x-custom-api-key': 'sk-custom-test-key', + 'api-key': 'sk-custom-test-key', + }); + }); + + it('should resolve user placeholders in custom endpoint headers', async () => { + const llmConfig = { + provider: 'custom', + model: 'gpt-4o-mini', + configuration: { + defaultHeaders: { + 'X-User-Identifier': '{{LIBRECHAT_USER_EMAIL}}', + 'X-User-ID': '{{LIBRECHAT_USER_ID}}', + }, + }, + }; + + await processMemory({ + res: mockRes, + userId: 'user-123', + setMemory: mockMemoryMethods.setMemory, + deleteMemory: mockMemoryMethods.deleteMemory, + messages: [], + memory: 'test memory', + messageId: 'msg-123', + conversationId: 'conv-123', + validKeys: ['preferences'], + instructions: 'test instructions', + llmConfig, + user: testUser, + }); + + expect(Run.create as jest.Mock).toHaveBeenCalled(); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0]; + expect(runConfig.graphConfig.llmConfig.configuration.defaultHeaders).toEqual({ + 'X-User-Identifier': 'test@example.com', + 'X-User-ID': 'user-123', + }); + }); + + it('should handle mixed environment variables and user placeholders', async () => { + const llmConfig = { + provider: 'custom', + model: 'gpt-4o-mini', + configuration: { + defaultHeaders: { + 'x-custom-api-key': '${CUSTOM_API_KEY}', + 'X-User-Identifier': '{{LIBRECHAT_USER_EMAIL}}', + 'X-Application-Identifier': 'LibreChat - Test', + }, + }, + }; + + await processMemory({ + res: mockRes, + userId: 'user-123', + setMemory: mockMemoryMethods.setMemory, + deleteMemory: mockMemoryMethods.deleteMemory, + messages: [], + memory: 'test memory', + messageId: 'msg-123', + conversationId: 'conv-123', + validKeys: ['preferences'], + instructions: 'test instructions', + llmConfig, + user: testUser, + }); + + expect(Run.create as jest.Mock).toHaveBeenCalled(); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0]; + expect(runConfig.graphConfig.llmConfig.configuration.defaultHeaders).toEqual({ + 'x-custom-api-key': 'sk-custom-test-key', + 'X-User-Identifier': 'test@example.com', + 'X-Application-Identifier': 'LibreChat - Test', + }); + }); + + it('should resolve env vars when user is undefined', async () => { + const llmConfig = { + provider: 'custom', + model: 'gpt-4o-mini', + configuration: { + defaultHeaders: { + 'x-custom-api-key': '${CUSTOM_API_KEY}', + }, + }, + }; + + await processMemory({ + res: mockRes, + userId: 'user-123', + setMemory: mockMemoryMethods.setMemory, + deleteMemory: mockMemoryMethods.deleteMemory, + messages: [], + memory: 'test memory', + messageId: 'msg-123', + conversationId: 'conv-123', + validKeys: ['preferences'], + instructions: 'test instructions', + llmConfig, + user: undefined, + }); + + expect(Run.create as jest.Mock).toHaveBeenCalled(); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0]; + expect(runConfig.graphConfig.llmConfig.configuration.defaultHeaders).toEqual({ + 'x-custom-api-key': 'sk-custom-test-key', + }); + }); + + it('should not throw when llmConfig has no configuration', async () => { + const llmConfig = { + provider: 'openai', + model: 'gpt-4o-mini', + }; + + await processMemory({ + res: mockRes, + userId: 'user-123', + setMemory: mockMemoryMethods.setMemory, + deleteMemory: mockMemoryMethods.deleteMemory, + messages: [], + memory: 'test memory', + messageId: 'msg-123', + conversationId: 'conv-123', + validKeys: ['preferences'], + instructions: 'test instructions', + llmConfig, + user: testUser, + }); + + expect(Run.create as jest.Mock).toHaveBeenCalled(); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0]; + expect(runConfig.graphConfig.llmConfig.configuration).toBeUndefined(); + }); + + it('should use createSafeUser to sanitize user data', async () => { + const userWithSensitiveData = createTestUser({ + id: 'user-123', + email: 'test@example.com', + password: 'sensitive-password', + refreshToken: 'sensitive-token', + } as unknown as Partial); + + const llmConfig = { + provider: 'openai', + model: 'gpt-4o-mini', + configuration: { + defaultHeaders: { + 'X-User-ID': '{{LIBRECHAT_USER_ID}}', + }, + }, + }; + + await processMemory({ + res: mockRes, + userId: 'user-123', + setMemory: mockMemoryMethods.setMemory, + deleteMemory: mockMemoryMethods.deleteMemory, + messages: [], + memory: 'test memory', + messageId: 'msg-123', + conversationId: 'conv-123', + validKeys: ['preferences'], + instructions: 'test instructions', + llmConfig, + user: userWithSensitiveData, + }); + + expect(Run.create as jest.Mock).toHaveBeenCalled(); + + // Verify createSafeUser was used - the user object passed to Run.create should not have sensitive fields + const safeUser = createSafeUser(userWithSensitiveData); + expect(safeUser).not.toHaveProperty('password'); + expect(safeUser).not.toHaveProperty('refreshToken'); + expect(safeUser).toHaveProperty('id'); + expect(safeUser).toHaveProperty('email'); + }); +}); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index 2d5076381a..dcf26a8666 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -14,10 +14,11 @@ import type { LLMConfig, } from '@librechat/agents'; import type { TAttachment, MemoryArtifact } from 'librechat-data-provider'; -import type { ObjectId, MemoryMethods } from '@librechat/data-schemas'; +import type { ObjectId, MemoryMethods, IUser } from '@librechat/data-schemas'; import type { BaseMessage, ToolMessage } from '@langchain/core/messages'; import type { Response as ServerResponse } from 'express'; import { GenerationJobManager } from '~/stream/GenerationJobManager'; +import { resolveHeaders, createSafeUser } from '~/utils/env'; import { Tokenizer } from '~/utils'; type RequiredMemoryMethods = Pick< @@ -285,6 +286,7 @@ export async function processMemory({ tokenLimit, totalTokens = 0, streamId = null, + user, }: { res: ServerResponse; setMemory: MemoryMethods['setMemory']; @@ -300,6 +302,7 @@ export async function processMemory({ totalTokens?: number; llmConfig?: Partial; streamId?: string | null; + user?: IUser; }): Promise<(TAttachment | null)[] | undefined> { try { const memoryTool = createMemoryTool({ @@ -366,6 +369,14 @@ ${memory ?? 'No existing memories'}`; } } + const llmConfigWithHeaders = finalLLMConfig as OpenAIClientOptions; + if (llmConfigWithHeaders?.configuration?.defaultHeaders != null) { + llmConfigWithHeaders.configuration.defaultHeaders = resolveHeaders({ + headers: llmConfigWithHeaders.configuration.defaultHeaders as Record, + user: user ? createSafeUser(user) : undefined, + }); + } + const artifactPromises: Promise[] = []; const memoryCallback = createMemoryCallback({ res, artifactPromises, streamId }); const customHandlers = { @@ -421,6 +432,7 @@ export async function createMemoryProcessor({ conversationId, config = {}, streamId = null, + user, }: { res: ServerResponse; messageId: string; @@ -429,6 +441,7 @@ export async function createMemoryProcessor({ memoryMethods: RequiredMemoryMethods; config?: MemoryConfig; streamId?: string | null; + user?: IUser; }): Promise<[string, (messages: BaseMessage[]) => Promise<(TAttachment | null)[] | undefined>]> { const { validKeys, instructions, llmConfig, tokenLimit } = config; const finalInstructions = instructions || getDefaultInstructions(validKeys, tokenLimit); @@ -456,6 +469,7 @@ export async function createMemoryProcessor({ instructions: finalInstructions, setMemory: memoryMethods.setMemory, deleteMemory: memoryMethods.deleteMemory, + user, }); } catch (error) { logger.error('Memory Agent failed to process memory', error);