@office-ai/aioncli-core 0.2.3 → 0.18.4

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

@@ -3,21 +3,43 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { afterEach, describe, expect, it, vi } from 'vitest';
-import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { populateMcpServerCommand, createTransport, isEnabled, hasValidTypes, McpClient, } from './mcp-client.js';
+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 * as ClientLib from '@modelcontextprotocol/sdk/client/index.js';
-import * as GenAiLib from '@google/genai';
-import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
+import { StreamableHTTPClientTransport } 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 { 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';
 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('../utils/events.js', () => ({
+    coreEvents: {
+        emitFeedback: vi.fn(),
+        emitConsoleLog: vi.fn(),
+    },
+}));
 describe('mcp-client', () => {
+    let workspaceContext;
+    let testWorkspace;
+    beforeEach(() => {
+        // create a tmp dir for this test
+        // Create a unique temporary directory for the workspace to avoid conflicts
+        testWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'gemini-agent-test-'));
+        workspaceContext = new WorkspaceContext(testWorkspace);
+    });
     afterEach(() => {
         vi.restoreAllMocks();
     });
@@ -30,29 +52,38 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
-            };
-            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
-            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
-            const mockedMcpToTool = vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
-                tool: () => ({
-                    functionDeclarations: [
+                getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
                         {
                             name: 'testFunction',
+                            inputSchema: {
+                                type: 'object',
+                                properties: {},
+                            },
                         },
                     ],
                 }),
-            });
+                listPrompts: vi.fn().mockResolvedValue({
+                    prompts: [],
+                }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            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 client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, {}, false);
+            }, mockedToolRegistry, {}, workspaceContext, false);
             await client.connect();
-            await client.discover();
-            expect(mockedMcpToTool).toHaveBeenCalledOnce();
+            await client.discover({});
+            expect(mockedClient.listTools).toHaveBeenCalledWith({});
         });
-        it('should skip tools if a parameter is missing a type', async () => {
+        it('should not skip tools even if a parameter is missing a type', async () => {
             const consoleWarnSpy = vi
                 .spyOn(console, 'warn')
                 .mockImplementation(() => { });
@@ -63,16 +94,12 @@ describe('mcp-client', () => {
                 getStatus: vi.fn(),
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
-                tool: vi.fn(),
-            };
-            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
-            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
-            vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
-                tool: () => Promise.resolve({
-                    functionDeclarations: [
+                getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
                         {
                             name: 'validTool',
-                            parametersJsonSchema: {
+                            inputSchema: {
                                 type: 'object',
                                 properties: {
                                     param1: { type: 'string' },
@@ -81,7 +108,7 @@ describe('mcp-client', () => {
                         },
                         {
                             name: 'invalidTool',
-                            parametersJsonSchema: {
+                            inputSchema: {
                                 type: 'object',
                                 properties: {
                                     param1: { description: 'a param with no type' },
@@ -90,25 +117,28 @@ describe('mcp-client', () => {
                         },
                     ],
                 }),
-            });
+                listPrompts: vi.fn().mockResolvedValue({
+                    prompts: [],
+                }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            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 client = new McpClient('test-server', {
                 command: 'test-command',
-            }, mockedToolRegistry, {}, {}, false);
+            }, mockedToolRegistry, {}, workspaceContext, false);
             await client.connect();
-            await client.discover();
-            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
-            expect(consoleWarnSpy).toHaveBeenCalledOnce();
-            expect(consoleWarnSpy).toHaveBeenCalledWith(`Skipping tool 'invalidTool' from MCP server 'test-server' because it has ` +
-                `missing types in its parameter schema. Please file an issue with the owner of the MCP server.`);
+            await client.discover({});
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
+            expect(consoleWarnSpy).not.toHaveBeenCalled();
             consoleWarnSpy.mockRestore();
         });
         it('should handle errors when discovering prompts', async () => {
-            const consoleErrorSpy = vi
-                .spyOn(console, 'error')
-                .mockImplementation(() => { });
             const mockedClient = {
                 connect: vi.fn(),
                 discover: vi.fn(),
@@ -117,20 +147,198 @@ describe('mcp-client', () => {
                 registerCapabilities: vi.fn(),
                 setRequestHandler: vi.fn(),
                 getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
-                request: vi.fn().mockRejectedValue(new Error('Test error')),
+                listTools: vi.fn().mockResolvedValue({ tools: [] }),
+                listPrompts: vi.fn().mockRejectedValue(new Error('Test error')),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const client = new McpClient('test-server', {
+                command: 'test-command',
+            }, mockedToolRegistry, {}, workspaceContext, false);
+            await client.connect();
+            await expect(client.discover({})).rejects.toThrow('No prompts or tools 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 () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
             };
             vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
             vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
-            vi.mocked(GenAiLib.mcpToTool).mockReturnValue({
-                tool: () => Promise.resolve({ functionDeclarations: [] }),
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                sortTools: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+            };
+            const client = new McpClient('test-server', {
+                command: 'test-command',
+            }, mockedToolRegistry, {}, workspaceContext, false);
+            await client.connect();
+            await expect(client.discover({})).rejects.toThrow('No prompts or tools found on the server.');
+        });
+        it('should discover tools if server supports them', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'testTool',
+                            description: 'A test tool',
+                            inputSchema: { type: 'object', properties: {} },
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            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 client = new McpClient('test-server', {
+                command: 'test-command',
+            }, mockedToolRegistry, {}, workspaceContext, false);
+            await client.connect();
+            await client.discover({});
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
+        });
+        it('should discover tools with $defs and $ref in schema', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                discover: vi.fn(),
+                disconnect: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'toolWithDefs',
+                            description: 'A tool using $defs',
+                            inputSchema: {
+                                type: 'object',
+                                properties: {
+                                    param1: {
+                                        $ref: '#/$defs/MyType',
+                                    },
+                                },
+                                $defs: {
+                                    MyType: {
+                                        type: 'string',
+                                        description: 'A defined type',
+                                    },
+                                },
+                            },
+                        },
+                    ],
+                }),
+                listPrompts: vi.fn().mockResolvedValue({
+                    prompts: [],
+                }),
+                request: vi.fn().mockResolvedValue({}),
+            };
+            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 client = new McpClient('test-server', {
+                command: 'test-command',
+            }, mockedToolRegistry, {}, workspaceContext, false);
+            await client.connect();
+            await client.discover({});
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
+            const registeredTool = vi.mocked(mockedToolRegistry.registerTool).mock
+                .calls[0][0];
+            expect(registeredTool.schema.parametersJsonSchema).toEqual({
+                type: 'object',
+                properties: {
+                    param1: {
+                        $ref: '#/$defs/MyType',
+                    },
+                },
+                $defs: {
+                    MyType: {
+                        type: 'string',
+                        description: 'A defined type',
+                    },
+                },
             });
+        });
+        it('should remove tools and prompts on disconnect', async () => {
+            const mockedClient = {
+                connect: vi.fn(),
+                close: vi.fn(),
+                getStatus: vi.fn(),
+                registerCapabilities: vi.fn(),
+                setRequestHandler: vi.fn(),
+                getServerCapabilities: vi
+                    .fn()
+                    .mockReturnValue({ tools: {}, prompts: {} }),
+                listPrompts: vi.fn().mockResolvedValue({
+                    prompts: [{ id: 'prompt1', text: 'a prompt' }],
+                }),
+                request: vi.fn().mockResolvedValue({}),
+                listTools: vi.fn().mockResolvedValue({
+                    tools: [
+                        {
+                            name: 'testTool',
+                            description: 'A test tool',
+                            inputSchema: { type: 'object', properties: {} },
+                        },
+                    ],
+                }),
+            };
+            vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
+            vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({});
+            const mockedToolRegistry = {
+                registerTool: vi.fn(),
+                unregisterTool: vi.fn(),
+                getMessageBus: vi.fn().mockReturnValue(undefined),
+                removeMcpToolsByServer: vi.fn(),
+                sortTools: vi.fn(),
+            };
+            const mockedPromptRegistry = {
+                registerPrompt: vi.fn(),
+                unregisterPrompt: vi.fn(),
+                removePromptsByServer: vi.fn(),
+            };
             const client = new McpClient('test-server', {
                 command: 'test-command',
-            }, {}, {}, {}, false);
+            }, mockedToolRegistry, mockedPromptRegistry, workspaceContext, false);
             await client.connect();
-            await expect(client.discover()).rejects.toThrow('No prompts or tools found on the server.');
-            expect(consoleErrorSpy).toHaveBeenCalledWith(`Error discovering prompts from test-server: Test error`);
-            consoleErrorSpy.mockRestore();
+            await client.discover({});
+            expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
+            expect(mockedPromptRegistry.registerPrompt).toHaveBeenCalledOnce();
+            await client.disconnect();
+            expect(mockedClient.close).toHaveBeenCalledOnce();
+            expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledOnce();
+            expect(mockedPromptRegistry.removePromptsByServer).toHaveBeenCalledOnce();
         });
     });
     describe('appendMcpServerCommand', () => {
@@ -158,18 +366,20 @@ describe('mcp-client', () => {
                 const transport = await createTransport('test-server', {
                     httpUrl: 'http://test-server',
                 }, false);
-                expect(transport).toEqual(new StreamableHTTPClientTransport(new URL('http://test-server'), {}));
+                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
             });
             it('with headers', async () => {
+                // We need this to be an any type because we dig into its private state.
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 const transport = await createTransport('test-server', {
                     httpUrl: 'http://test-server',
                     headers: { Authorization: 'derp' },
                 }, false);
-                expect(transport).toEqual(new StreamableHTTPClientTransport(new URL('http://test-server'), {
-                    requestInit: {
-                        headers: { Authorization: 'derp' },
-                    },
-                }));
+                expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
+                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
+                const authHeader = transport._requestInit?.headers?.['Authorization'];
+                expect(authHeader).toBe('derp');
             });
         });
         describe('should connect via url', () => {
@@ -177,18 +387,20 @@ describe('mcp-client', () => {
                 const transport = await createTransport('test-server', {
                     url: 'http://test-server',
                 }, false);
-                expect(transport).toEqual(new SSEClientTransport(new URL('http://test-server'), {}));
+                expect(transport).toBeInstanceOf(SSEClientTransport);
+                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
             });
             it('with headers', async () => {
+                // We need this to be an any type because we dig into its private state.
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 const transport = await createTransport('test-server', {
                     url: 'http://test-server',
                     headers: { Authorization: 'derp' },
                 }, false);
-                expect(transport).toEqual(new SSEClientTransport(new URL('http://test-server'), {
-                    requestInit: {
-                        headers: { Authorization: 'derp' },
-                    },
-                }));
+                expect(transport).toBeInstanceOf(SSEClientTransport);
+                expect(transport).toHaveProperty('_url', new URL('http://test-server/'));
+                const authHeader = transport._requestInit?.headers?.['Authorization'];
+                expect(authHeader).toBe('derp');
             });
         });
         it('should connect via command', async () => {
@@ -217,11 +429,17 @@ describe('mcp-client', () => {
                     oauth: {
                         scopes: ['scope1'],
                     },
+                    headers: {
+                        'X-Goog-User-Project': 'myproject',
+                    },
                 }, false);
                 expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
                 // eslint-disable-next-line @typescript-eslint/no-explicit-any
                 const authProvider = transport._authProvider;
                 expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                const googUserProject = transport._requestInit?.headers?.['X-Goog-User-Project'];
+                expect(googUserProject).toBe('myproject');
             });
             it('should use GoogleCredentialProvider with SSE transport', async () => {
                 const transport = await createTransport('test-server', {
@@ -282,148 +500,117 @@ describe('mcp-client', () => {
             expect(isEnabled(namelessFuncDecl, serverName, mcpServerConfig)).toBe(false);
         });
     });
-    describe('hasValidTypes', () => {
-        it('should return true for a valid schema with anyOf', () => {
-            const schema = {
-                anyOf: [{ type: 'string' }, { type: 'number' }],
-            };
-            expect(hasValidTypes(schema)).toBe(true);
-        });
-        it('should return false for an invalid schema with anyOf', () => {
-            const schema = {
-                anyOf: [{ type: 'string' }, { description: 'no type' }],
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return true for a valid schema with allOf', () => {
-            const schema = {
-                allOf: [
-                    { type: 'string' },
-                    { type: 'object', properties: { foo: { type: 'string' } } },
-                ],
-            };
-            expect(hasValidTypes(schema)).toBe(true);
-        });
-        it('should return false for an invalid schema with allOf', () => {
-            const schema = {
-                allOf: [{ type: 'string' }, { description: 'no type' }],
+    describe('hasNetworkTransport', () => {
+        it('should return true if only url is provided', () => {
+            const config = { url: 'http://example.com' };
+            expect(hasNetworkTransport(config)).toBe(true);
+        });
+        it('should return true if only httpUrl is provided', () => {
+            const config = { httpUrl: 'http://example.com' };
+            expect(hasNetworkTransport(config)).toBe(true);
+        });
+        it('should return true if both url and httpUrl are provided', () => {
+            const config = {
+                url: 'http://example.com/sse',
+                httpUrl: 'http://example.com/http',
             };
-            expect(hasValidTypes(schema)).toBe(false);
+            expect(hasNetworkTransport(config)).toBe(true);
         });
-        it('should return true for a valid schema with oneOf', () => {
-            const schema = {
-                oneOf: [{ type: 'string' }, { type: 'number' }],
-            };
-            expect(hasValidTypes(schema)).toBe(true);
+        it('should return false if neither url nor httpUrl is provided', () => {
+            const config = { command: 'do-something' };
+            expect(hasNetworkTransport(config)).toBe(false);
         });
-        it('should return false for an invalid schema with oneOf', () => {
-            const schema = {
-                oneOf: [{ type: 'string' }, { description: 'no type' }],
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return true for a valid schema with nested subschemas', () => {
-            const schema = {
-                anyOf: [
-                    { type: 'string' },
-                    {
-                        allOf: [
-                            { type: 'object', properties: { a: { type: 'string' } } },
-                            { type: 'object', properties: { b: { type: 'number' } } },
-                        ],
-                    },
-                ],
-            };
-            expect(hasValidTypes(schema)).toBe(true);
-        });
-        it('should return false for an invalid schema with nested subschemas', () => {
-            const schema = {
-                anyOf: [
-                    { type: 'string' },
-                    {
-                        allOf: [
-                            { type: 'object', properties: { a: { type: 'string' } } },
-                            { description: 'no type' },
-                        ],
-                    },
-                ],
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return true for a schema with a type and subschemas', () => {
-            const schema = {
-                type: 'string',
-                anyOf: [{ minLength: 1 }, { maxLength: 5 }],
-            };
-            expect(hasValidTypes(schema)).toBe(true);
-        });
-        it('should return false for a schema with no type and no subschemas', () => {
-            const schema = {
-                description: 'a schema with no type',
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return true for a valid schema', () => {
-            const schema = {
-                type: 'object',
-                properties: {
-                    param1: { type: 'string' },
-                },
-            };
-            expect(hasValidTypes(schema)).toBe(true);
-        });
-        it('should return false if a parameter is missing a type', () => {
-            const schema = {
-                type: 'object',
-                properties: {
-                    param1: { description: 'a param with no type' },
-                },
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return false if a nested parameter is missing a type', () => {
-            const schema = {
-                type: 'object',
-                properties: {
-                    param1: {
-                        type: 'object',
-                        properties: {
-                            nestedParam: {
-                                description: 'a nested param with no type',
-                            },
-                        },
-                    },
-                },
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return false if an array item is missing a type', () => {
-            const schema = {
-                type: 'object',
-                properties: {
-                    param1: {
-                        type: 'array',
-                        items: {
-                            description: 'an array item with no type',
-                        },
-                    },
-                },
-            };
-            expect(hasValidTypes(schema)).toBe(false);
-        });
-        it('should return true for a schema with no properties', () => {
-            const schema = {
-                type: 'object',
-            };
-            expect(hasValidTypes(schema)).toBe(true);
-        });
-        it('should return true for a schema with an empty properties object', () => {
-            const schema = {
-                type: 'object',
-                properties: {},
-            };
-            expect(hasValidTypes(schema)).toBe(true);
+        it('should return false for an empty config object', () => {
+            const config = {};
+            expect(hasNetworkTransport(config)).toBe(false);
         });
     });
 });
+describe('connectToMcpServer with OAuth', () => {
+    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);
+    });
+    afterEach(() => {
+        vi.clearAllMocks();
+    });
+    it('should handle automatic OAuth flow on 401 with www-authenticate header', async () => {
+        const serverUrl = 'http://test-server.com/';
+        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(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
+            authorizationUrl: authUrl,
+            tokenUrl,
+            scopes: ['test-scope'],
+        });
+        // We need this to be an any type because we dig into its private state.
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        let capturedTransport;
+        vi.mocked(mockedClient.connect).mockImplementationOnce(async (transport) => {
+            capturedTransport = transport;
+            return Promise.resolve();
+        });
+        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl }, false, workspaceContext);
+        expect(client).toBe(mockedClient);
+        expect(mockedClient.connect).toHaveBeenCalledTimes(2);
+        expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
+        const authHeader = capturedTransport._requestInit?.headers?.['Authorization'];
+        expect(authHeader).toBe('Bearer test-access-token');
+    });
+    it('should discover oauth config if not in www-authenticate header', async () => {
+        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(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
+            authorizationUrl: authUrl,
+            tokenUrl,
+            scopes: ['test-scope'],
+        });
+        vi.mocked(mockAuthProvider.getValidToken).mockResolvedValue('test-access-token-from-discovery');
+        // We need this to be an any type because we dig into its private state.
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        let capturedTransport;
+        vi.mocked(mockedClient.connect).mockImplementationOnce(async (transport) => {
+            capturedTransport = transport;
+            return Promise.resolve();
+        });
+        const client = await connectToMcpServer('test-server', { httpUrl: serverUrl }, false, workspaceContext);
+        expect(client).toBe(mockedClient);
+        expect(mockedClient.connect).toHaveBeenCalledTimes(2);
+        expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
+        expect(OAuthUtils.discoverOAuthConfig).toHaveBeenCalledWith(serverUrl);
+        const authHeader = capturedTransport._requestInit?.headers?.['Authorization'];
+        expect(authHeader).toBe('Bearer test-access-token-from-discovery');
+    });
+});
 //# sourceMappingURL=mcp-client.test.js.map