From c0cb48256ec2dde1699c7abf3c43ce876fe17894 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 17 Nov 2025 16:57:51 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20refactor:=20Improve=20Agent=20Ha?= =?UTF-8?q?ndoff=20Context=20Tracking=20(#10553)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: update @librechat/agents dependency to version 3.0.18 * refactor: add optional metadata field to message schema and types * chore: update @librechat/agents to v3.0.19 * refactor: update return type of sendCompletion method to include metadata * chore: linting * chore: update @librechat/agents dependency to v3.0.20 * refactor: implement agent labeling for conversation history in multi-agent scenarios * refactor: improve error handling for capturing agent ID map in AgentClient * refactor: clear agentIdMap and related properties during client disposal to prevent memory leaks * chore: update sendCompletion method for FakeClient to return an object with completion and metadata fields --- api/app/clients/BaseClient.js | 5 +- .../prompts/formatAgentMessages.spec.js | 20 ++-- api/app/clients/specs/FakeClient.js | 5 +- api/package.json | 2 +- api/server/cleanup.js | 3 + api/server/controllers/agents/client.js | 91 ++++++++++++++++++- api/typedefs.js | 2 +- package-lock.json | 10 +- packages/api/package.json | 2 +- packages/data-provider/src/schemas.ts | 2 + packages/data-schemas/src/schema/message.ts | 1 + packages/data-schemas/src/types/message.ts | 1 + 12 files changed, 122 insertions(+), 22 deletions(-) diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 185e1c964f..149d331df1 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -81,6 +81,7 @@ class BaseClient { throw new Error("Method 'getCompletion' must be implemented."); } + /** @type {sendCompletion} */ async sendCompletion() { throw new Error("Method 'sendCompletion' must be implemented."); } @@ -689,8 +690,7 @@ class BaseClient { }); } - /** @type {string|string[]|undefined} */ - const completion = await this.sendCompletion(payload, opts); + const { completion, metadata } = await this.sendCompletion(payload, opts); if (this.abortController) { this.abortController.requestCompleted = true; } @@ -708,6 +708,7 @@ class BaseClient { iconURL: this.options.iconURL, endpoint: this.options.endpoint, ...(this.metadata ?? {}), + metadata, }; if (typeof completion === 'string') { diff --git a/api/app/clients/prompts/formatAgentMessages.spec.js b/api/app/clients/prompts/formatAgentMessages.spec.js index 360fa00a34..1aee3edf71 100644 --- a/api/app/clients/prompts/formatAgentMessages.spec.js +++ b/api/app/clients/prompts/formatAgentMessages.spec.js @@ -130,7 +130,7 @@ describe('formatAgentMessages', () => { content: [ { type: ContentTypes.TEXT, - [ContentTypes.TEXT]: 'I\'ll search for that information.', + [ContentTypes.TEXT]: "I'll search for that information.", tool_call_ids: ['search_1'], }, { @@ -144,7 +144,7 @@ describe('formatAgentMessages', () => { }, { type: ContentTypes.TEXT, - [ContentTypes.TEXT]: 'Now, I\'ll convert the temperature.', + [ContentTypes.TEXT]: "Now, I'll convert the temperature.", tool_call_ids: ['convert_1'], }, { @@ -156,7 +156,7 @@ describe('formatAgentMessages', () => { output: '23.89°C', }, }, - { type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Here\'s your answer.' }, + { type: ContentTypes.TEXT, [ContentTypes.TEXT]: "Here's your answer." }, ], }, ]; @@ -171,7 +171,7 @@ describe('formatAgentMessages', () => { expect(result[4]).toBeInstanceOf(AIMessage); // Check first AIMessage - expect(result[0].content).toBe('I\'ll search for that information.'); + expect(result[0].content).toBe("I'll search for that information."); expect(result[0].tool_calls).toHaveLength(1); expect(result[0].tool_calls[0]).toEqual({ id: 'search_1', @@ -187,7 +187,7 @@ describe('formatAgentMessages', () => { ); // Check second AIMessage - expect(result[2].content).toBe('Now, I\'ll convert the temperature.'); + expect(result[2].content).toBe("Now, I'll convert the temperature."); expect(result[2].tool_calls).toHaveLength(1); expect(result[2].tool_calls[0]).toEqual({ id: 'convert_1', @@ -202,7 +202,7 @@ describe('formatAgentMessages', () => { // Check final AIMessage expect(result[4].content).toStrictEqual([ - { [ContentTypes.TEXT]: 'Here\'s your answer.', type: ContentTypes.TEXT }, + { [ContentTypes.TEXT]: "Here's your answer.", type: ContentTypes.TEXT }, ]); }); @@ -217,7 +217,7 @@ describe('formatAgentMessages', () => { role: 'assistant', content: [{ type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'How can I help you?' }], }, - { role: 'user', content: 'What\'s the weather?' }, + { role: 'user', content: "What's the weather?" }, { role: 'assistant', content: [ @@ -240,7 +240,7 @@ describe('formatAgentMessages', () => { { role: 'assistant', content: [ - { type: ContentTypes.TEXT, [ContentTypes.TEXT]: 'Here\'s the weather information.' }, + { type: ContentTypes.TEXT, [ContentTypes.TEXT]: "Here's the weather information." }, ], }, ]; @@ -265,12 +265,12 @@ describe('formatAgentMessages', () => { { [ContentTypes.TEXT]: 'How can I help you?', type: ContentTypes.TEXT }, ]); expect(result[2].content).toStrictEqual([ - { [ContentTypes.TEXT]: 'What\'s the weather?', type: ContentTypes.TEXT }, + { [ContentTypes.TEXT]: "What's the weather?", type: ContentTypes.TEXT }, ]); expect(result[3].content).toBe('Let me check that for you.'); expect(result[4].content).toBe('Sunny, 75°F'); expect(result[5].content).toStrictEqual([ - { [ContentTypes.TEXT]: 'Here\'s the weather information.', type: ContentTypes.TEXT }, + { [ContentTypes.TEXT]: "Here's the weather information.", type: ContentTypes.TEXT }, ]); // Check that there are no consecutive AIMessages diff --git a/api/app/clients/specs/FakeClient.js b/api/app/clients/specs/FakeClient.js index 8c79847069..d1d07a967d 100644 --- a/api/app/clients/specs/FakeClient.js +++ b/api/app/clients/specs/FakeClient.js @@ -82,7 +82,10 @@ const initializeFakeClient = (apiKey, options, fakeMessages) => { }); TestClient.sendCompletion = jest.fn(async () => { - return 'Mock response text'; + return { + completion: 'Mock response text', + metadata: undefined, + }; }); TestClient.getCompletion = jest.fn().mockImplementation(async (..._args) => { diff --git a/api/package.json b/api/package.json index 6b47e6656c..e64165348b 100644 --- a/api/package.json +++ b/api/package.json @@ -47,7 +47,7 @@ "@langchain/google-genai": "^0.2.13", "@langchain/google-vertexai": "^0.2.13", "@langchain/textsplitters": "^0.1.0", - "@librechat/agents": "^3.0.17", + "@librechat/agents": "^3.0.20", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", diff --git a/api/server/cleanup.js b/api/server/cleanup.js index c482a2267e..8e19c853ea 100644 --- a/api/server/cleanup.js +++ b/api/server/cleanup.js @@ -350,6 +350,9 @@ function disposeClient(client) { if (client.agentConfigs) { client.agentConfigs = null; } + if (client.agentIdMap) { + client.agentIdMap = null; + } if (client.artifactPromises) { client.artifactPromises = null; } diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index d76dc4bb6e..5dea281bb0 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -20,6 +20,7 @@ const { Providers, TitleMethod, formatMessage, + labelContentByAgent, formatAgentMessages, getTokenCountForMessage, createMetadataAggregator, @@ -92,6 +93,61 @@ function logToolError(graph, error, toolId) { }); } +/** + * Applies agent labeling to conversation history when multi-agent patterns are detected. + * Labels content parts by their originating agent to prevent identity confusion. + * + * @param {TMessage[]} orderedMessages - The ordered conversation messages + * @param {Agent} primaryAgent - The primary agent configuration + * @param {Map} agentConfigs - Map of additional agent configurations + * @returns {TMessage[]} Messages with agent labels applied where appropriate + */ +function applyAgentLabelsToHistory(orderedMessages, primaryAgent, agentConfigs) { + const shouldLabelByAgent = (primaryAgent.edges?.length ?? 0) > 0 || (agentConfigs?.size ?? 0) > 0; + + if (!shouldLabelByAgent) { + return orderedMessages; + } + + const processedMessages = []; + + for (let i = 0; i < orderedMessages.length; i++) { + const message = orderedMessages[i]; + + /** @type {Record} */ + const agentNames = { [primaryAgent.id]: primaryAgent.name || 'Assistant' }; + + if (agentConfigs) { + for (const [agentId, agentConfig] of agentConfigs.entries()) { + agentNames[agentId] = agentConfig.name || agentConfig.id; + } + } + + if ( + !message.isCreatedByUser && + message.metadata?.agentIdMap && + Array.isArray(message.content) + ) { + try { + const labeledContent = labelContentByAgent( + message.content, + message.metadata.agentIdMap, + agentNames, + ); + + processedMessages.push({ ...message, content: labeledContent }); + } catch (error) { + logger.error('[AgentClient] Error applying agent labels to message:', error); + processedMessages.push(message); + } + } else { + processedMessages.push(message); + } + } + + return processedMessages; +} + class AgentClient extends BaseClient { constructor(options = {}) { super(null, options); @@ -141,6 +197,8 @@ class AgentClient extends BaseClient { this.indexTokenCountMap = {}; /** @type {(messages: BaseMessage[]) => Promise} */ this.processMemory; + /** @type {Record | null} */ + this.agentIdMap = null; } /** @@ -233,6 +291,12 @@ class AgentClient extends BaseClient { summary: this.shouldSummarize, }); + orderedMessages = applyAgentLabelsToHistory( + orderedMessages, + this.options.agent, + this.agentConfigs, + ); + let payload; /** @type {number | undefined} */ let promptTokens; @@ -612,7 +676,11 @@ class AgentClient extends BaseClient { userMCPAuthMap: opts.userMCPAuthMap, abortController: opts.abortController, }); - return filterMalformedContentParts(this.contentParts); + + const completion = filterMalformedContentParts(this.contentParts); + const metadata = this.agentIdMap ? { agentIdMap: this.agentIdMap } : undefined; + + return { completion, metadata }; } /** @@ -902,6 +970,24 @@ class AgentClient extends BaseClient { ); }); } + + try { + /** Capture agent ID map if we have edges or multiple agents */ + const shouldStoreAgentMap = + (this.options.agent.edges?.length ?? 0) > 0 || (this.agentConfigs?.size ?? 0) > 0; + if (shouldStoreAgentMap && run?.Graph) { + const contentPartAgentMap = run.Graph.getContentPartAgentMap(); + if (contentPartAgentMap && contentPartAgentMap.size > 0) { + this.agentIdMap = Object.fromEntries(contentPartAgentMap); + logger.debug('[AgentClient] Captured agent ID map:', { + totalParts: this.contentParts.length, + mappedParts: Object.keys(this.agentIdMap).length, + }); + } + } + } catch (error) { + logger.error('[AgentClient] Error capturing agent ID map:', error); + } } catch (err) { logger.error( '[api/server/controllers/agents/client.js #sendCompletion] Operation aborted', @@ -935,6 +1021,9 @@ class AgentClient extends BaseClient { err, ); } + run = null; + config = null; + memoryPromise = null; } } diff --git a/api/typedefs.js b/api/typedefs.js index a75ab7d8c1..b6385c69a9 100644 --- a/api/typedefs.js +++ b/api/typedefs.js @@ -1828,7 +1828,7 @@ * @param {onTokenProgress} opts.onProgress - Callback function to handle token progress * @param {AbortController} opts.abortController - AbortController instance * @param {Record>} [opts.userMCPAuthMap] - * @returns {Promise} + * @returns {Promise<{ content: Promise; metadata: Record; }>} * @memberof typedefs */ diff --git a/package-lock.json b/package-lock.json index e6d3853bd3..3bc7520a70 100644 --- a/package-lock.json +++ b/package-lock.json @@ -63,7 +63,7 @@ "@langchain/google-genai": "^0.2.13", "@langchain/google-vertexai": "^0.2.13", "@langchain/textsplitters": "^0.1.0", - "@librechat/agents": "^3.0.17", + "@librechat/agents": "^3.0.20", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", @@ -16351,9 +16351,9 @@ } }, "node_modules/@librechat/agents": { - "version": "3.0.17", - "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.0.17.tgz", - "integrity": "sha512-gnom77oW10cdGmmQ6rExc+Blrfzib9JqrZk2fDPwkSOWXQIK4nMm6vWS+UkhH/YfN996mMnffpWoQQ6QQvkTGg==", + "version": "3.0.20", + "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.0.20.tgz", + "integrity": "sha512-GSFpTYIylN/01c4QksMlsKBkptB4v9qYHzcT7rXxHzj568W7mz0ZPBfs1N6PwZRyp+1RgDwo3v38CS4oQmdOTQ==", "license": "MIT", "dependencies": { "@langchain/anthropic": "^0.3.26", @@ -47334,7 +47334,7 @@ "@azure/storage-blob": "^12.27.0", "@keyv/redis": "^4.3.3", "@langchain/core": "^0.3.79", - "@librechat/agents": "^3.0.17", + "@librechat/agents": "^3.0.20", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.21.0", "axios": "^1.12.1", diff --git a/packages/api/package.json b/packages/api/package.json index 44fe3963ac..1197aa9881 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -84,7 +84,7 @@ "@azure/storage-blob": "^12.27.0", "@keyv/redis": "^4.3.3", "@langchain/core": "^0.3.79", - "@librechat/agents": "^3.0.17", + "@librechat/agents": "^3.0.20", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.21.0", "axios": "^1.12.1", diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 9d1761d468..afc07880da 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -610,6 +610,8 @@ export const tMessageSchema = z.object({ /* frontend components */ iconURL: z.string().nullable().optional(), feedback: feedbackSchema.optional(), + /** metadata */ + metadata: z.record(z.unknown()).optional(), }); export type MemoryArtifact = { diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index c11252cb87..f287f14ca4 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -132,6 +132,7 @@ const messageSchema: Schema = new Schema( iconURL: { type: String, }, + metadata: { type: mongoose.Schema.Types.Mixed }, attachments: { type: [{ type: mongoose.Schema.Types.Mixed }], default: undefined }, /* attachments: { diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index 2743f7242e..f69bcff6b9 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -37,6 +37,7 @@ export interface IMessage extends Document { content?: unknown[]; thread_id?: string; iconURL?: string; + metadata?: Record; attachments?: unknown[]; expiredAt?: Date; createdAt?: Date;