diff --git a/client/src/Providers/BadgeRowContext.tsx b/client/src/Providers/BadgeRowContext.tsx index 17e8f3b6e8..0f9647f8f8 100644 --- a/client/src/Providers/BadgeRowContext.tsx +++ b/client/src/Providers/BadgeRowContext.tsx @@ -270,7 +270,12 @@ export default function BadgeRowProvider({ isAuthenticated: true, }); - const mcpServerManager = useMCPServerManager({ conversationId, storageContextKey }); + const mcpServerManager = useMCPServerManager({ + conversationId, + storageContextKey, + specName, + ownsChatSelection: true, + }); const value: BadgeRowContextType = { skills, diff --git a/client/src/components/Chat/Input/MCPSelect.tsx b/client/src/components/Chat/Input/MCPSelect.tsx index ff9f6007d1..ec5f8e450c 100644 --- a/client/src/components/Chat/Input/MCPSelect.tsx +++ b/client/src/components/Chat/Input/MCPSelect.tsx @@ -23,22 +23,23 @@ function MCPSelectContent() { return []; } const selectedSet = new Set(manager.mcpValues); - return manager.selectableServers?.filter((s) => selectedSet.has(s.serverName)); + return manager.selectableServers?.filter((s) => selectedSet.has(s.serverName)) ?? []; }, [manager?.selectableServers, manager?.mcpValues]); + /** Counts what the menu actually offers, never the raw selection: a name the + * catalog has not returned — or one the admin has hidden — renders no row, + * and billing it to the badge reads as a server that cannot be turned off. */ const displayText = useMemo(() => { - const selectedCount = manager?.mcpValues?.length ?? 0; + const selectedCount = selectedServers.length; if (selectedCount === 0) { return null; } if (selectedCount === 1) { - const server = manager?.selectableServers?.find( - (s) => s.serverName === manager?.mcpValues?.[0], - ); - return server?.config?.title || manager?.mcpValues?.[0]; + const server = selectedServers[0]; + return server.config?.title || server.serverName; } return localize('com_ui_x_selected', { 0: selectedCount }); - }, [manager?.selectableServers, manager?.mcpValues, localize]); + }, [selectedServers, localize]); if (!manager) { return null; diff --git a/client/src/components/Chat/Input/__tests__/MCPSelect.spec.tsx b/client/src/components/Chat/Input/__tests__/MCPSelect.spec.tsx index 421b8cf473..28152fc853 100644 --- a/client/src/components/Chat/Input/__tests__/MCPSelect.spec.tsx +++ b/client/src/components/Chat/Input/__tests__/MCPSelect.spec.tsx @@ -126,6 +126,28 @@ describe('MCPSelect', () => { expect(items[0]).toHaveFocus(); }); + it('labels the badge from the servers the menu offers, not the raw selection', () => { + mockMcpServerManager = { + ...defaultMcpServerManager, + mcpValues: ['server-a', 'hidden-server'], + }; + render(); + + expect(screen.getByRole('button', { name: /Server A/i })).toBeInTheDocument(); + expect(screen.queryByText('com_ui_x_selected')).not.toBeInTheDocument(); + }); + + it('shows no label when every selected server is hidden from the menu', () => { + mockMcpServerManager = { + ...defaultMcpServerManager, + mcpValues: ['hidden-server'], + }; + render(); + + expect(screen.getByRole('button', { name: /MCP Servers/i })).toBeInTheDocument(); + expect(screen.queryByText('hidden-server')).not.toBeInTheDocument(); + }); + it('renders nothing when user lacks MCP access', () => { mockCanUseMcp = false; const { container } = render(); diff --git a/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx b/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx index bbdace485e..93ca85f519 100644 --- a/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx +++ b/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx @@ -16,7 +16,12 @@ jest.mock('~/utils/timestamps', () => ({ jest.mock('lodash/isEqual', () => jest.fn((a, b) => JSON.stringify(a) === JSON.stringify(b))); // Mutable startup config so tests can vary `interface.defaultPinnedTools` -let mockStartupConfig: { interface?: { defaultPinnedTools?: string[] } } | undefined; +let mockStartupConfig: + | { + interface?: { defaultPinnedTools?: string[] }; + modelSpecs?: { list?: Array<{ name: string; mcpServers?: string[] }> }; + } + | undefined; jest.mock('~/data-provider', () => ({ ...jest.requireActual('~/data-provider'), @@ -58,7 +63,7 @@ describe('useMCPSelect', () => { describe('Basic Functionality', () => { it('should initialize with default values', () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -71,18 +76,24 @@ describe('useMCPSelect', () => { it('should use conversationId when provided', () => { const conversationId = 'test-convo-123'; const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ conversationId, servers }), { - wrapper: Wrapper, - }); + const { result } = renderHook( + () => useMCPSelect({ conversationId, servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); expect(result.current.mcpValues).toEqual([]); }); it('should use NEW_CONVO constant when conversationId is null', () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ conversationId: null, servers }), { - wrapper: Wrapper, - }); + const { result } = renderHook( + () => useMCPSelect({ conversationId: null, servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); expect(result.current.mcpValues).toEqual([]); }); @@ -91,7 +102,7 @@ describe('useMCPSelect', () => { describe('State Updates', () => { it('should update mcpValues when setMCPValues is called', async () => { const { Wrapper, servers } = createWrapper(['value1', 'value2']); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -108,7 +119,7 @@ describe('useMCPSelect', () => { it('should not update mcpValues if non-array is passed', () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -122,7 +133,7 @@ describe('useMCPSelect', () => { it('should update isPinned state', () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -149,9 +160,12 @@ describe('useMCPSelect', () => { it('should set timestamp when mcpValues is updated with values', async () => { const conversationId = 'test-convo'; const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ conversationId, servers }), { - wrapper: Wrapper, - }); + const { result } = renderHook( + () => useMCPSelect({ conversationId, servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); const newValues = ['value1', 'value2']; @@ -167,7 +181,7 @@ describe('useMCPSelect', () => { it('should not set timestamp when mcpValues is empty', async () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -184,9 +198,12 @@ describe('useMCPSelect', () => { describe('Race Conditions and Infinite Loops Prevention', () => { it('should not create infinite loop when syncing between Jotai and Recoil states', async () => { const { Wrapper, servers } = createWrapper(); - const { result, rerender } = renderHook(() => useMCPSelect({ servers }), { - wrapper: Wrapper, - }); + const { result, rerender } = renderHook( + () => useMCPSelect({ servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); let renderCount = 0; const maxRenders = 10; @@ -217,7 +234,7 @@ describe('useMCPSelect', () => { it('should handle rapid consecutive updates without race conditions', async () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -244,9 +261,12 @@ describe('useMCPSelect', () => { it('should maintain stable setter function reference', () => { const { Wrapper, servers } = createWrapper(); - const { result, rerender } = renderHook(() => useMCPSelect({ servers }), { - wrapper: Wrapper, - }); + const { result, rerender } = renderHook( + () => useMCPSelect({ servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); const firstSetMCPValues = result.current.setMCPValues; @@ -262,7 +282,7 @@ describe('useMCPSelect', () => { it('should handle switching conversation IDs without issues', async () => { const { Wrapper, servers } = createWrapper(['convo1-value', 'convo2-value']); const { result, rerender } = renderHook( - ({ conversationId }) => useMCPSelect({ conversationId, servers }), + ({ conversationId }) => useMCPSelect({ conversationId, servers, ownsChatSelection: true }), { wrapper: Wrapper, initialProps: { conversationId: 'convo1' }, @@ -310,7 +330,7 @@ describe('useMCPSelect', () => { // Create a component that uses both hooks to ensure they share state const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const [ephemeralAgent, setEphemeralAgent] = useRecoilState( ephemeralAgentByConvoId(Constants.NEW_CONVO), ); @@ -337,7 +357,7 @@ describe('useMCPSelect', () => { const { Wrapper, servers } = createWrapper(['server1', 'server2']); const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -359,7 +379,7 @@ describe('useMCPSelect', () => { const { Wrapper, servers } = createWrapper(['server1', 'server2']); const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -381,7 +401,7 @@ describe('useMCPSelect', () => { const { Wrapper, servers } = createWrapper(['server1', 'server2', 'server3']); const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -405,7 +425,7 @@ describe('useMCPSelect', () => { // Create a component that uses both the hook and accesses Recoil state const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, ephemeralAgent }; }; @@ -431,7 +451,7 @@ describe('useMCPSelect', () => { // Create a component that uses both hooks const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -467,7 +487,7 @@ describe('useMCPSelect', () => { // Create a component that uses both hooks const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -506,7 +526,7 @@ describe('useMCPSelect', () => { ]); const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -544,9 +564,12 @@ describe('useMCPSelect', () => { describe('Edge Cases', () => { it('should handle undefined conversationId', () => { const { Wrapper, servers } = createWrapper(['test']); - const { result } = renderHook(() => useMCPSelect({ conversationId: undefined, servers }), { - wrapper: Wrapper, - }); + const { result } = renderHook( + () => useMCPSelect({ conversationId: undefined, servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); expect(result.current.mcpValues).toEqual([]); @@ -559,9 +582,12 @@ describe('useMCPSelect', () => { it('should handle empty string conversationId', () => { const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ conversationId: '', servers }), { - wrapper: Wrapper, - }); + const { result } = renderHook( + () => useMCPSelect({ conversationId: '', servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); expect(result.current.mcpValues).toEqual([]); }); @@ -569,7 +595,7 @@ describe('useMCPSelect', () => { it('should handle very large arrays without performance issues', async () => { const largeArray = Array.from({ length: 1000 }, (_, i) => `value-${i}`); const { Wrapper, servers } = createWrapper(largeArray); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -592,7 +618,7 @@ describe('useMCPSelect', () => { it('should cleanup properly on unmount', () => { const { Wrapper, servers } = createWrapper(); - const { unmount } = renderHook(() => useMCPSelect({ servers }), { + const { unmount } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -608,7 +634,13 @@ describe('useMCPSelect', () => { // Hook A: new conversation with storageContextKey const { result: resultA } = renderHook( - () => useMCPSelect({ conversationId: null, storageContextKey, servers }), + () => + useMCPSelect({ + conversationId: null, + storageContextKey, + servers, + ownsChatSelection: true, + }), { wrapper: Wrapper }, ); @@ -622,7 +654,7 @@ describe('useMCPSelect', () => { // Hook B: new conversation WITHOUT storageContextKey (different environment) const { result: resultB } = renderHook( - () => useMCPSelect({ conversationId: null, servers }), + () => useMCPSelect({ conversationId: null, servers, ownsChatSelection: true }), { wrapper: Wrapper }, ); @@ -636,7 +668,7 @@ describe('useMCPSelect', () => { const storageContextKey = '__defaults__'; const { result } = renderHook( - () => useMCPSelect({ conversationId, storageContextKey, servers }), + () => useMCPSelect({ conversationId, storageContextKey, servers, ownsChatSelection: true }), { wrapper: Wrapper }, ); @@ -658,7 +690,13 @@ describe('useMCPSelect', () => { const storageContextKey = '__defaults__'; const { result } = renderHook( - () => useMCPSelect({ conversationId: null, storageContextKey, servers }), + () => + useMCPSelect({ + conversationId: null, + storageContextKey, + servers, + ownsChatSelection: true, + }), { wrapper: Wrapper }, ); @@ -678,9 +716,12 @@ describe('useMCPSelect', () => { const conversationId = 'convo-no-specs'; const { Wrapper, servers } = createWrapper(['server1']); - const { result } = renderHook(() => useMCPSelect({ conversationId, servers }), { - wrapper: Wrapper, - }); + const { result } = renderHook( + () => useMCPSelect({ conversationId, servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); act(() => { result.current.setMCPValues(['server1']); @@ -701,7 +742,13 @@ describe('useMCPSelect', () => { // Set environment defaults via new conversation const { result: newConvoResult } = renderHook( - () => useMCPSelect({ conversationId: null, storageContextKey, servers }), + () => + useMCPSelect({ + conversationId: null, + storageContextKey, + servers, + ownsChatSelection: true, + }), { wrapper: Wrapper }, ); @@ -715,7 +762,13 @@ describe('useMCPSelect', () => { // Existing conversation should have its own isolated state const { result: existingResult } = renderHook( - () => useMCPSelect({ conversationId: 'existing-convo', storageContextKey, servers }), + () => + useMCPSelect({ + conversationId: 'existing-convo', + storageContextKey, + servers, + ownsChatSelection: true, + }), { wrapper: Wrapper }, ); @@ -742,7 +795,12 @@ describe('useMCPSelect', () => { const storageContextKey = '__defaults__'; const TestComponent = ({ ctxKey }: { ctxKey?: string }) => { - const mcpHook = useMCPSelect({ conversationId: null, storageContextKey: ctxKey, servers }); + const mcpHook = useMCPSelect({ + conversationId: null, + storageContextKey: ctxKey, + servers, + ownsChatSelection: true, + }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -775,7 +833,7 @@ describe('useMCPSelect', () => { const { Wrapper, servers } = createWrapper(['spec-server1', 'spec-server2']); const TestComponent = () => { - const mcpHook = useMCPSelect({ conversationId: null, servers }); + const mcpHook = useMCPSelect({ conversationId: null, servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -799,7 +857,7 @@ describe('useMCPSelect', () => { const { Wrapper, servers } = createWrapper(['server1', 'server2']); const TestComponent = () => { - const mcpHook = useMCPSelect({ servers }); + const mcpHook = useMCPSelect({ servers, ownsChatSelection: true }); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); return { mcpHook, setEphemeralAgent }; }; @@ -828,11 +886,202 @@ describe('useMCPSelect', () => { }); }); + describe('Chat-hidden servers (chatMenu: false)', () => { + it('clears a persisted selection once the catalog reports every server hidden', async () => { + const { Wrapper } = createWrapper(); + const storageKey = `${LocalStorageKeys.LAST_MCP_}${Constants.NEW_CONVO}`; + localStorage.setItem(storageKey, JSON.stringify(['hidden-server'])); + + const { result } = renderHook( + () => + useMCPSelect({ + servers: [], + allServers: createMCPServers(['hidden-server']), + ownsChatSelection: true, + }), + { wrapper: Wrapper }, + ); + + await waitFor(() => { + expect(result.current.mcpValues).toEqual([]); + }); + }); + + it('clears the ephemeral agent too, so the request stops carrying the server', async () => { + const { Wrapper } = createWrapper(); + + const TestComponent = () => { + const mcpHook = useMCPSelect({ + ownsChatSelection: true, + servers: createMCPServers(['visible']), + allServers: createMCPServers(['visible', 'hidden-server']), + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { mcpHook, ephemeralAgent, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ mcp: ['visible', 'hidden-server'] }); + }); + + await waitFor(() => { + expect(result.current.ephemeralAgent?.mcp).toEqual(['visible']); + expect(result.current.mcpHook.mcpValues).toEqual(['visible']); + }); + }); + + it('keeps a server a model spec pins, even though the menu hides it', async () => { + mockStartupConfig = { + modelSpecs: { list: [{ name: 'pins-hidden', mcpServers: ['spec-server'] }] }, + }; + const { Wrapper } = createWrapper(); + + const TestComponent = () => { + const mcpHook = useMCPSelect({ + ownsChatSelection: true, + servers: createMCPServers(['visible']), + allServers: createMCPServers(['visible', 'spec-server', 'stale-hidden']), + specName: 'pins-hidden', + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { mcpHook, ephemeralAgent, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ + mcp: ['visible', 'spec-server', 'stale-hidden'], + }); + }); + + await waitFor(() => { + expect(result.current.ephemeralAgent?.mcp).toEqual(['visible', 'spec-server']); + expect(result.current.mcpHook.mcpValues).toEqual(['visible', 'spec-server']); + }); + }); + + it("lets a catalog-only instance leave the picker instance's selection alone", async () => { + mockStartupConfig = { + modelSpecs: { list: [{ name: 'pins-hidden', mcpServers: ['spec-server'] }] }, + }; + const { Wrapper } = createWrapper(); + + /** Mirrors `useSideNavLinks`, which is always mounted and shares the + * NEW_CONVO selection but knows nothing about the active spec. */ + const TestComponent = () => { + const picker = useMCPSelect({ + servers: createMCPServers(['visible']), + allServers: createMCPServers(['visible', 'spec-server']), + specName: 'pins-hidden', + ownsChatSelection: true, + }); + useMCPSelect({ + servers: createMCPServers(['visible']), + allServers: createMCPServers(['visible', 'spec-server']), + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { picker, ephemeralAgent, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ mcp: ['visible', 'spec-server'] }); + }); + + await waitFor(() => { + expect(result.current.picker.mcpValues).toEqual(['visible', 'spec-server']); + }); + expect(result.current.ephemeralAgent?.mcp).toEqual(['visible', 'spec-server']); + }); + + it('mirrors the selection into a non-owner instance so its actions build on it', async () => { + const { Wrapper } = createWrapper(); + + /** Mirrors `MCPServerCard`, which mounts the manager for its actions only; + * those actions derive the next selection from `mcpValues`. */ + const TestComponent = () => { + const actions = useMCPSelect({ servers: createMCPServers(['server1', 'server2']) }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { actions, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ mcp: ['server1'] }); + }); + + await waitFor(() => { + expect(result.current.actions.mcpValues).toEqual(['server1']); + }); + + act(() => { + result.current.actions.setMCPValues([...result.current.actions.mcpValues, 'server2']); + }); + + await waitFor(() => { + expect(result.current.actions.mcpValues).toEqual(['server1', 'server2']); + }); + }); + + it('keeps a selection while the catalog is empty, so a degraded read cannot wipe it', async () => { + const { Wrapper } = createWrapper(); + const storageKey = `${LocalStorageKeys.LAST_MCP_}${Constants.NEW_CONVO}`; + localStorage.setItem(storageKey, JSON.stringify(['server1', 'server2'])); + + const { result, rerender } = renderHook( + ({ servers, allServers }) => useMCPSelect({ servers, allServers, ownsChatSelection: true }), + { + initialProps: { + servers: [] as MCPServerDefinition[], + allServers: [] as MCPServerDefinition[], + }, + wrapper: Wrapper, + }, + ); + + await waitFor(() => { + expect(result.current.mcpValues).toEqual(['server1', 'server2']); + }); + + rerender({ + servers: createMCPServers(['server1', 'server2']), + allServers: createMCPServers(['server1', 'server2']), + }); + + await waitFor(() => { + expect(result.current.mcpValues).toEqual(['server1', 'server2']); + }); + }); + + it('falls back to the selectable list when no unfiltered catalog is passed', async () => { + const { Wrapper } = createWrapper(); + const storageKey = `${LocalStorageKeys.LAST_MCP_}${Constants.NEW_CONVO}`; + localStorage.setItem(storageKey, JSON.stringify(['server1', 'stale'])); + + const { result } = renderHook( + () => useMCPSelect({ servers: createMCPServers(['server1']), ownsChatSelection: true }), + { wrapper: Wrapper }, + ); + + await waitFor(() => { + expect(result.current.mcpValues).toEqual(['server1']); + }); + }); + }); + describe('Memory Leak Prevention', () => { it('should not leak memory on repeated updates', async () => { const values = Array.from({ length: 100 }, (_, i) => `value-${i}`); const { Wrapper, servers } = createWrapper(values); - const { result } = renderHook(() => useMCPSelect({ servers }), { + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { wrapper: Wrapper, }); @@ -849,9 +1098,12 @@ describe('useMCPSelect', () => { it('should handle component remounting', () => { const { Wrapper, servers } = createWrapper(); - const { result, unmount } = renderHook(() => useMCPSelect({ servers }), { - wrapper: Wrapper, - }); + const { result, unmount } = renderHook( + () => useMCPSelect({ servers, ownsChatSelection: true }), + { + wrapper: Wrapper, + }, + ); act(() => { result.current.setMCPValues(['before-unmount']); @@ -861,9 +1113,12 @@ describe('useMCPSelect', () => { // Remount const { Wrapper: Wrapper2, servers: servers2 } = createWrapper(); - const { result: newResult } = renderHook(() => useMCPSelect({ servers: servers2 }), { - wrapper: Wrapper2, - }); + const { result: newResult } = renderHook( + () => useMCPSelect({ servers: servers2, ownsChatSelection: true }), + { + wrapper: Wrapper2, + }, + ); // Should handle remounting gracefully expect(newResult.current.mcpValues).toBeDefined(); @@ -874,7 +1129,9 @@ describe('useMCPSelect', () => { it('keeps the MCP dropdown pinned by default when "mcp" is listed', async () => { mockStartupConfig = { interface: { defaultPinnedTools: ['artifacts', 'mcp'] } }; const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { wrapper: Wrapper }); + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { + wrapper: Wrapper, + }); await waitFor(() => { expect(result.current.isPinned).toBe(true); @@ -884,7 +1141,9 @@ describe('useMCPSelect', () => { it('unpins the MCP dropdown by default when configured without "mcp"', async () => { mockStartupConfig = { interface: { defaultPinnedTools: ['artifacts'] } }; const { Wrapper, servers } = createWrapper(['serverA']); - const { result } = renderHook(() => useMCPSelect({ servers }), { wrapper: Wrapper }); + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { + wrapper: Wrapper, + }); await waitFor(() => { expect(result.current.isPinned).toBe(false); @@ -894,7 +1153,9 @@ describe('useMCPSelect', () => { it('pins the MCP dropdown when a configured server name is listed', async () => { mockStartupConfig = { interface: { defaultPinnedTools: ['serverA'] } }; const { Wrapper, servers } = createWrapper(['serverA', 'serverB']); - const { result } = renderHook(() => useMCPSelect({ servers }), { wrapper: Wrapper }); + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { + wrapper: Wrapper, + }); await waitFor(() => { expect(result.current.isPinned).toBe(true); @@ -904,7 +1165,9 @@ describe('useMCPSelect', () => { it('keeps the legacy pinned default when defaultPinnedTools is not configured', async () => { mockStartupConfig = { interface: {} }; const { Wrapper, servers } = createWrapper(['serverA']); - const { result } = renderHook(() => useMCPSelect({ servers }), { wrapper: Wrapper }); + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { + wrapper: Wrapper, + }); await waitFor(() => { expect(result.current.isPinned).toBe(true); @@ -915,7 +1178,9 @@ describe('useMCPSelect', () => { localStorage.setItem(LocalStorageKeys.PIN_MCP_, JSON.stringify(false)); mockStartupConfig = { interface: { defaultPinnedTools: ['mcp'] } }; const { Wrapper, servers } = createWrapper(); - const { result } = renderHook(() => useMCPSelect({ servers }), { wrapper: Wrapper }); + const { result } = renderHook(() => useMCPSelect({ servers, ownsChatSelection: true }), { + wrapper: Wrapper, + }); await waitFor(() => { expect(result.current.isPinned).toBe(false); diff --git a/client/src/hooks/MCP/useMCPSelect.ts b/client/src/hooks/MCP/useMCPSelect.ts index 1a5fd97b0a..f4e08efaee 100644 --- a/client/src/hooks/MCP/useMCPSelect.ts +++ b/client/src/hooks/MCP/useMCPSelect.ts @@ -3,10 +3,11 @@ import { useAtom } from 'jotai'; import isEqual from 'lodash/isEqual'; import { useRecoilState } from 'recoil'; import { Constants, LocalStorageKeys } from 'librechat-data-provider'; +import type { MCPServerDefinition } from './useMCPServerManager'; import { ephemeralAgentByConvoId, mcpValuesAtomFamily, mcpPinnedAtom } from '~/store'; -import { MCPServerDefinition } from './useMCPServerManager'; import { useGetStartupConfig } from '~/data-provider'; import { setTimestamp } from '~/utils/timestamps'; +import { getModelSpec } from '~/utils'; /** Sentinel in `interface.defaultPinnedTools` that pins the MCP dropdown to the prompt bar. */ const MCP_PIN_KEYWORD = 'mcp'; @@ -15,15 +16,62 @@ export function useMCPSelect({ conversationId, storageContextKey, servers, + allServers, + specName, + ownsChatSelection = false, }: { conversationId?: string | null; storageContextKey?: string; + /** Chat-selectable servers, i.e. the subset the dropdown offers. */ servers: MCPServerDefinition[]; + /** Every server the catalog returned, selectable or not. Defaults to `servers`. */ + allServers?: MCPServerDefinition[]; + /** Active model spec, whose pinned servers are exempt from pruning. */ + specName?: string | null; + /** + * Whether this instance drives the chat picker and may therefore rewrite the + * shared selection. Off by default: every instance keyed to a conversation + * shares one atom and one ephemeral agent, so a caller mounted for the catalog + * alone — with no spec context to exempt from — would otherwise prune away a + * selection the picker's own instance is deliberately keeping. + */ + ownsChatSelection?: boolean; }) { const key = conversationId ?? Constants.NEW_CONVO; const configuredServers = useMemo(() => { return new Set(servers?.map((s) => s.serverName)); }, [servers]); + /** + * Whether the catalog has told us enough to forget a stale selection. + * + * Gating on the UNFILTERED list is what makes a `chatMenu: false` server + * clearable: when every configured server is hidden, `servers` is empty and a + * guard on it can never fire, so a selection persisted while the server was + * visible stays active forever. The unfiltered list is non-empty in that case. + * + * It stays a guard rather than a load flag because `mcpValues` writes through + * to localStorage: an empty catalog — still loading, or a degraded read — must + * never be read as "the admin removed everything" and wipe the selection. + */ + const canPruneSelections = ownsChatSelection && (allServers ?? servers).length > 0; + const { data: startupConfig } = useGetStartupConfig(); + /** + * Selections that survive pruning: what the dropdown offers, plus whatever the + * active model spec pins. `chatMenu` hides a server from the picker; it does + * not override an admin's spec, so a spec-assigned server stays selected even + * when the picker would never have offered it. + */ + const retainedServers = useMemo(() => { + const specServers = getModelSpec({ specName, startupConfig })?.mcpServers; + if (!specServers?.length) { + return configuredServers; + } + const retained = new Set(configuredServers); + for (const serverName of specServers) { + retained.add(serverName); + } + return retained; + }, [configuredServers, specName, startupConfig]); /** * For new conversations, key the MCP atom by environment (spec or defaults) @@ -33,7 +81,6 @@ export function useMCPSelect({ const isNewConvo = key === Constants.NEW_CONVO; const mcpAtomKey = isNewConvo && storageContextKey ? storageContextKey : key; - const { data: startupConfig } = useGetStartupConfig(); const [isPinned, setIsPinned] = useAtom(mcpPinnedAtom); const [mcpValues, setMCPValuesRaw] = useAtom(mcpValuesAtomFamily(mcpAtomKey)); const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key)); @@ -71,11 +118,41 @@ export function useMCPSelect({ } }, [startupConfig, servers, isPinned, setIsPinned]); - // Sync ephemeral agent MCP → Jotai atom (strip unconfigured servers) + /** Drop persisted selections the chat menu no longer offers. Runs on the atom + * itself so a selection survives with no ephemeral agent to sync from. */ + useEffect(() => { + if (!canPruneSelections || mcpValues.length === 0) { + return; + } + const activeMcpValues = mcpValues.filter((mcp) => retainedServers.has(mcp)); + if (activeMcpValues.length !== mcpValues.length) { + setMCPValuesRaw(activeMcpValues); + } + }, [canPruneSelections, mcpValues, retainedServers, setMCPValuesRaw]); + + /** + * Mirror the ephemeral agent's MCP list into this instance's atom. + * + * Every instance mirrors, owner or not: the action paths (`initializeServer`, + * the revoke handler) build their next selection from `mcpValues`, so an + * instance left unmirrored would write a stale list back through + * `setMCPValues` — authenticating one server would drop another. Only the + * pruning inside is owner-gated. + */ useEffect(() => { const mcps = ephemeralAgent?.mcp; - if (Array.isArray(mcps) && mcps.length > 0 && configuredServers.size > 0) { - const activeMcps = mcps.filter((mcp) => configuredServers.has(mcp)); + if (Array.isArray(mcps) && mcps.length > 0) { + const activeMcps = canPruneSelections ? mcps.filter((mcp) => retainedServers.has(mcp)) : mcps; + /** The ephemeral agent is what carries the selection to the server, so a + * hidden name has to leave it too, not just the dropdown's atom. */ + if (activeMcps.length !== mcps.length) { + setEphemeralAgent((prev) => { + if (!Array.isArray(prev?.mcp) || isEqual(prev.mcp, activeMcps)) { + return prev; + } + return { ...prev, mcp: activeMcps }; + }); + } if (!isEqual(activeMcps, mcpValues)) { setMCPValuesRaw(activeMcps); } @@ -83,7 +160,14 @@ export function useMCPSelect({ // Ephemeral agent explicitly has empty MCP (e.g., spec with no MCP servers) — clear atom setMCPValuesRaw([]); } - }, [ephemeralAgent?.mcp, setMCPValuesRaw, configuredServers, mcpValues]); + }, [ + ephemeralAgent?.mcp, + setEphemeralAgent, + setMCPValuesRaw, + retainedServers, + canPruneSelections, + mcpValues, + ]); // Write timestamp when MCP values change useEffect(() => { diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index e72186c64c..6afa424c30 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -62,7 +62,21 @@ type PollIntervals = Record; export function useMCPServerManager({ conversationId, storageContextKey, -}: { conversationId?: string | null; storageContextKey?: string } = {}) { + specName, + ownsChatSelection = false, +}: { + conversationId?: string | null; + storageContextKey?: string; + specName?: string | null; + /** + * Opt in to managing the chat MCP selection. Most callers mount this hook for + * the catalog, the server actions, or the status icons and never read the + * selection, so it defaults off: every instance keyed to a conversation shares + * one selection, and only the one rendering the picker knows the spec context + * needed to prune it correctly. + */ + ownsChatSelection?: boolean; +} = {}) { const localize = useLocalize(); const queryClient = useQueryClient(); const { showToast } = useToastContext(); @@ -121,6 +135,9 @@ export function useMCPServerManager({ conversationId, storageContextKey, servers: selectableServers, + allServers: availableMCPServers, + specName, + ownsChatSelection, }); const mcpValuesRef = useRef(mcpValues);