@office-ai/aioncli-core 0.18.7 → 0.24.0-preview.1

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

@@ -6,25 +6,33 @@
 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';
-import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { StreamableHTTPClientTransport, StreamableHTTPError, } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { AuthProviderType } from '../config/config.js';
 import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
 import { MCPOAuthProvider } from '../mcp/oauth-provider.js';
 import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
 import { OAuthUtils } from '../mcp/oauth-utils.js';
+import { ToolListChangedNotificationSchema } from '@modelcontextprotocol/sdk/types.js';
 import { WorkspaceContext } from '../utils/workspaceContext.js';
 import { connectToMcpServer, createTransport, hasNetworkTransport, isEnabled, McpClient, populateMcpServerCommand, } from './mcp-client.js';
 import * as fs from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
 import { coreEvents } from '../utils/events.js';
+const EMPTY_CONFIG = {
+    enableEnvironmentVariableRedaction: true,
+    allowedEnvironmentVariables: [],
+    blockedEnvironmentVariables: [],
+};
 vi.mock('@modelcontextprotocol/sdk/client/stdio.js');
 vi.mock('@modelcontextprotocol/sdk/client/index.js');
 vi.mock('@google/genai');
 vi.mock('../mcp/oauth-provider.js');
 vi.mock('../mcp/oauth-token-storage.js');
 vi.mock('../mcp/oauth-utils.js');
+vi.mock('google-auth-library');
+import { GoogleAuth } from 'google-auth-library';
 vi.mock('../utils/events.js', () => ({
     coreEvents: {
         emitFeedback: vi.fn(),
@@ -52,6 +60,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
                 listTools: vi.fn().mockResolvedValue({
                     tools: [
@@ -76,12 +85,20 @@ describe('mcp-client', () => {
                 sortTools: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, workspaceContext, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
             await client.discover({});
-            expect(mockedClient.listTools).toHaveBeenCalledWith({});
+            expect(mockedClient.listTools).toHaveBeenCalledWith({}, { timeout: 600000 });
         });
         it('should not skip tools even if a parameter is missing a type', async () => {
             const consoleWarnSpy = vi
@@ -94,6 +111,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
                 listTools: vi.fn().mockResolvedValue({
                     tools: [
@@ -129,9 +147,17 @@ describe('mcp-client', () => {
                 sortTools: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, workspaceContext, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
@@ -146,6 +172,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
                 listTools: vi.fn().mockResolvedValue({ tools: [] }),
                 listPrompts: vi.fn().mockRejectedValue(new Error('Test error')),
@@ -157,11 +184,19 @@ describe('mcp-client', () => {
                 registerTool: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, workspaceContext, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
-            await expect(client.discover({})).rejects.toThrow('No prompts or tools found on the server.');
+            await expect(client.discover({})).rejects.toThrow('No prompts, tools, or resources found on the server.');
             expect(coreEvents.emitFeedback).toHaveBeenCalledWith('error', `Error discovering prompts from test-server: Test error`, expect.any(Error));
         });
         it('should not discover tools if server does not support them', async () => {
@@ -172,6 +207,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
                 listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
                 request: vi.fn().mockResolvedValue({}),
@@ -183,11 +219,19 @@ describe('mcp-client', () => {
                 sortTools: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, workspaceContext, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
-            await expect(client.discover({})).rejects.toThrow('No prompts or tools found on the server.');
+            await expect(client.discover({})).rejects.toThrow('No prompts, tools, or resources found on the server.');
         });
         it('should discover tools if server supports them', async () => {
             const mockedClient = {
@@ -197,6 +241,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
                 listTools: vi.fn().mockResolvedValue({
                     tools: [
@@ -217,9 +262,17 @@ describe('mcp-client', () => {
                 sortTools: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, workspaceContext, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -232,6 +285,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
                 listTools: vi.fn().mockResolvedValue({
                     tools: [
@@ -267,9 +321,17 @@ describe('mcp-client', () => {
                 sortTools: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, workspaceContext, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -290,6 +352,126 @@ describe('mcp-client', () => {
                 },
             });
         });
+        it('should discover resources when a server only exposes resources', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
+                getServerCapabilities: vi.fn().mockReturnValue({ resources: {} }),
+                request: vi.fn().mockImplementation(({ method }) => {
+                    if (method === 'resources/list') {
+                        return Promise.resolve({
+                            resources: [
+                                {
+                                    uri: 'file:///tmp/resource.txt',
+                                    name: 'resource',
+                                    description: 'Test Resource',
+                                    mimeType: 'text/plain',
+                                },
+                            ],
+                        });
+                    }
+                    return Promise.resolve({ prompts: [] });
+                }),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
+            const client = new McpClient('test-server', {
+                command: 'test-command',
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            await client.connect();
+            await client.discover({});
+            expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith('test-server', [
+                expect.objectContaining({
+                    uri: 'file:///tmp/resource.txt',
+                    name: 'resource',
+                }),
+            ]);
+        });
+        it('refreshes registry when resource list change notification is received', async () => {
+            let listCallCount = 0;
+            let resourceListHandler;
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn((_, handler) => {
+                    resourceListHandler = handler;
+                }),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ resources: { listChanged: true } }),
+                request: vi.fn().mockImplementation(({ method }) => {
+                    if (method === 'resources/list') {
+                        listCallCount += 1;
+                        if (listCallCount === 1) {
+                            return Promise.resolve({
+                                resources: [
+                                    {
+                                        uri: 'file:///tmp/one.txt',
+                                    },
+                                ],
+                            });
+                        }
+                        return Promise.resolve({
+                            resources: [
+                                {
+                                    uri: 'file:///tmp/two.txt',
+                                },
+                            ],
+                        });
+                    }
+                    return Promise.resolve({ prompts: [] });
+                }),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const promptRegistry = {
+                registerPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
+            const client = new McpClient('test-server', {
+                command: 'test-command',
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            await client.connect();
+            await client.discover({});
+            expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce();
+            expect(resourceListHandler).toBeDefined();
+            await resourceListHandler?.({
+                method: 'notifications/resources/list_changed',
+            });
+            expect(resourceRegistry.setResourcesForServer).toHaveBeenLastCalledWith('test-server', [expect.objectContaining({ uri: 'file:///tmp/two.txt' })]);
+            expect(coreEvents.emitFeedback).toHaveBeenCalledWith('info', 'Resources updated for server: test-server');
+        });
         it('should remove tools and prompts on disconnect', async () => {
             const mockedClient = {
                 connect: vi.fn(),
@@ -297,6 +479,7 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn(),
                 getServerCapabilities: vi
                     .fn()
                     .mockReturnValue({ tools: {}, prompts: {} }),
@@ -328,9 +511,13 @@ describe('mcp-client', () => {
                 unregisterPrompt: vi.fn(),
                 removePromptsByServer: vi.fn(),
             };
+            const resourceRegistry = {
+                setResourcesForServer: vi.fn(),
+                removeResourcesByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, mockedPromptRegistry, workspaceContext, false);
+            }, mockedToolRegistry, mockedPromptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -339,6 +526,251 @@ describe('mcp-client', () => {
             expect(mockedClient.close).toHaveBeenCalledOnce();
             expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledOnce();
             expect(mockedPromptRegistry.removePromptsByServer).toHaveBeenCalledOnce();
+            expect(resourceRegistry.removeResourcesByServer).toHaveBeenCalledOnce();
+        });
+    });
+    describe('Dynamic Tool Updates', () => {
+        it('should set up notification handler if server supports tool list changes', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                // Capability enables the listener
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: { listChanged: true } }),
+                setNotificationHandler: vi.fn(),
+                listTools: vi.fn().mockResolvedValue({ tools: [] }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const client = new McpClient('test-server', { command: 'test-command' }, {}, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            await client.connect();
+            expect(mockedClient.setNotificationHandler).toHaveBeenCalledWith(ToolListChangedNotificationSchema, expect.any(Function));
+        });
+        it('should NOT set up notification handler if server lacks capability', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }), // No listChanged
+                setNotificationHandler: vi.fn(),
+                request: vi.fn().mockResolvedValue({}),
+                registerCapabilities: vi.fn().mockResolvedValue({}),
+                setRequestHandler: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const client = new McpClient('test-server', { command: 'test-command' }, {}, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            await client.connect();
+            expect(mockedClient.setNotificationHandler).not.toHaveBeenCalled();
+        });
+        it('should refresh tools and notify manager when notification is received', async () => {
+            // Setup mocks
+            const mockedClient = {
+                connect: vi.fn(),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: { listChanged: true } }),
+                setNotificationHandler: vi.fn(),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'newTool',
+                            inputSchema: { type: 'object', properties: {} },
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+                registerCapabilities: vi.fn().mockResolvedValue({}),
+                setRequestHandler: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                removeMcpToolsByServer: vi.fn(),
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const onToolsUpdatedSpy = vi.fn().mockResolvedValue(undefined);
+            // Initialize client with onToolsUpdated callback
+            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, onToolsUpdatedSpy);
+            // 1. Connect (sets up listener)
+            await client.connect();
+            // 2. Extract the callback passed to setNotificationHandler
+            const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
+            // 3. Trigger the notification manually
+            await notificationCallback();
+            // 4. Assertions
+            // It should clear old tools
+            expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledWith('test-server');
+            // It should fetch new tools (listTools called inside discoverTools)
+            expect(mockedClient.listTools).toHaveBeenCalled();
+            // It should register the new tool
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalled();
+            // It should notify the manager
+            expect(onToolsUpdatedSpy).toHaveBeenCalled();
+            // It should emit feedback event
+            expect(coreEvents.emitFeedback).toHaveBeenCalledWith('info', 'Tools updated for server: test-server');
+        });
+        it('should handle errors during tool refresh gracefully', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: { listChanged: true } }),
+                setNotificationHandler: vi.fn(),
+                // Simulate error during discovery
+                listTools: vi.fn().mockRejectedValue(new Error('Network blip')),
+                request: vi.fn().mockResolvedValue({}),
+                registerCapabilities: vi.fn().mockResolvedValue({}),
+                setRequestHandler: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                removeMcpToolsByServer: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            await client.connect();
+            const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
+            // Trigger notification - should fail internally but catch the error
+            await notificationCallback();
+            // Should try to remove tools
+            expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalled();
+            // Should NOT emit success feedback
+            expect(coreEvents.emitFeedback).not.toHaveBeenCalledWith('info', expect.stringContaining('Tools updated'));
+        });
+        it('should handle concurrent updates from multiple servers', async () => {
+            const createMockSdkClient = (toolName) => ({
+                connect: vi.fn(),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: { listChanged: true } }),
+                setNotificationHandler: vi.fn(),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: toolName,
+                            inputSchema: { type: 'object', properties: {} },
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+                registerCapabilities: vi.fn().mockResolvedValue({}),
+                setRequestHandler: vi.fn().mockResolvedValue({}),
+            });
+            const mockClientA = createMockSdkClient('tool-from-A');
+            const mockClientB = createMockSdkClient('tool-from-B');
+            vi.mocked(ClientLib.Client)
+                .mockReturnValueOnce(mockClientA)
+                .mockReturnValueOnce(mockClientB);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                removeMcpToolsByServer: vi.fn(),
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const onToolsUpdatedSpy = vi.fn().mockResolvedValue(undefined);
+            const clientA = new McpClient('server-A', { command: 'cmd-a' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, onToolsUpdatedSpy);
+            const clientB = new McpClient('server-B', { command: 'cmd-b' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, onToolsUpdatedSpy);
+            await clientA.connect();
+            await clientB.connect();
+            const handlerA = mockClientA.setNotificationHandler.mock.calls[0][1];
+            const handlerB = mockClientB.setNotificationHandler.mock.calls[0][1];
+            // Trigger burst updates simultaneously
+            await Promise.all([handlerA(), handlerB()]);
+            expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledWith('server-A');
+            expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledWith('server-B');
+            // Verify fetching happened on both clients
+            expect(mockClientA.listTools).toHaveBeenCalled();
+            expect(mockClientB.listTools).toHaveBeenCalled();
+            // Verify tools from both servers were registered (2 total calls)
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
+            // Verify the update callback was triggered for both
+            expect(onToolsUpdatedSpy).toHaveBeenCalledTimes(2);
+        });
+        it('should abort discovery and log error if timeout is exceeded during refresh', async () => {
+            vi.useFakeTimers();
+            const mockedClient = {
+                connect: vi.fn(),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: { listChanged: true } }),
+                setNotificationHandler: vi.fn(),
+                // Mock listTools to simulate a long running process that respects the abort signal
+                listTools: vi.fn().mockImplementation(async (params, options) => new Promise((resolve, reject) => {
+                    if (options?.signal?.aborted) {
+                        return reject(new Error('Operation aborted'));
+                    }
+                    options?.signal?.addEventListener('abort', () => {
+                        reject(new Error('Operation aborted'));
+                    });
+                    // Intentionally do not resolve immediately to simulate lag
+                })),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+                registerCapabilities: vi.fn().mockResolvedValue({}),
+                setRequestHandler: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                removeMcpToolsByServer: vi.fn(),
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const client = new McpClient('test-server', 
+            // Set a short timeout
+            { command: 'test-command', timeout: 100 }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            await client.connect();
+            const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
+            const refreshPromise = notificationCallback();
+            vi.advanceTimersByTime(150);
+            await refreshPromise;
+            expect(mockedClient.listTools).toHaveBeenCalledWith(expect.anything(), expect.objectContaining({
+                signal: expect.any(AbortSignal),
+            }));
+            expect(mockedToolRegistry.registerTool).not.toHaveBeenCalled();
+            vi.useRealTimers();
+        });
+        it('should pass abort signal to onToolsUpdated callback', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: { listChanged: true } }),
+                setNotificationHandler: vi.fn(),
+                listTools: vi.fn().mockResolvedValue({ tools: [] }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+                registerCapabilities: vi.fn().mockResolvedValue({}),
+                setRequestHandler: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                removeMcpToolsByServer: vi.fn(),
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const onToolsUpdatedSpy = vi.fn().mockResolvedValue(undefined);
+            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, onToolsUpdatedSpy);
+            await client.connect();
+            const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
+            await notificationCallback();
+            expect(onToolsUpdatedSpy).toHaveBeenCalledWith(expect.any(AbortSignal));
+            // Verify the signal passed was not aborted (happy path)
+            const signal = onToolsUpdatedSpy.mock.calls[0][0];
+            expect(signal.aborted).toBe(false);
         });
     });
     describe('appendMcpServerCommand', () => {
@@ -365,9 +797,10 @@ describe('mcp-client', () => {
             it('without headers', async () => {
                 const transport = await createTransport('test-server', {
                     httpUrl: 'http://test-server',
-                }, false);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
+                }, false, EMPTY_CONFIG);
+                expect(transport).toEqual(new StreamableHTTPClientTransport(new URL('http://test-server'), {
+                    requestInit: { headers: {} },
+                }));
             });
             it('with headers', async () => {
                 // We need this to be an any type because we dig into its private state.
@@ -375,20 +808,22 @@ describe('mcp-client', () => {
                 const transport = await createTransport('test-server', {
                     httpUrl: 'http://test-server',
                     headers: { Authorization: 'derp' },
-                }, false);
-                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
-                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
-                const authHeader = transport._requestInit?.headers?.['Authorization'];
-                expect(authHeader).toBe('derp');
+                }, false, EMPTY_CONFIG);
+                expect(transport).toEqual(new StreamableHTTPClientTransport(new URL('http://test-server'), {
+                    requestInit: {
+                        headers: { Authorization: 'derp' },
+                    },
+                }));
             });
         });
         describe('should connect via url', () => {
             it('without headers', async () => {
                 const transport = await createTransport('test-server', {
                     url: 'http://test-server',
-                }, false);
-                expect(transport).toBeInstanceOf(SSEClientTransport);
-                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
+                }, false, EMPTY_CONFIG);
+                expect(transport).toEqual(new SSEClientTransport(new URL('http://test-server'), {
+                    requestInit: { headers: {} },
+                }));
             });
             it('with headers', async () => {
                 // We need this to be an any type because we dig into its private state.
@@ -396,11 +831,12 @@ describe('mcp-client', () => {
                 const transport = await createTransport('test-server', {
                     url: 'http://test-server',
                     headers: { Authorization: 'derp' },
-                }, false);
-                expect(transport).toBeInstanceOf(SSEClientTransport);
-                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
-                const authHeader = transport._requestInit?.headers?.['Authorization'];
-                expect(authHeader).toBe('derp');
+                }, false, EMPTY_CONFIG);
+                expect(transport).toEqual(new SSEClientTransport(new URL('http://test-server'), {
+                    requestInit: {
+                        headers: { Authorization: 'derp' },
+                    },
+                }));
             });
         });
         it('should connect via command', async () => {
@@ -412,16 +848,24 @@ describe('mcp-client', () => {
                 args: ['--foo', 'bar'],
                 env: { FOO: 'bar' },
                 cwd: 'test/cwd',
-            }, false);
+            }, false, EMPTY_CONFIG);
             expect(mockedTransport).toHaveBeenCalledWith({
                 command: 'test-command',
                 args: ['--foo', 'bar'],
                 cwd: 'test/cwd',
-                env: { ...process.env, FOO: 'bar' },
+                env: expect.objectContaining({ FOO: 'bar' }),
                 stderr: 'pipe',
             });
         });
         describe('useGoogleCredentialProvider', () => {
+            beforeEach(() => {
+                // Mock GoogleAuth client
+                const mockClient = {
+                    getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }),
+                    quotaProjectId: 'myproject',
+                };
+                GoogleAuth.prototype.getClient = vi.fn().mockResolvedValue(mockClient);
+            });
             it('should use GoogleCredentialProvider when specified', async () => {
                 const transport = await createTransport('test-server', {
                     httpUrl: 'http://test.googleapis.com',
@@ -432,7 +876,7 @@ describe('mcp-client', () => {
                     headers: {
                         'X-Goog-User-Project': 'myproject',
                     },
-                }, false);
+                }, false, EMPTY_CONFIG);
                 expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
                 // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 const authProvider = transport._authProvider;
@@ -441,14 +885,53 @@ describe('mcp-client', () => {
                 const googUserProject = transport._requestInit?.headers?.['X-Goog-User-Project'];
                 expect(googUserProject).toBe('myproject');
             });
+            it('should use headers from GoogleCredentialProvider', async () => {
+                const mockGetRequestHeaders = vi.fn().mockResolvedValue({
+                    'X-Goog-User-Project': 'provider-project',
+                });
+                vi.spyOn(GoogleCredentialProvider.prototype, 'getRequestHeaders').mockImplementation(mockGetRequestHeaders);
+                const transport = await createTransport('test-server', {
+                    httpUrl: 'http://test.googleapis.com',
+                    authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
+                    oauth: {
+                        scopes: ['scope1'],
+                    },
+                }, false, EMPTY_CONFIG);
+                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(mockGetRequestHeaders).toHaveBeenCalled();
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                const headers = transport._requestInit?.headers;
+                expect(headers['X-Goog-User-Project']).toBe('provider-project');
+            });
+            it('should prioritize provider headers over config headers', async () => {
+                const mockGetRequestHeaders = vi.fn().mockResolvedValue({
+                    'X-Goog-User-Project': 'provider-project',
+                });
+                vi.spyOn(GoogleCredentialProvider.prototype, 'getRequestHeaders').mockImplementation(mockGetRequestHeaders);
+                const transport = await createTransport('test-server', {
+                    httpUrl: 'http://test.googleapis.com',
+                    authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
+                    oauth: {
+                        scopes: ['scope1'],
+                    },
+                    headers: {
+                        'X-Goog-User-Project': 'config-project',
+                    },
+                }, false, EMPTY_CONFIG);
+                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                const headers = transport._requestInit?.headers;
+                expect(headers['X-Goog-User-Project']).toBe('provider-project');
+            });
             it('should use GoogleCredentialProvider with SSE transport', async () => {
                 const transport = await createTransport('test-server', {
                     url: 'http://test.googleapis.com',
+                    type: 'sse',
                     authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
                     oauth: {
                         scopes: ['scope1'],
                     },
-                }, false);
+                }, false, EMPTY_CONFIG);
                 expect(transport).toBeInstanceOf(SSEClientTransport);
                 // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 const authProvider = transport._authProvider;
@@ -460,7 +943,7 @@ describe('mcp-client', () => {
                     oauth: {
                         scopes: ['scope1'],
                     },
-                }, false)).rejects.toThrow('URL must be provided in the config for Google Credentials provider');
+                }, false, EMPTY_CONFIG)).rejects.toThrow('URL must be provided in the config for Google Credentials provider');
             });
         });
     });
@@ -566,7 +1049,7 @@ describe('connectToMcpServer with OAuth', () => {
         const authUrl = 'http://auth.example.com/auth';
         const tokenUrl = 'http://auth.example.com/token';
         const wwwAuthHeader = `Bearer realm="test", resource_metadata="http://test-server.com/.well-known/oauth-protected-resource"`;
-        vi.mocked(mockedClient.connect).mockRejectedValueOnce(new Error(`401 Unauthorized\nwww-authenticate: ${wwwAuthHeader}`));
+        vi.mocked(mockedClient.connect).mockRejectedValueOnce(new StreamableHTTPError(401, `Unauthorized\nwww-authenticate: ${wwwAuthHeader}`));
         vi.mocked(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
             authorizationUrl: authUrl,
             tokenUrl,
@@ -579,7 +1062,7 @@ describe('connectToMcpServer with OAuth', () => {
             capturedTransport = transport;
             return Promise.resolve();
         });
-        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl }, false, workspaceContext);
+        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
         expect(client).toBe(mockedClient);
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
         expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
@@ -590,7 +1073,7 @@ describe('connectToMcpServer with OAuth', () => {
         const serverUrl = 'http://test-server.com';
         const authUrl = 'http://auth.example.com/auth';
         const tokenUrl = 'http://auth.example.com/token';
-        vi.mocked(mockedClient.connect).mockRejectedValueOnce(new Error('401 Unauthorized'));
+        vi.mocked(mockedClient.connect).mockRejectedValueOnce(new StreamableHTTPError(401, 'Unauthorized'));
         vi.mocked(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
             authorizationUrl: authUrl,
             tokenUrl,
@@ -604,7 +1087,7 @@ describe('connectToMcpServer with OAuth', () => {
             capturedTransport = transport;
             return Promise.resolve();
         });
-        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl }, false, workspaceContext);
+        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
         expect(client).toBe(mockedClient);
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
         expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
@@ -613,4 +1096,118 @@ describe('connectToMcpServer with OAuth', () => {
         expect(authHeader).toBe('Bearer test-access-token-from-discovery');
     });
 });
+describe('connectToMcpServer - HTTP→SSE fallback', () => {
+    let mockedClient;
+    let workspaceContext;
+    let testWorkspace;
+    beforeEach(() => {
+        mockedClient = {
+            connect: vi.fn(),
+            close: vi.fn(),
+            registerCapabilities: vi.fn(),
+            setRequestHandler: vi.fn(),
+            onclose: vi.fn(),
+            notification: vi.fn(),
+        };
+        vi.mocked(ClientLib.Client).mockImplementation(() => mockedClient);
+        testWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'gemini-agent-test-'));
+        workspaceContext = new WorkspaceContext(testWorkspace);
+        vi.spyOn(console, 'log').mockImplementation(() => { });
+        vi.spyOn(console, 'warn').mockImplementation(() => { });
+        vi.spyOn(console, 'error').mockImplementation(() => { });
+    });
+    afterEach(() => {
+        vi.clearAllMocks();
+    });
+    it('should NOT trigger fallback when type="http" is explicit', async () => {
+        vi.mocked(mockedClient.connect).mockRejectedValueOnce(new Error('Connection failed'));
+        await expect(connectToMcpServer('test-server', { url: 'http://test-server', type: 'http' }, false, workspaceContext, EMPTY_CONFIG)).rejects.toThrow('Connection failed');
+        // Should only try once (no fallback)
+        expect(mockedClient.connect).toHaveBeenCalledTimes(1);
+    });
+    it('should NOT trigger fallback when type="sse" is explicit', async () => {
+        vi.mocked(mockedClient.connect).mockRejectedValueOnce(new Error('Connection failed'));
+        await expect(connectToMcpServer('test-server', { url: 'http://test-server', type: 'sse' }, false, workspaceContext, EMPTY_CONFIG)).rejects.toThrow('Connection failed');
+        // Should only try once (no fallback)
+        expect(mockedClient.connect).toHaveBeenCalledTimes(1);
+    });
+    it('should trigger fallback when url provided without type and HTTP fails', async () => {
+        vi.mocked(mockedClient.connect)
+            .mockRejectedValueOnce(new StreamableHTTPError(500, 'Server error'))
+            .mockResolvedValueOnce(undefined);
+        const client = await connectToMcpServer('test-server', { url: 'http://test-server' }, false, workspaceContext, EMPTY_CONFIG);
+        expect(client).toBe(mockedClient);
+        // First HTTP attempt fails, second SSE attempt succeeds
+        expect(mockedClient.connect).toHaveBeenCalledTimes(2);
+    });
+    it('should throw original HTTP error when both HTTP and SSE fail (non-401)', async () => {
+        const httpError = new StreamableHTTPError(500, 'Server error');
+        const sseError = new Error('SSE connection failed');
+        vi.mocked(mockedClient.connect)
+            .mockRejectedValueOnce(httpError)
+            .mockRejectedValueOnce(sseError);
+        await expect(connectToMcpServer('test-server', { url: 'http://test-server' }, false, workspaceContext, EMPTY_CONFIG)).rejects.toThrow('Server error');
+        expect(mockedClient.connect).toHaveBeenCalledTimes(2);
+    });
+    it('should handle HTTP 404 followed by SSE success', async () => {
+        vi.mocked(mockedClient.connect)
+            .mockRejectedValueOnce(new StreamableHTTPError(404, 'Not Found'))
+            .mockResolvedValueOnce(undefined);
+        const client = await connectToMcpServer('test-server', { url: 'http://test-server' }, false, workspaceContext, EMPTY_CONFIG);
+        expect(client).toBe(mockedClient);
+        expect(mockedClient.connect).toHaveBeenCalledTimes(2);
+    });
+});
+describe('connectToMcpServer - OAuth with transport fallback', () => {
+    let mockedClient;
+    let workspaceContext;
+    let testWorkspace;
+    let mockAuthProvider;
+    let mockTokenStorage;
+    beforeEach(() => {
+        mockedClient = {
+            connect: vi.fn(),
+            close: vi.fn(),
+            registerCapabilities: vi.fn(),
+            setRequestHandler: vi.fn(),
+            onclose: vi.fn(),
+            notification: vi.fn(),
+        };
+        vi.mocked(ClientLib.Client).mockImplementation(() => mockedClient);
+        testWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'gemini-agent-test-'));
+        workspaceContext = new WorkspaceContext(testWorkspace);
+        vi.spyOn(console, 'log').mockImplementation(() => { });
+        vi.spyOn(console, 'warn').mockImplementation(() => { });
+        vi.spyOn(console, 'error').mockImplementation(() => { });
+        mockTokenStorage = {
+            getCredentials: vi.fn().mockResolvedValue({ clientId: 'test-client' }),
+        };
+        vi.mocked(MCPOAuthTokenStorage).mockReturnValue(mockTokenStorage);
+        mockAuthProvider = {
+            authenticate: vi.fn().mockResolvedValue(undefined),
+            getValidToken: vi.fn().mockResolvedValue('test-access-token'),
+            tokenStorage: mockTokenStorage,
+        };
+        vi.mocked(MCPOAuthProvider).mockReturnValue(mockAuthProvider);
+        vi.mocked(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
+            authorizationUrl: 'http://auth.example.com/auth',
+            tokenUrl: 'http://auth.example.com/token',
+            scopes: ['test-scope'],
+        });
+    });
+    afterEach(() => {
+        vi.clearAllMocks();
+    });
+    it('should handle HTTP 404 → SSE 401 → OAuth → SSE+OAuth succeeds', async () => {
+        // Tests that OAuth flow works when SSE (not HTTP) requires auth
+        vi.mocked(mockedClient.connect)
+            .mockRejectedValueOnce(new StreamableHTTPError(404, 'Not Found'))
+            .mockRejectedValueOnce(new StreamableHTTPError(401, 'Unauthorized'))
+            .mockResolvedValueOnce(undefined);
+        const client = await connectToMcpServer('test-server', { url: 'http://test-server', oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
+        expect(client).toBe(mockedClient);
+        expect(mockedClient.connect).toHaveBeenCalledTimes(3);
+        expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
+    });
+});
 //# sourceMappingURL=mcp-client.test.js.map