🤖 refactor: Improve Agent Handoff Context Tracking (#10553)

* 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

@@ -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') {

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

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) => {

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",

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;
     }

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<string, Agent>} 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<string, string>} */
+    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<void>} */
     this.processMemory;
+    /** @type {Record<number, string> | 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;
     }
   }
 

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<string, Record<string, string>>} [opts.userMCPAuthMap]
- * @returns {Promise<string>}
+ * @returns {Promise<{ content: Promise<MessageContentComplex[]>; metadata: Record<string, unknown>; }>}
  * @memberof typedefs
  */