@office-ai/aioncli-core 0.26.0 → 0.30.1

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

@@ -14,6 +14,7 @@ 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 { ApprovalMode, PolicyDecision } from '../policy/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';
@@ -95,7 +96,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(mockedClient.listTools).toHaveBeenCalledWith({}, { timeout: 600000 });
@@ -157,14 +158,14 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
             expect(consoleWarnSpy).not.toHaveBeenCalled();
             consoleWarnSpy.mockRestore();
         });
-        it('should handle errors when discovering prompts', async () => {
+        it('should propagate errors when discovering prompts', async () => {
             const mockedClient = {
                 connect: vi.fn(),
                 discover: vi.fn(),
@@ -194,9 +195,9 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
-            await expect(client.discover({})).rejects.toThrow('No prompts, tools, or resources found on the server.');
+            await expect(client.discover({})).rejects.toThrow('Test error');
             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 () => {
@@ -229,7 +230,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await expect(client.discover({})).rejects.toThrow('No prompts, tools, or resources found on the server.');
         });
@@ -272,11 +273,121 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
         });
+        it('should register tool with readOnlyHint and add policy rule', 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({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'readOnlyTool',
+                            description: 'A read-only tool',
+                            inputSchema: { type: 'object', properties: {} },
+                            annotations: { readOnlyHint: true },
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockPolicyEngine = {
+                addRule: vi.fn(),
+            };
+            const mockConfig = {
+                getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
+            };
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+                removeMcpToolsByServer: vi.fn(),
+            };
+            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, '0.0.1');
+            await client.connect();
+            await client.discover(mockConfig);
+            // Verify tool registration
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
+            // Verify policy rule addition
+            expect(mockPolicyEngine.addRule).toHaveBeenCalledWith({
+                toolName: 'test-server__readOnlyTool',
+                decision: PolicyDecision.ASK_USER,
+                priority: 50,
+                modes: [ApprovalMode.PLAN],
+                source: 'MCP Annotation (readOnlyHint) - test-server',
+            });
+        });
+        it('should not add policy rule for tool without readOnlyHint', 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({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'writeTool',
+                            description: 'A write tool',
+                            inputSchema: { type: 'object', properties: {} },
+                            // No annotations or readOnlyHint: false
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockPolicyEngine = {
+                addRule: vi.fn(),
+            };
+            const mockConfig = {
+                getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
+            };
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+                removeMcpToolsByServer: vi.fn(),
+            };
+            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, '0.0.1');
+            await client.connect();
+            await client.discover(mockConfig);
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
+            expect(mockPolicyEngine.addRule).not.toHaveBeenCalled();
+        });
         it('should discover tools with $defs and $ref in schema', async () => {
             const mockedClient = {
                 connect: vi.fn(),
@@ -331,7 +442,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -395,7 +506,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith('test-server', [
@@ -461,7 +572,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, promptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce();
@@ -472,6 +583,64 @@ describe('mcp-client', () => {
             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('refreshes prompts when prompt list change notification is received', async () => {
+            let listCallCount = 0;
+            let promptListHandler;
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                setNotificationHandler: vi.fn((_, handler) => {
+                    promptListHandler = handler;
+                }),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ prompts: { listChanged: true } }),
+                listPrompts: vi.fn().mockImplementation(() => {
+                    listCallCount += 1;
+                    if (listCallCount === 1) {
+                        return Promise.resolve({
+                            prompts: [{ name: 'one', description: 'first' }],
+                        });
+                    }
+                    return Promise.resolve({
+                        prompts: [{ name: 'two', description: 'second' }],
+                    });
+                }),
+                request: vi.fn().mockResolvedValue({ 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, '0.0.1');
+            await client.connect();
+            await client.discover({ sanitizationConfig: EMPTY_CONFIG });
+            expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce();
+            expect(promptListHandler).toBeDefined();
+            await promptListHandler?.({
+                method: 'notifications/prompts/list_changed',
+            });
+            expect(promptRegistry.removePromptsByServer).toHaveBeenCalledWith('test-server');
+            expect(promptRegistry.registerPrompt).toHaveBeenLastCalledWith(expect.objectContaining({ name: 'two' }));
+            expect(coreEvents.emitFeedback).toHaveBeenCalledWith('info', 'Prompts updated for server: test-server');
+        });
         it('should remove tools and prompts on disconnect', async () => {
             const mockedClient = {
                 connect: vi.fn(),
@@ -517,7 +686,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, mockedPromptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            }, mockedToolRegistry, mockedPromptRegistry, resourceRegistry, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             await client.discover({});
             expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -547,7 +716,7 @@ describe('mcp-client', () => {
             };
             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);
+            const client = new McpClient('test-server', { command: 'test-command' }, {}, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             expect(mockedClient.setNotificationHandler).toHaveBeenCalledWith(ToolListChangedNotificationSchema, expect.any(Function));
         });
@@ -562,7 +731,7 @@ describe('mcp-client', () => {
             };
             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);
+            const client = new McpClient('test-server', { command: 'test-command' }, {}, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             expect(mockedClient.setNotificationHandler).not.toHaveBeenCalled();
         });
@@ -597,7 +766,7 @@ describe('mcp-client', () => {
             };
             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);
+            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1', onToolsUpdatedSpy);
             // 1. Connect (sets up listener)
             await client.connect();
             // 2. Extract the callback passed to setNotificationHandler
@@ -635,7 +804,7 @@ describe('mcp-client', () => {
                 removeMcpToolsByServer: vi.fn(),
                 getMessageBus: vi.fn().mockReturnValue(undefined),
             };
-            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
             // Trigger notification - should fail internally but catch the error
@@ -678,8 +847,8 @@ describe('mcp-client', () => {
                 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);
+            const clientA = new McpClient('server-A', { command: 'cmd-a' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1', onToolsUpdatedSpy);
+            const clientB = new McpClient('server-B', { command: 'cmd-b' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1', onToolsUpdatedSpy);
             await clientA.connect();
             await clientB.connect();
             const handlerA = mockClientA.setNotificationHandler.mock.calls[0][1];
@@ -711,8 +880,7 @@ describe('mcp-client', () => {
                     }
                     options?.signal?.addEventListener('abort', () => {
                         reject(new Error('Operation aborted'));
-                    });
-                    // Intentionally do not resolve immediately to simulate lag
+                    }, { once: true });
                 })),
                 listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
                 request: vi.fn().mockResolvedValue({}),
@@ -729,7 +897,7 @@ describe('mcp-client', () => {
             };
             const client = new McpClient('test-server', 
             // Set a short timeout
-            { command: 'test-command', timeout: 100 }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false);
+            { command: 'test-command', timeout: 100 }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1');
             await client.connect();
             const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
             const refreshPromise = notificationCallback();
@@ -763,7 +931,7 @@ describe('mcp-client', () => {
                 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);
+            const client = new McpClient('test-server', { command: 'test-command' }, mockedToolRegistry, {}, {}, workspaceContext, { sanitizationConfig: EMPTY_CONFIG }, false, '0.0.1', onToolsUpdatedSpy);
             await client.connect();
             const notificationCallback = mockedClient.setNotificationHandler.mock.calls[0][1];
             await notificationCallback();
@@ -857,6 +1025,47 @@ describe('mcp-client', () => {
                 stderr: 'pipe',
             });
         });
+        it('sets an env variable GEMINI_CLI=1 for stdio MCP servers', async () => {
+            const mockedTransport = vi
+                .spyOn(SdkClientStdioLib, 'StdioClientTransport')
+                .mockReturnValue({});
+            await createTransport('test-server', {
+                command: 'test-command',
+                args: ['--foo', 'bar'],
+                env: {},
+                cwd: 'test/cwd',
+            }, false, EMPTY_CONFIG);
+            const callArgs = mockedTransport.mock.calls[0][0];
+            expect(callArgs.env).toBeDefined();
+            expect(callArgs.env['GEMINI_CLI']).toBe('1');
+        });
+        it('should exclude extension settings with undefined values from environment', async () => {
+            const mockedTransport = vi
+                .spyOn(SdkClientStdioLib, 'StdioClientTransport')
+                .mockReturnValue({});
+            await createTransport('test-server', {
+                command: 'test-command',
+                extension: {
+                    name: 'test-ext',
+                    resolvedSettings: [
+                        {
+                            envVar: 'GEMINI_CLI_EXT_VAR',
+                            value: undefined,
+                            sensitive: false,
+                            name: 'ext-setting',
+                        },
+                    ],
+                    version: '',
+                    isActive: false,
+                    path: '',
+                    contextFiles: [],
+                    id: '',
+                },
+            }, false, EMPTY_CONFIG);
+            const callArgs = mockedTransport.mock.calls[0][0];
+            expect(callArgs.env).toBeDefined();
+            expect(callArgs.env['GEMINI_CLI_EXT_VAR']).toBeUndefined();
+        });
         describe('useGoogleCredentialProvider', () => {
             beforeEach(() => {
                 // Mock GoogleAuth client
@@ -1062,7 +1271,7 @@ describe('connectToMcpServer with OAuth', () => {
             capturedTransport = transport;
             return Promise.resolve();
         });
-        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
+        const client = await connectToMcpServer('0.0.1', 'test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
         expect(client).toBe(mockedClient);
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
         expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
@@ -1087,7 +1296,7 @@ describe('connectToMcpServer with OAuth', () => {
             capturedTransport = transport;
             return Promise.resolve();
         });
-        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
+        const client = await connectToMcpServer('0.0.1', 'test-server', { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, EMPTY_CONFIG);
         expect(client).toBe(mockedClient);
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
         expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
@@ -1121,13 +1330,13 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => {
     });
     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');
+        await expect(connectToMcpServer('0.0.1', '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');
+        await expect(connectToMcpServer('0.0.1', '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);
     });
@@ -1135,7 +1344,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => {
         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);
+        const client = await connectToMcpServer('0.0.1', '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);
@@ -1146,14 +1355,14 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => {
         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');
+        await expect(connectToMcpServer('0.0.1', '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);
+        const client = await connectToMcpServer('0.0.1', 'test-server', { url: 'http://test-server' }, false, workspaceContext, EMPTY_CONFIG);
         expect(client).toBe(mockedClient);
         expect(mockedClient.connect).toHaveBeenCalledTimes(2);
     });
@@ -1179,6 +1388,15 @@ describe('connectToMcpServer - OAuth with transport fallback', () => {
         vi.spyOn(console, 'log').mockImplementation(() => { });
         vi.spyOn(console, 'warn').mockImplementation(() => { });
         vi.spyOn(console, 'error').mockImplementation(() => { });
+        // Mock fetch to prevent real network calls during OAuth discovery fallback.
+        // When a 401 error lacks a www-authenticate header, the code attempts to
+        // fetch the header directly from the server, which would hang without this mock.
+        vi.stubGlobal('fetch', vi.fn().mockResolvedValue({
+            status: 401,
+            headers: new Headers({
+                'www-authenticate': `Bearer realm="test", resource_metadata="http://test-server/.well-known/oauth-protected-resource"`,
+            }),
+        }));
         mockTokenStorage = {
             getCredentials: vi.fn().mockResolvedValue({ clientId: 'test-client' }),
         };
@@ -1197,6 +1415,7 @@ describe('connectToMcpServer - OAuth with transport fallback', () => {
     });
     afterEach(() => {
         vi.clearAllMocks();
+        vi.unstubAllGlobals();
     });
     it('should handle HTTP 404 → SSE 401 → OAuth → SSE+OAuth succeeds', async () => {
         // Tests that OAuth flow works when SSE (not HTTP) requires auth
@@ -1204,7 +1423,7 @@ describe('connectToMcpServer - OAuth with transport fallback', () => {
             .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);
+        const client = await connectToMcpServer('0.0.1', '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();