diff --git a/api/server/controllers/agents/__tests__/responses.unit.spec.js b/api/server/controllers/agents/__tests__/responses.unit.spec.js index 3e980a057e..4bc6d19e7e 100644 --- a/api/server/controllers/agents/__tests__/responses.unit.spec.js +++ b/api/server/controllers/agents/__tests__/responses.unit.spec.js @@ -54,6 +54,9 @@ const mockCanAuthorSkillFiles = jest.fn( scopedEditableSkillIds.length > 0 || skillCreateAllowed === true, ); const mockGetSkillToolDeps = jest.fn(() => ({})); +const mockBuildAgentScopedContext = jest.fn().mockResolvedValue(new Map()); +const mockBuildAgentContextAttachmentsByAgentId = jest.fn().mockReturnValue(new Map()); +const mockApplyContextToAgent = jest.fn().mockResolvedValue(undefined); jest.mock('nanoid', () => ({ nanoid: jest.fn(() => 'mock-nanoid-123'), @@ -84,7 +87,11 @@ jest.mock('@librechat/api', () => ({ createRun: jest.fn().mockResolvedValue({ processStream: jest.fn().mockResolvedValue(undefined), }), + applyContextToAgent: (...args) => mockApplyContextToAgent(...args), buildToolSet: jest.fn().mockReturnValue(new Set()), + buildAgentScopedContext: (...args) => mockBuildAgentScopedContext(...args), + buildAgentContextAttachmentsByAgentId: (...args) => + mockBuildAgentContextAttachmentsByAgentId(...args), scopeSkillIds: jest.fn().mockImplementation((ids) => ids), resolveAgentScopedSkillIds: jest .fn() @@ -97,12 +104,21 @@ jest.mock('@librechat/api', () => ({ model_parameters: {}, toolRegistry: {}, edges: [], + agentContextAttachments: [], }), - discoverConnectedAgents: jest.fn().mockResolvedValue({ - agentConfigs: new Map(), - edges: [], - skippedAgentIds: new Set(), - userMCPAuthMap: undefined, + discoverConnectedAgents: jest.fn().mockImplementation(async (computedParams, deps) => { + // Call onAgentInitialized for each agent config if provided by the mock setup + if (deps?.onAgentInitialized && mockGlobalDiscoveredAgentConfigs) { + for (const [agentId, config] of mockGlobalDiscoveredAgentConfigs) { + deps.onAgentInitialized(agentId, config, config); + } + } + return { + agentConfigs: mockGlobalDiscoveredAgentConfigs ?? new Map(), + edges: [], + skippedAgentIds: new Set(), + userMCPAuthMap: undefined, + }; }), getBalanceConfig: mockGetBalanceConfig, getTransactionsConfig: mockGetTransactionsConfig, @@ -196,6 +212,14 @@ jest.mock('~/server/controllers/ModelController', () => ({ getModelsConfig: jest.fn().mockResolvedValue({}), })); +jest.mock('~/server/services/MCP', () => ({ + resolveConfigServers: jest.fn().mockResolvedValue({}), +})); + +jest.mock('~/config', () => ({ + getMCPManager: jest.fn().mockReturnValue({}), +})); + jest.mock('~/server/services/Files/permissions', () => ({ filterFilesByAgentAccess: jest.fn(), })); @@ -253,12 +277,15 @@ jest.mock('~/models', () => ({ getConvo: jest.fn().mockResolvedValue(null), })); +let mockGlobalDiscoveredAgentConfigs = null; + describe('createResponse controller', () => { let createResponse; let req, res; beforeEach(() => { jest.clearAllMocks(); + mockGlobalDiscoveredAgentConfigs = null; const controller = require('../responses'); createResponse = controller.createResponse; @@ -449,6 +476,89 @@ describe('createResponse controller', () => { }); }); + describe('agent context parity with UI path', () => { + it('applies agent-scoped attachment context before createRun', async () => { + const api = require('@librechat/api'); + api.initializeAgent.mockResolvedValueOnce({ + id: 'agent-123', + model: 'claude-3', + model_parameters: {}, + toolRegistry: {}, + edges: [], + agentContextAttachments: [{ file_id: 'file-1', filename: 'ocr_file.pdf' }], + }); + mockBuildAgentContextAttachmentsByAgentId.mockReturnValueOnce( + new Map([['agent-123', [{ file_id: 'file-1', filename: 'ocr_file.pdf' }]]]), + ); + mockBuildAgentScopedContext.mockResolvedValueOnce( + new Map([['agent-123', 'PDF context: ocr_file.pdf']]), + ); + + await createResponse(req, res); + + expect(mockBuildAgentContextAttachmentsByAgentId).toHaveBeenCalledWith([ + expect.objectContaining({ id: 'agent-123' }), + ]); + expect(mockBuildAgentScopedContext).toHaveBeenCalledWith( + expect.objectContaining({ + agentIds: ['agent-123'], + attachmentsByAgentId: expect.any(Map), + req, + }), + ); + expect(mockApplyContextToAgent).toHaveBeenCalledWith( + expect.objectContaining({ + agent: expect.objectContaining({ id: 'agent-123' }), + agentId: 'agent-123', + sharedRunContext: 'PDF context: ocr_file.pdf', + }), + ); + }); + + it('applies context to primary and discovered handoff agents', async () => { + const api = require('@librechat/api'); + const handoffConfig = { + id: 'agent-handoff', + model: 'claude-3', + model_parameters: {}, + toolRegistry: {}, + edges: [], + agentContextAttachments: [{ file_id: 'file-2', filename: 'handoff_context.pdf' }], + }; + + // Set primary agent to have edges pointing to handoff agent + api.initializeAgent.mockResolvedValueOnce({ + id: 'agent-123', + model: 'claude-3', + model_parameters: {}, + toolRegistry: {}, + edges: [{ source: 'agent-123', target: 'agent-handoff' }], + agentContextAttachments: [{ file_id: 'file-1', filename: 'primary_context.pdf' }], + }); + + // Set global config so discoverConnectedAgents mock can invoke onAgentInitialized + mockGlobalDiscoveredAgentConfigs = new Map([['agent-handoff', handoffConfig]]); + + mockBuildAgentScopedContext.mockResolvedValueOnce( + new Map([ + ['agent-123', 'Primary context'], + ['agent-handoff', 'Handoff context'], + ]), + ); + + await createResponse(req, res); + + const appliedAgentIds = mockApplyContextToAgent.mock.calls.map((call) => call[0].agentId); + expect(appliedAgentIds).toEqual(expect.arrayContaining(['agent-123', 'agent-handoff'])); + expect(mockApplyContextToAgent).toHaveBeenCalledWith( + expect.objectContaining({ + agentId: 'agent-handoff', + sharedRunContext: 'Handoff context', + }), + ); + }); + }); + describe('token usage recording - streaming', () => { beforeEach(() => { req.body.stream = true; diff --git a/api/server/controllers/agents/responses.js b/api/server/controllers/agents/responses.js index 212b245a49..c88545c3e6 100644 --- a/api/server/controllers/agents/responses.js +++ b/api/server/controllers/agents/responses.js @@ -11,7 +11,10 @@ const { } = require('librechat-data-provider'); const { createRun, + applyContextToAgent, buildToolSet, + buildAgentScopedContext, + buildAgentContextAttachmentsByAgentId, createSafeUser, initializeAgent, loadSkillStates, @@ -65,6 +68,8 @@ const { enrichLoadedToolsWithAgentContext, } = require('~/server/services/Endpoints/agents/skillDeps'); const { getModelsConfig } = require('~/server/controllers/ModelController'); +const { resolveConfigServers } = require('~/server/services/MCP'); +const { getMCPManager } = require('~/config'); const { logViolation } = require('~/cache'); const db = require('~/models'); @@ -562,6 +567,29 @@ const createResponse = async (req, res) => { const runAgents = [primaryConfig, ...handoffAgentConfigs.values()]; const mergedMCPAuthMap = discoveredMCPAuthMap ?? primaryConfig.userMCPAuthMap; + const agentContextAttachmentsByAgentId = buildAgentContextAttachmentsByAgentId(runAgents); + const agentScopedContext = await buildAgentScopedContext({ + agentIds: runAgents.map(({ id }) => id), + attachmentsByAgentId: agentContextAttachmentsByAgentId, + req, + }); + + const mcpManager = getMCPManager(); + const configServers = await resolveConfigServers(req); + + await Promise.all( + runAgents.map((runAgent) => + applyContextToAgent({ + agent: runAgent, + agentId: runAgent.id, + logger, + mcpManager, + configServers, + sharedRunContext: agentScopedContext.get(runAgent.id) ?? '', + }), + ), + ); + // Determine if streaming is enabled (check both request and agent config) const streamingDisabled = !!primaryConfig.model_parameters?.disableStreaming; const actuallyStreaming = isStreaming && !streamingDisabled;