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

client/src/components/Chat/Messages/Content/Parts/Reasoning.tsx

@@ -5,7 +5,13 @@ import { Lightbulb, ChevronDown } from 'lucide-react';
 import { ContentTypes } from 'librechat-data-provider';
 import { Button, disclosureChevronVariants } from '@librechat/client';
 import type { MouseEvent, FocusEvent } from 'react';
-import { ThinkingContent, ThinkingButton, ThinkingLabel, FloatingThinkingBar, useInViewport } from './Thinking';
+import {
+  ThinkingContent,
+  ThinkingButton,
+  ThinkingLabel,
+  FloatingThinkingBar,
+  useInViewport,
+} from './Thinking';
 import { useLocalize, useExpandCollapse, useLazyCollapseBody } from '~/hooks';
 import useSmoothStreaming from '~/hooks/Messages/useSmoothStreaming';
 import CopyButton from '~/components/Messages/Content/CopyButton';

client/src/components/Chat/Messages/Content/Parts/Thinking.tsx

@@ -9,7 +9,7 @@ import {
   type MouseEvent,
 } from 'react';
 import { useAtomValue } from 'jotai';
-import { Lightbulb, ChevronDown, ChevronUp } from 'lucide-react';
+import { Lightbulb, ChevronDown } from 'lucide-react';
 import { Button, MorphIcon, TooltipAnchor } from '@librechat/client';
 import { ChevronUp as ChevronUpNode, ChevronDown as ChevronDownNode } from 'lucide';
 import type { FocusEvent, FC } from 'react';
@@ -255,24 +255,21 @@ export const FloatingThinkingBar = memo(
         <TooltipAnchor
           description={collapseTooltip}
           render={
-            <button
+            <Button
+              variant="ghost"
+              size="icon-sm"
               type="button"
               tabIndex={isVisible ? 0 : -1}
               onClick={onClick}
               aria-label={collapseTooltip}
               aria-expanded={isExpanded}
               aria-controls={contentId}
-              className={cn(
-                'flex items-center justify-center rounded-lg p-1.5 text-text-secondary',
-                'hover:bg-surface-hover hover:text-text-primary',
-                'focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-border-heavy',
-              )}
             >
               <MorphIcon
                 icon={isExpanded ? ChevronUpNode : ChevronDownNode}
                 className="h-[18px] w-[18px]"
               />
-            </button>
+            </Button>
           }
         />
         {content && (

client/src/components/Chat/Messages/Content/Parts/__tests__/ThinkingButton.test.tsx

@@ -1,45 +0,0 @@
-import React from 'react';
-import { render, screen } from '@testing-library/react';
-import { ThinkingButton } from '../Thinking';
-
-jest.mock('jotai', () => ({
-  ...jest.requireActual('jotai'),
-  useAtomValue: () => 'text-base',
-}));
-
-jest.mock('~/hooks', () => ({
-  useLocalize: () => (key: string) => key,
-}));
-
-jest.mock('~/components/Messages/Content/CopyButton', () => ({
-  __esModule: true,
-  default: ({ className }: { className?: string }) => (
-    <button type="button" data-testid="copy-thoughts" className={className} />
-  ),
-}));
-
-describe('ThinkingButton', () => {
-  const props = {
-    onClick: jest.fn(),
-    label: 'Thoughts',
-    content: 'A useful thought',
-    contentId: 'thoughts-content',
-  };
-
-  it('keeps the expanded copy control out of the header layout flow', () => {
-    const { rerender } = render(<ThinkingButton {...props} isExpanded={false} />);
-
-    const toggle = screen.getByRole('button', { name: 'Thoughts' });
-    expect(toggle.parentElement).toHaveClass('relative');
-    expect(screen.queryByTestId('copy-thoughts')).not.toBeInTheDocument();
-
-    rerender(<ThinkingButton {...props} isExpanded />);
-
-    expect(screen.getByTestId('copy-thoughts')).toHaveClass(
-      'absolute',
-      'right-0',
-      'top-1/2',
-      '-translate-y-1/2',
-    );
-  });
-});

client/src/components/Chat/Messages/Content/ToolCall.tsx

@@ -20,7 +20,6 @@ import { AttachmentGroup } from './Parts';
 import ToolCallInfo from './ToolCallInfo';
 import ProgressText from './ProgressText';
 import { TOOL_ROW_CLASSES } from './rows';
-import { logger } from '~/utils';
 import { ToolAuthWarning } from './auth';
 import store from '~/store';
 
@@ -216,14 +215,11 @@ export default function ToolCall({
 
   const handleToggleInfo = useCallback(() => {
     mountBody();
-    setShowInfo((prev) => {
-      const next = !prev;
-      if (next) {
-        onExpand?.();
-      }
-      return next;
-    });
-  }, [mountBody, onExpand]);
+    if (!showInfo) {
+      onExpand?.();
+    }
+    setShowInfo((prev) => !prev);
+  }, [mountBody, onExpand, showInfo]);
 
   const subtitle = useMemo(() => {
     if (isMCPToolCall && mcpServerName) {

client/src/components/Chat/Messages/Content/__tests__/ContentParts.integration.test.tsx

@@ -348,7 +348,7 @@ describe('ContentParts integration: MCP image hoist and grouping', () => {
       </RecoilRoot>,
     );
 
-    const toggles = screen.getAllByRole('button', { name: 'Used 2 tools' });
+    const toggles = screen.getAllByRole('button', { name: /Everything/ });
     fireEvent.click(toggles[1]);
     expect(toggles[0]).toHaveAttribute('aria-expanded', 'false');
     expect(toggles[1]).toHaveAttribute('aria-expanded', 'true');
@@ -358,7 +358,7 @@ describe('ContentParts integration: MCP image hoist and grouping', () => {
         <ContentParts {...baseProps} content={nextContent} />
       </RecoilRoot>,
     );
-    const shiftedToggles = screen.getAllByRole('button', { name: 'Used 2 tools' });
+    const shiftedToggles = screen.getAllByRole('button', { name: /Everything/ });
     expect(shiftedToggles[0]).toHaveAttribute('aria-expanded', 'false');
     expect(shiftedToggles[1]).toHaveAttribute('aria-expanded', 'true');
   });
@@ -389,7 +389,7 @@ describe('ContentParts integration: MCP image hoist and grouping', () => {
 
     fireEvent.click(screen.getByRole('button', { name: 'First phase' }));
     fireEvent.click(screen.getByRole('button', { name: 'Second phase' }));
-    const toggles = screen.getAllByRole('button', { name: 'Used 2 tools' });
+    const toggles = screen.getAllByRole('button', { name: /Everything/ });
     fireEvent.click(toggles[1]);
 
     rerender(
@@ -403,7 +403,7 @@ describe('ContentParts integration: MCP image hoist and grouping', () => {
         fireEvent.click(phaseToggle);
       }
     }
-    const shiftedToggles = screen.getAllByRole('button', { name: 'Used 2 tools' });
+    const shiftedToggles = screen.getAllByRole('button', { name: /Everything/ });
     expect(shiftedToggles[0]).toHaveAttribute('aria-expanded', 'false');
     expect(shiftedToggles[1]).toHaveAttribute('aria-expanded', 'true');
   });
@@ -849,7 +849,7 @@ describe('ContentParts — synthesized activity folds', () => {
       content: [makeMcpToolCall('t1'), makeMcpToolCall('t2')],
     });
 
-    expect(screen.getByRole('button', { name: 'Used 2 tools' })).toBeInTheDocument();
+    expect(screen.getByRole('button')).toHaveAttribute('aria-expanded', 'false');
     expect(screen.queryByTestId('activity-phase-panel')).toBeNull();
   });
 });
@@ -975,7 +975,10 @@ describe('ContentParts integration: lane groups backed by one agent', () => {
 
     renderContentParts({ ...baseProps, content });
 
-    expect(screen.getByRole('button', { name: 'Used 2 tools' })).toBeInTheDocument();
+    const toggle = screen.getByRole('button');
+    expect(toggle).toHaveAttribute('aria-expanded', 'false');
+    fireEvent.click(toggle);
+    expect(toggle).toHaveAttribute('aria-expanded', 'true');
   });
 
   it('still renders columns once a second agent shares the group', () => {

client/src/components/Chat/Messages/Content/__tests__/ToolCallGroup.test.tsx

@@ -440,10 +440,7 @@ describe('ToolCallGroup image hoisting', () => {
       withinActivityPhase: true,
     });
 
-    expect(screen.getByRole('button', { name: 'Used 2 tools' })).toHaveAttribute(
-      'aria-expanded',
-      'true',
-    );
+    expect(screen.getByRole('button')).toHaveAttribute('aria-expanded', 'true');
     expect(screen.getByTestId('inner-0')).toBeInTheDocument();
     expect(screen.getByTestId('inner-1')).toBeInTheDocument();
   });
@@ -451,7 +448,7 @@ describe('ToolCallGroup image hoisting', () => {
   it('still expands on user toggle inside a completed phase', () => {
     renderGroup({ ...baseProps, withinActivityPhase: true });
 
-    fireEvent.click(screen.getByRole('button', { name: 'Used 2 tools' }));
+    fireEvent.click(screen.getByRole('button'));
 
     expect(screen.getByTestId('inner-0')).toBeInTheDocument();
     expect(screen.getByTestId('inner-1')).toBeInTheDocument();