@google/gemini-cli-core 0.35.0-preview.2 → 0.36.0-nightly.20260318.e2658ccda

package/dist/src/tools/mcp-client.test.js

@@ -3,6 +3,7 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
+/* eslint-disable @typescript-eslint/no-explicit-any */
 import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js';
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js';
@@ -112,9 +113,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedClient.listTools).toHaveBeenCalledWith({}, expect.objectContaining({ timeout: 600000, progressReporter: client }));
         });
         it('should not skip tools even if a parameter is missing a type', async () => {
@@ -177,9 +182,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
             expect(consoleWarnSpy).not.toHaveBeenCalled();
             consoleWarnSpy.mockRestore();
@@ -214,9 +223,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await expect(client.discover(MOCK_CONTEXT)).rejects.toThrow('Test error');
+            await expect(client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            })).rejects.toThrow('Test error');
             expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith('error', `Error discovering prompts from test-server: Test error`, expect.any(Error), 'test-server');
         });
         it('should not discover tools if server does not support them', async () => {
@@ -252,9 +265,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await expect(client.discover(MOCK_CONTEXT)).rejects.toThrow('No prompts, tools, or resources found on the server.');
+            await expect(client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            })).rejects.toThrow('No prompts, tools, or resources found on the server.');
         });
         it('should discover tools if server supports them', async () => {
             const mockedClient = {
@@ -298,9 +315,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
         });
         it('should register tool with readOnlyHint and preserve annotations', async () => {
@@ -348,9 +369,13 @@ describe('mcp-client', () => {
                 setResourcesForServer: vi.fn(),
                 removeResourcesByServer: vi.fn(),
             };
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(mockConfig);
+            await client.discoverInto(mockConfig, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             // Verify tool registration
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
             // Verify addRule is NOT called (annotation-based rules are in plan.toml now)
@@ -405,9 +430,13 @@ describe('mcp-client', () => {
                 setResourcesForServer: vi.fn(),
                 removeResourcesByServer: vi.fn(),
             };
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(mockConfig);
+            await client.discoverInto(mockConfig, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
             expect(mockPolicyEngine.addRule).not.toHaveBeenCalled();
             // Verify annotations are undefined for tools without annotations
@@ -460,9 +489,13 @@ describe('mcp-client', () => {
                 setResourcesForServer: vi.fn(),
                 removeResourcesByServer: vi.fn(),
             };
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(mockConfig);
+            await client.discoverInto(mockConfig, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
             const registeredTool = mockedToolRegistry.registerTool.mock.calls[0][0];
             expect(registeredTool.toolAnnotations).toEqual({
@@ -529,9 +562,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
             const registeredTool = vi.mocked(mockedToolRegistry.registerTool).mock
                 .calls[0][0];
@@ -600,9 +637,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith('test-server', [
                 expect.objectContaining({
                     uri: 'file:///tmp/resource.txt',
@@ -671,9 +712,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2);
             expect(resourceListHandler).toBeDefined();
             await resourceListHandler?.({
@@ -733,9 +778,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry,
+                resourceRegistry,
+            });
             expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2);
             expect(promptListHandler).toBeDefined();
             await promptListHandler?.({
@@ -790,9 +839,13 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, mockedPromptRegistry, resourceRegistry, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
-            await client.discover(MOCK_CONTEXT);
+            await client.discoverInto(MOCK_CONTEXT, {
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: mockedPromptRegistry,
+                resourceRegistry,
+            });
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
             expect(mockedPromptRegistry.registerPrompt).toHaveBeenCalledOnce();
             await client.disconnect();
@@ -827,16 +880,22 @@ describe('mcp-client', () => {
                 getToolsByServer: vi.fn().mockReturnValue([]),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {
-                getPromptsByServer: vi.fn().mockReturnValue([]),
-                registerPrompt: vi.fn(),
-            }, {
-                getResourcesByServer: vi.fn().mockReturnValue([]),
-                registerResource: vi.fn(),
-                removeResourcesByServer: vi.fn(),
-                setResourcesForServer: vi.fn(),
-            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {
+                    getPromptsByServer: vi.fn().mockReturnValue([]),
+                    registerPrompt: vi.fn(),
+                },
+                resourceRegistry: {
+                    getResourcesByServer: vi.fn().mockReturnValue([]),
+                    registerResource: vi.fn(),
+                    removeResourcesByServer: vi.fn(),
+                    setResourcesForServer: vi.fn(),
+                },
+            });
             expect(mockedClient.setNotificationHandler).toHaveBeenCalledWith(ToolListChangedNotificationSchema, expect.any(Function));
         });
         it('should NOT set up notification handler if server lacks capability', async () => {
@@ -850,20 +909,26 @@ describe('mcp-client', () => {
             };
             vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
             vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
-            const client = new McpClient('test-server', { command: 'test-command' }, {
-                getToolsByServer: vi.fn().mockReturnValue([]),
-                registerTool: vi.fn(),
-                sortTools: vi.fn(),
-            }, {
-                getPromptsByServer: vi.fn().mockReturnValue([]),
-                registerPrompt: vi.fn(),
-            }, {
-                getResourcesByServer: vi.fn().mockReturnValue([]),
-                registerResource: vi.fn(),
-                removeResourcesByServer: vi.fn(),
-                setResourcesForServer: vi.fn(),
-            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: {
+                    getToolsByServer: vi.fn().mockReturnValue([]),
+                    registerTool: vi.fn(),
+                    sortTools: vi.fn(),
+                },
+                promptRegistry: {
+                    getPromptsByServer: vi.fn().mockReturnValue([]),
+                    registerPrompt: vi.fn(),
+                },
+                resourceRegistry: {
+                    getResourcesByServer: vi.fn().mockReturnValue([]),
+                    registerResource: vi.fn(),
+                    removeResourcesByServer: vi.fn(),
+                    setResourcesForServer: vi.fn(),
+                },
+            });
             // Should be called for ProgressNotificationSchema, even if no other capabilities
             expect(mockedClient.setNotificationHandler).toHaveBeenCalled();
             const progressCall = mockedClient.setNotificationHandler.mock.calls.find((call) => call[0] === ProgressNotificationSchema);
@@ -882,20 +947,26 @@ describe('mcp-client', () => {
                 setRequestHandler: vi.fn().mockResolvedValue({}),
             };
             vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
-            const client = new McpClient('test-server', { command: 'test-command' }, {
-                getToolsByServer: vi.fn().mockReturnValue([]),
-                registerTool: vi.fn(),
-                sortTools: vi.fn(),
-            }, {
-                getPromptsByServer: vi.fn().mockReturnValue([]),
-                registerPrompt: vi.fn(),
-            }, {
-                getResourcesByServer: vi.fn().mockReturnValue([]),
-                registerResource: vi.fn(),
-                removeResourcesByServer: vi.fn(),
-                setResourcesForServer: vi.fn(),
-            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: {
+                    getToolsByServer: vi.fn().mockReturnValue([]),
+                    registerTool: vi.fn(),
+                    sortTools: vi.fn(),
+                },
+                promptRegistry: {
+                    getPromptsByServer: vi.fn().mockReturnValue([]),
+                    registerPrompt: vi.fn(),
+                },
+                resourceRegistry: {
+                    getResourcesByServer: vi.fn().mockReturnValue([]),
+                    registerResource: vi.fn(),
+                    removeResourcesByServer: vi.fn(),
+                    setResourcesForServer: vi.fn(),
+                },
+            });
             const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
             expect(toolUpdateCall).toBeDefined();
         });
@@ -932,12 +1003,18 @@ describe('mcp-client', () => {
             };
             const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);
             // Initialize client with onContextUpdated callback
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {
-                removeMcpResourcesByServer: vi.fn(),
-                registerResource: vi.fn(),
-            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
             // 1. Connect (sets up listener)
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {},
+                resourceRegistry: {
+                    removeMcpResourcesByServer: vi.fn(),
+                    registerResource: vi.fn(),
+                },
+            });
             // 2. Extract the callback passed to setNotificationHandler for tools
             const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
             const notificationCallback = toolUpdateCall[1];
@@ -975,8 +1052,14 @@ describe('mcp-client', () => {
                 getToolsByServer: vi.fn().mockReturnValue([]),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {},
+                resourceRegistry: {},
+            });
             const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
             const notificationCallback = toolUpdateCall[1];
             // Trigger notification - should fail internally but catch the error
@@ -1020,10 +1103,22 @@ describe('mcp-client', () => {
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
             const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);
-            const clientA = new McpClient('server-A', { command: 'cmd-a' }, mockedToolRegistry, {}, {}, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
-            const clientB = new McpClient('server-B', { command: 'cmd-b' }, mockedToolRegistry, {}, {}, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
+            const clientA = new McpClient('server-A', { command: 'cmd-a' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
+            const clientB = new McpClient('server-B', { command: 'cmd-b' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
             await clientA.connect();
+            // INJECTED REGISTRIES
+            clientA.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {},
+                resourceRegistry: {},
+            });
             await clientB.connect();
+            // INJECTED REGISTRIES
+            clientB.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {},
+                resourceRegistry: {},
+            });
             const toolUpdateCallA = mockClientA.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
             const handlerA = toolUpdateCallA[1];
             const toolUpdateCallB = mockClientB.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
@@ -1074,17 +1169,23 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', 
             // Set a very short timeout
-            { command: 'test-command', timeout: 50 }, mockedToolRegistry, {
-                getPromptsByServer: vi.fn().mockReturnValue([]),
-                registerPrompt: vi.fn(),
-                removePromptsByServer: vi.fn(),
-            }, {
-                getResourcesByServer: vi.fn().mockReturnValue([]),
-                registerResource: vi.fn(),
-                removeResourcesByServer: vi.fn(),
-                setResourcesForServer: vi.fn(),
-            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
+            { command: 'test-command', timeout: 50 }, workspaceContext, MOCK_CONTEXT, false, '0.0.1');
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {
+                    getPromptsByServer: vi.fn().mockReturnValue([]),
+                    registerPrompt: vi.fn(),
+                    removePromptsByServer: vi.fn(),
+                },
+                resourceRegistry: {
+                    getResourcesByServer: vi.fn().mockReturnValue([]),
+                    registerResource: vi.fn(),
+                    removeResourcesByServer: vi.fn(),
+                    setResourcesForServer: vi.fn(),
+                },
+            });
             const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
             const notificationCallback = toolUpdateCall[1];
             const refreshPromise = notificationCallback();
@@ -1119,17 +1220,23 @@ describe('mcp-client', () => {
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
             const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {
-                getPromptsByServer: vi.fn().mockReturnValue([]),
-                registerPrompt: vi.fn(),
-                removePromptsByServer: vi.fn(),
-            }, {
-                getResourcesByServer: vi.fn().mockReturnValue([]),
-                registerResource: vi.fn(),
-                removeResourcesByServer: vi.fn(),
-                setResourcesForServer: vi.fn(),
-            }, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
+            const client = new McpClient('test-server', { command: 'test-command' }, workspaceContext, MOCK_CONTEXT, false, '0.0.1', onContextUpdatedSpy);
             await client.connect();
+            // INJECTED REGISTRIES
+            client.registeredRegistries?.add({
+                toolRegistry: mockedToolRegistry,
+                promptRegistry: {
+                    getPromptsByServer: vi.fn().mockReturnValue([]),
+                    registerPrompt: vi.fn(),
+                    removePromptsByServer: vi.fn(),
+                },
+                resourceRegistry: {
+                    getResourcesByServer: vi.fn().mockReturnValue([]),
+                    registerResource: vi.fn(),
+                    removeResourcesByServer: vi.fn(),
+                    setResourcesForServer: vi.fn(),
+                },
+            });
             const toolUpdateCall = mockedClient.setNotificationHandler.mock.calls.find((call) => call[0] === ToolListChangedNotificationSchema);
             const notificationCallback = toolUpdateCall[1];
             vi.useFakeTimers();