fix: address remaining Codex findings after dev rebase

Stamp agent ownership on streaming and final search snapshots. Compose the shared floating thought button. Preserve rebased group behavior, correct stale test selectors, and keep expansion callbacks outside state updaters.

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

@@ -468,6 +468,36 @@ describe('createToolEndCallback', () => {
     artifactPromises = [];
   });
 
+  it('preserves separate owners for final search artifacts with repeated tool-call IDs', async () => {
+    const toolEndCallback = createToolEndCallback({ req, res, artifactPromises });
+    for (const agentId of ['agent-a', 'agent-b']) {
+      await toolEndCallback(
+        {
+          output: {
+            tool_call_id: 'call_0',
+            artifact: {
+              [Tools.web_search]: {
+                turn: 0,
+                organic: [{ link: `https://example.com/${agentId}` }],
+              },
+            },
+          },
+        },
+        { run_id: 'run456', thread_id: 'thread789', agent_id: agentId },
+      );
+    }
+    const results = await Promise.all(artifactPromises);
+    expect(
+      results.map((attachment) => ({
+        owner: attachment.agentId,
+        link: attachment[Tools.web_search].organic[0].link,
+      })),
+    ).toEqual([
+      { owner: 'agent-a', link: 'https://example.com/agent-a' },
+      { owner: 'agent-b', link: 'https://example.com/agent-b' },
+    ]);
+  });
+
   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/callbacks.js

@@ -1006,6 +1006,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo
             type: Tools.web_search,
             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] },
           };

api/server/services/Tools/search.js

@@ -34,6 +34,7 @@ function createOnSearchResults(res, streamId = null, jobCreatedAt) {
     sourceMap: new Map(),
     searchResultData: undefined,
     toolCallId: undefined,
+    agentId: undefined,
     attachmentName: undefined,
     messageId: undefined,
     conversationId: undefined,
@@ -83,6 +84,7 @@ function createOnSearchResults(res, streamId = null, jobCreatedAt) {
     }
 
     context.toolCallId = runnableConfig.toolCall.id;
+    context.agentId = runnableConfig.metadata.agent_id;
     context.messageId = runnableConfig.metadata.run_id;
     context.conversationId = runnableConfig.metadata.thread_id;
     context.attachmentName = `${runnableConfig.toolCall.name}_${context.toolCallId}_${nanoid()}`;
@@ -132,6 +134,7 @@ function buildAttachment(context) {
   return {
     messageId: context.messageId,
     toolCallId: context.toolCallId,
+    agentId: context.agentId,
     conversationId: context.conversationId,
     name: context.attachmentName,
     type: Tools.web_search,

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

@@ -60,4 +60,35 @@ describe('createOnSearchResults', () => {
       ]);
     }
   });
+
+  it('keeps repeated search calls owned by their handoff agent through highlights', () => {
+    const snapshots = [];
+    const res = {
+      headersSent: true,
+      write: (event) => snapshots.push(JSON.parse(event.split('\ndata: ')[1])),
+    };
+    for (const agentId of ['agent-a', 'agent-b']) {
+      const callbacks = createOnSearchResults(res);
+      callbacks.onSearchResults(
+        {
+          success: true,
+          data: { organic: [{ link: `https://example.com/${agentId}` }], topStories: [] },
+        },
+        {
+          metadata: { agent_id: agentId, run_id: 'response-1', thread_id: 'conversation-1' },
+          toolCall: { id: 'call_0', name: 'web_search', turn: 0 },
+        },
+      );
+      callbacks.onGetHighlights(`https://example.com/${agentId}`);
+    }
+
+    expect(snapshots.map(({ agentId }) => agentId)).toEqual([
+      'agent-a',
+      'agent-a',
+      'agent-b',
+      'agent-b',
+    ]);
+    expect(snapshots[1].web_search.organic[0].processed).toBe(true);
+    expect(snapshots[3].web_search.organic[0].processed).toBe(true);
+  });
 });