Resolve follow-up review findings

api/server/controllers/agents/__tests__/callbacks.spec.js

@@ -8,6 +8,7 @@ jest.mock('nanoid', () => ({
 jest.mock('@librechat/api', () => ({
   sendEvent: jest.fn(),
   writeAttachmentEvent: jest.fn(),
+  createOwnedToolEndHandler: jest.fn((callback) => ({ handle: callback })),
   GenerationJobManager: {
     emitChunk: jest.fn(),
   },
@@ -498,6 +499,42 @@ describe('createToolEndCallback', () => {
     ]);
   });
 
+  it.each(['createToolEndCallback', 'createResponsesToolEndCallback'])(
+    '%s preserves memory payloads and their execution owner',
+    async (factoryName) => {
+      const toolEndCallback = require('../callbacks')[factoryName]({ req, res, artifactPromises });
+      await toolEndCallback(
+        {
+          output: {
+            tool_call_id: 'call_0',
+            artifact: {
+              [Tools.memory]: {
+                key: 'project',
+                type: 'update',
+                value: 'owned',
+              },
+            },
+          },
+        },
+        {
+          run_id: 'run456',
+          thread_id: 'thread789',
+          agent_id: 'agent-a',
+          stepId: 'step-memory-1',
+        },
+      );
+
+      const [attachment] = await Promise.all(artifactPromises);
+      expect(attachment).toMatchObject({
+        type: Tools.memory,
+        toolCallId: 'call_0',
+        agentId: 'agent-a',
+        stepId: 'step-memory-1',
+        [Tools.memory]: { key: 'project', type: 'update', value: 'owned' },
+      });
+    },
+  );
+
   describe('ui_resources artifact handling', () => {
     it('should process ui_resources artifact and return attachment when headers not sent', async () => {
       const toolEndCallback = createToolEndCallback({ req, res, artifactPromises });

api/server/controllers/agents/__tests__/openai.spec.js

@@ -194,6 +194,9 @@ jest.mock('@librechat/api', () => ({
    *  export or the call throws before the assertions run. */
   stripActivityLabelParts: jest.fn((payload) => payload),
   writeSSE: jest.fn(),
+  createOwnedToolEndHandler: jest.fn(
+    (...args) => new (require('@librechat/agents').ToolEndHandler)(...args),
+  ),
   createRun: jest.fn().mockResolvedValue({
     processStream: mockProcessStream,
   }),

api/server/controllers/agents/__tests__/responses.unit.spec.js

@@ -203,6 +203,9 @@ jest.mock('@librechat/api', () => ({
    *  before SDK formatting; the mock must expose it like any other used
    *  export or the call throws before the assertions run. */
   stripActivityLabelParts: jest.fn((payload) => payload),
+  createOwnedToolEndHandler: jest.fn(
+    (...args) => new (require('@librechat/agents').ToolEndHandler)(...args),
+  ),
   collectReachableAgents: (roots) => {
     const agents = [];
     const pending = [...roots];

api/server/controllers/agents/callbacks.js

@@ -13,7 +13,6 @@ const {
 const {
   GraphEvents,
   GraphNodeKeys,
-  ToolEndHandler,
   createContentAggregator,
   summarizeEvent,
 } = require('@librechat/agents');
@@ -23,6 +22,7 @@ const {
   GenerationJobManager,
   writeAttachmentEvent,
   createToolExecuteHandler,
+  createOwnedToolEndHandler,
   createBackgroundCodeResultHandler: createCodeHarvestHandler,
   HOST_FILE_AUTHORING_ARTIFACT_KEY,
   isCodeSessionToolName,
@@ -42,6 +42,15 @@ function isCodeArtifactToolOutput(output) {
   return isCodeSessionToolName(output.name) || isHostFileAuthoringArtifact(output.artifact);
 }
 
+function getAttachmentOwnership(metadata) {
+  const agentId = metadata?.agent_id ?? metadata?.agentId;
+  const stepId = metadata?.stepId;
+  return {
+    ...(typeof agentId === 'string' && agentId.length > 0 ? { agentId } : {}),
+    ...(typeof stepId === 'string' && stepId.length > 0 ? { stepId } : {}),
+  };
+}
+
 function addStatefulWorkspaceChange(attachment, artifact, executionProfile) {
   if (!attachment || executionProfile !== 'stateful' || !isHostFileAuthoringArtifact(artifact)) {
     return attachment;
@@ -506,7 +515,7 @@ function getDefaultHandlers({
       collectedThoughtSignatures,
       emitTokenUsage,
     ),
-    [GraphEvents.TOOL_END]: new ToolEndHandler(toolEndCallback, logger),
+    [GraphEvents.TOOL_END]: createOwnedToolEndHandler(toolEndCallback, logger),
     [GraphEvents.ON_RUN_STEP]: {
       /**
        * Handle ON_RUN_STEP event.
@@ -1004,9 +1013,9 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo
         (async () => {
           const attachment = {
             type: Tools.web_search,
+            ...getAttachmentOwnership(metadata),
             messageId: metadata.run_id,
             toolCallId: output.tool_call_id,
-            agentId: metadata.agent_id,
             conversationId: metadata.thread_id,
             [Tools.web_search]: { ...output.artifact[Tools.web_search] },
           };
@@ -1027,6 +1036,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo
         (async () => {
           const attachment = {
             type: Tools.memory,
+            ...getAttachmentOwnership(metadata),
             messageId: metadata.run_id,
             toolCallId: output.tool_call_id,
             conversationId: metadata.thread_id,
@@ -1368,6 +1378,7 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises })
           const attachment = {
             type: Tools.web_search,
             toolCallId: output.tool_call_id,
+            ...getAttachmentOwnership(metadata),
             [Tools.web_search]: { ...output.artifact[Tools.web_search] },
           };
           // For Responses API, always emit attachment during streaming
@@ -1382,6 +1393,26 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises })
       );
     }
 
+    if (output.artifact[Tools.memory]) {
+      artifactPromises.push(
+        (async () => {
+          const attachment = {
+            type: Tools.memory,
+            toolCallId: output.tool_call_id,
+            ...getAttachmentOwnership(metadata),
+            [Tools.memory]: output.artifact[Tools.memory],
+          };
+          if (res.headersSent && !res.writableEnded) {
+            writeResponsesAttachment(res, tracker, attachment, metadata);
+          }
+          return attachment;
+        })().catch((error) => {
+          logger.error('Error processing memory artifact content:', error);
+          return null;
+        }),
+      );
+    }
+
     if (output.artifact.content) {
       /** @type {FormattedContent[]} */
       const content = output.artifact.content;

api/server/controllers/agents/openai.js

@@ -1,6 +1,6 @@
 const { nanoid } = require('nanoid');
 const { logger } = require('@librechat/data-schemas');
-const { Callback, ToolEndHandler, formatAgentMessages } = require('@librechat/agents');
+const { Callback, formatAgentMessages } = require('@librechat/agents');
 const {
   EModelEndpoint,
   ResourceType,
@@ -56,6 +56,7 @@ const {
   getUserFacingProviderError,
   getRemoteAgentPermissions,
   createToolExecuteHandler,
+  createOwnedToolEndHandler,
   buildNonStreamingResponse,
   createOpenAIStreamTracker,
   resolveAgentScopedSkillIds,
@@ -1001,7 +1002,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
         },
         on_run_step_completed: createHandler(),
         // Use proper ToolEndHandler for processing artifacts (images, file citations, code output)
-        on_tool_end: new ToolEndHandler(toolEndCallback, logger),
+        on_tool_end: createOwnedToolEndHandler(toolEndCallback, logger),
         on_chain_stream: createHandler(),
         on_chain_end: createHandler(),
         on_agent_update: createHandler(),

api/server/controllers/agents/responses.js

@@ -1,7 +1,7 @@
 const { nanoid } = require('nanoid');
 const { v4: uuidv4 } = require('uuid');
 const { logger } = require('@librechat/data-schemas');
-const { Callback, ToolEndHandler, formatAgentMessages } = require('@librechat/agents');
+const { Callback, formatAgentMessages } = require('@librechat/agents');
 const {
   EModelEndpoint,
   ResourceType,
@@ -53,6 +53,7 @@ const {
   getSafeErrorMetadata,
   getUserFacingProviderError,
   createToolExecuteHandler,
+  createOwnedToolEndHandler,
   resolveRecursionLimit,
   getRemoteAgentPermissions,
   resolveAgentScopedSkillIds,
@@ -1155,7 +1156,7 @@ const executeResponse = async (envelope, { req, res }) => {
               }
             },
           },
-          on_tool_end: new ToolEndHandler(toolEndCallback, logger),
+          on_tool_end: createOwnedToolEndHandler(toolEndCallback, logger),
           on_run_step_completed: { handle: () => {} },
           on_chain_stream: { handle: () => {} },
           on_chain_end: { handle: () => {} },
@@ -1368,7 +1369,7 @@ const executeResponse = async (envelope, { req, res }) => {
               }
             },
           },
-          on_tool_end: new ToolEndHandler(toolEndCallback, logger),
+          on_tool_end: createOwnedToolEndHandler(toolEndCallback, logger),
           on_run_step_completed: { handle: () => {} },
           on_chain_stream: { handle: () => {} },
           on_chain_end: { handle: () => {} },

api/server/services/Tools/search.js

@@ -35,6 +35,7 @@ function createOnSearchResults(res, streamId = null, jobCreatedAt) {
     searchResultData: undefined,
     toolCallId: undefined,
     agentId: undefined,
+    stepId: undefined,
     attachmentName: undefined,
     messageId: undefined,
     conversationId: undefined,
@@ -84,7 +85,8 @@ function createOnSearchResults(res, streamId = null, jobCreatedAt) {
     }
 
     context.toolCallId = runnableConfig.toolCall.id;
-    context.agentId = runnableConfig.metadata.agent_id;
+    context.agentId = runnableConfig.metadata.agent_id ?? runnableConfig.metadata.agentId;
+    context.stepId = runnableConfig.toolCall?.stepId;
     context.messageId = runnableConfig.metadata.run_id;
     context.conversationId = runnableConfig.metadata.thread_id;
     context.attachmentName = `${runnableConfig.toolCall.name}_${context.toolCallId}_${nanoid()}`;
@@ -135,6 +137,7 @@ function buildAttachment(context) {
     messageId: context.messageId,
     toolCallId: context.toolCallId,
     agentId: context.agentId,
+    stepId: context.stepId,
     conversationId: context.conversationId,
     name: context.attachmentName,
     type: Tools.web_search,

api/server/services/Tools/search.spec.js

@@ -28,6 +28,7 @@ describe('createOnSearchResults', () => {
       toolCall: {
         id: 'tool-call-1',
         name: 'web_search',
+        stepId: 'step-search-1',
         turn: 0,
       },
     };
@@ -53,6 +54,7 @@ describe('createOnSearchResults', () => {
           data: expect.objectContaining({
             messageId: 'response-1',
             toolCallId: 'tool-call-1',
+            stepId: 'step-search-1',
             conversationId: 'conversation-1',
           }),
         },