/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

/* eslint-disable @typescript-eslint/no-explicit-any */
import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js';
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js';
import {
  StreamableHTTPClientTransport,
  StreamableHTTPError,
} from '@modelcontextprotocol/sdk/client/streamableHttp.js';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { AuthProviderType, type Config } 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 type { PromptRegistry } from '../prompts/prompt-registry.js';
import {
  ErrorCode,
  McpError,
  PromptListChangedNotificationSchema,
  ResourceListChangedNotificationSchema,
  ToolListChangedNotificationSchema,
  ProgressNotificationSchema,
} from '@modelcontextprotocol/sdk/types.js';
import type { DiscoveredMCPTool } from './mcp-tool.js';

import { WorkspaceContext } from '../utils/workspaceContext.js';
import {
  connectToMcpServer,
  createTransport,
  hasNetworkTransport,
  isEnabled,
  McpClient,
  populateMcpServerCommand,
  discoverPrompts,
  type McpContext,
} from './mcp-client.js';
import type { ToolRegistry } from './tool-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
import * as fs from 'node:fs';
import * as os from 'node:os';
import * as path from 'node:path';
import { cleanupTmpDir } from '@google/gemini-cli-test-utils';
import { coreEvents } from '../utils/events.js';
import type { EnvironmentSanitizationConfig } from '../services/environmentSanitization.js';

interface TestableTransport {
  _authProvider?: GoogleCredentialProvider;
  _requestInit?: {
    headers?: Record<string, string>;
  };
}

const EMPTY_CONFIG: EnvironmentSanitizationConfig = {
  enableEnvironmentVariableRedaction: true,
  allowedEnvironmentVariables: [],
  blockedEnvironmentVariables: [],
};

const MOCK_CONTEXT_DEFAULT = {
  sanitizationConfig: EMPTY_CONFIG,
  emitMcpDiagnostic: vi.fn(),
  setUserInteractedWithMcp: vi.fn(),
  isTrustedFolder: vi.fn().mockReturnValue(true),
};

let MOCK_CONTEXT: McpContext = MOCK_CONTEXT_DEFAULT;

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(),
    emitConsoleLog: vi.fn(),
  },
}));

describe('mcp-client', () => {
  let workspaceContext: WorkspaceContext;
  let testWorkspace: string;

  beforeEach(() => {
    MOCK_CONTEXT = {
      sanitizationConfig: EMPTY_CONFIG,
      emitMcpDiagnostic: vi.fn(),
      setUserInteractedWithMcp: vi.fn(),
      isTrustedFolder: vi.fn().mockReturnValue(true),
    };
    // 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(async () => {
    vi.useRealTimers();
    await cleanupTmpDir(testWorkspace);
    workspaceContext = null as unknown as WorkspaceContext;
    vi.restoreAllMocks();
  });

  describe('McpClient', () => {
    it('should discover tools', 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: 'testFunction',
              inputSchema: {
                type: 'object',
                properties: {},
              },
            },
          ],
        }),
        listPrompts: vi.fn().mockResolvedValue({
          prompts: [],
        }),
        request: vi.fn().mockResolvedValue({}),
      };
      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });
      expect(mockedClient.listTools).toHaveBeenCalledWith(
        {},
        expect.objectContaining({ timeout: 600000, progressReporter: client }),
      );
    });

    it('should not skip tools even if a parameter is missing a type', async () => {
      const consoleWarnSpy = vi
        .spyOn(console, 'warn')
        .mockImplementation(() => {});
      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: 'validTool',
              inputSchema: {
                type: 'object',
                properties: {
                  param1: { type: 'string' },
                },
              },
            },
            {
              name: 'invalidTool',
              inputSchema: {
                type: 'object',
                properties: {
                  param1: { description: 'a param with no type' },
                },
              },
            },
          ],
        }),
        listPrompts: vi.fn().mockResolvedValue({
          prompts: [],
        }),
        request: vi.fn().mockResolvedValue({}),
      };
      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });
      expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
      expect(consoleWarnSpy).not.toHaveBeenCalled();
      consoleWarnSpy.mockRestore();
    });

    it('should propagate errors when discovering prompts', 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({ prompts: {} }),
        listTools: vi.fn().mockResolvedValue({ tools: [] }),
        listPrompts: vi.fn().mockRejectedValue(new Error('Test error')),
        request: vi.fn().mockResolvedValue({}),
      };
      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await expect(
        client.discoverInto(MOCK_CONTEXT, {
          toolRegistry: mockedToolRegistry,
          promptRegistry,
          resourceRegistry,
        }),
      ).rejects.toThrow('Test error');
      expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith(
        'error',
        `Error discovering prompts from test-server: Test error`,
        expect.any(Error),
        'test-server',
      );
    });

    it('should return empty array for discoverPrompts on MethodNotFound error without diagnostic', async () => {
      const mockedClient = {
        getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
        listPrompts: vi
          .fn()
          .mockRejectedValue(
            new McpError(ErrorCode.MethodNotFound, 'Method not supported'),
          ),
      };
      const result = await discoverPrompts(
        'test-server',
        mockedClient as unknown as ClientLib.Client,
        MOCK_CONTEXT,
      );
      expect(result).toEqual([]);
      // MethodNotFound errors should be silently ignored regardless of message text
      expect(MOCK_CONTEXT.emitMcpDiagnostic).not.toHaveBeenCalled();
    });

    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(),
        setNotificationHandler: vi.fn(),
        getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
        listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
        request: vi.fn().mockResolvedValue({}),
      };
      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await expect(
        client.discoverInto(MOCK_CONTEXT, {
          toolRegistry: mockedToolRegistry,
          promptRegistry,
          resourceRegistry,
        }),
      ).rejects.toThrow('No prompts, tools, or resources found on the server.');
    });

    it('should discover tools if server supports them', async () => {
      const mockedClient = {
        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: 'testTool',
              description: 'A test tool',
              inputSchema: { type: 'object', properties: {} },
            },
          ],
        }),
        listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
        request: vi.fn().mockResolvedValue({}),
      };
      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });
      expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
    });

    it('should register tool with readOnlyHint and preserve annotations', 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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockPolicyEngine = {
        addRule: vi.fn(),
      };
      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      } as unknown as Config;

      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getMessageBus: vi.fn().mockReturnValue(undefined),
        removeMcpToolsByServer: vi.fn(),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      await client.discoverInto(mockConfig, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });

      // Verify tool registration
      expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();

      // Verify addRule is NOT called (annotation-based rules are in plan.toml now)
      expect(mockPolicyEngine.addRule).not.toHaveBeenCalled();

      // Verify annotations are preserved on the registered tool
      const registeredTool = (
        mockedToolRegistry.registerTool as ReturnType<typeof vi.fn>
      ).mock.calls[0][0] as DiscoveredMCPTool;
      expect(registeredTool.toolAnnotations).toEqual({ readOnlyHint: true });
      expect(registeredTool.isReadOnly).toBe(true);
    });

    it('should preserve undefined annotations 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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockPolicyEngine = {
        addRule: vi.fn(),
      };
      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
      } as unknown as Config;

      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getMessageBus: vi.fn().mockReturnValue(undefined),
        removeMcpToolsByServer: vi.fn(),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      await client.discoverInto(mockConfig, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });

      expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
      expect(mockPolicyEngine.addRule).not.toHaveBeenCalled();

      // Verify annotations are undefined for tools without annotations
      const registeredTool = (
        mockedToolRegistry.registerTool as ReturnType<typeof vi.fn>
      ).mock.calls[0][0] as DiscoveredMCPTool;
      expect(registeredTool.toolAnnotations).toBeUndefined();
    });

    it('should preserve full annotations object with multiple hints', 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: 'multiAnnotationTool',
              description: 'A tool with multiple annotations',
              inputSchema: { type: 'object', properties: {} },
              annotations: {
                readOnlyHint: true,
                destructiveHint: false,
                idempotentHint: true,
              },
            },
          ],
        }),
        listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
        request: vi.fn().mockResolvedValue({}),
      };
      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockConfig = {
        getPolicyEngine: vi.fn().mockReturnValue({ addRule: vi.fn() }),
      } as unknown as Config;

      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getMessageBus: vi.fn().mockReturnValue(undefined),
        removeMcpToolsByServer: vi.fn(),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      await client.discoverInto(mockConfig, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });

      expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();

      const registeredTool = (
        mockedToolRegistry.registerTool as ReturnType<typeof vi.fn>
      ).mock.calls[0][0] as DiscoveredMCPTool;
      expect(registeredTool.toolAnnotations).toEqual({
        readOnlyHint: true,
        destructiveHint: false,
        idempotentHint: true,
      });
      expect(registeredTool.isReadOnly).toBe(true);
    });

    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(),
        setNotificationHandler: 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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });
      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',
          },
          wait_for_previous: {
            type: 'boolean',
            description:
              'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
          },
        },
        $defs: {
          MyType: {
            type: 'string',
            description: 'A defined type',
          },
        },
      });
    });

    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: [] });
        }),
      } as unknown as ClientLib.Client;
      vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });
      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:
        | ((notification: unknown) => Promise<void> | void)
        | undefined;
      const mockedClient = {
        connect: vi.fn(),
        discover: vi.fn(),
        disconnect: vi.fn(),
        getStatus: vi.fn(),
        registerCapabilities: vi.fn(),
        setRequestHandler: vi.fn(),
        setNotificationHandler: vi.fn((schema, handler) => {
          if (schema === ResourceListChangedNotificationSchema) {
            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: [] });
        }),
      } as unknown as ClientLib.Client;
      vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });

      expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2);
      expect(resourceListHandler).toBeDefined();

      await resourceListHandler?.({
        method: 'notifications/resources/list_changed',
      });

      expect(resourceRegistry.setResourcesForServer).toHaveBeenLastCalledWith(
        'test-server',
        [expect.objectContaining({ uri: 'file:///tmp/two.txt' })],
      );

      expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith(
        'info',
        'Resources updated for server: test-server',
        undefined,
        'test-server',
      );
    });

    it('refreshes prompts when prompt list change notification is received', async () => {
      let listCallCount = 0;
      let promptListHandler:
        | ((notification: unknown) => Promise<void> | void)
        | undefined;
      const mockedClient = {
        connect: vi.fn(),
        discover: vi.fn(),
        disconnect: vi.fn(),
        getStatus: vi.fn(),
        registerCapabilities: vi.fn(),
        setRequestHandler: vi.fn(),
        setNotificationHandler: vi.fn((schema, handler) => {
          if (schema === PromptListChangedNotificationSchema) {
            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: [] }),
      } as unknown as ClientLib.Client;
      vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;
      const promptRegistry = {
        registerPrompt: vi.fn(),
        getPromptsByServer: vi.fn().mockReturnValue([]),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        getResourcesByServer: vi.fn().mockReturnValue([]),
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry,
        resourceRegistry,
      });

      expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2);
      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(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith(
        'info',
        'Prompts updated for server: test-server',
        undefined,
        'test-server',
      );
    });

    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(),
        setNotificationHandler: 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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );
      const mockedToolRegistry = {
        registerTool: vi.fn(),
        unregisterTool: vi.fn(),
        getMessageBus: vi.fn().mockReturnValue(undefined),
        removeMcpToolsByServer: vi.fn(),
        sortTools: vi.fn(),
      } as unknown as ToolRegistry;
      const mockedPromptRegistry = {
        registerPrompt: vi.fn(),
        unregisterPrompt: vi.fn(),
        removePromptsByServer: vi.fn(),
      } as unknown as PromptRegistry;
      const resourceRegistry = {
        setResourcesForServer: vi.fn(),
        removeResourcesByServer: vi.fn(),
      } as unknown as ResourceRegistry;
      const client = new McpClient(
        'test-server',
        {
          command: 'test-command',
        },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );
      await client.connect();
      await client.discoverInto(MOCK_CONTEXT, {
        toolRegistry: mockedToolRegistry,
        promptRegistry: mockedPromptRegistry,
        resourceRegistry,
      });

      expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
      expect(mockedPromptRegistry.registerPrompt).toHaveBeenCalledOnce();

      await client.disconnect();

      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: [] }),
        listResources: vi.fn().mockResolvedValue({ resources: [] }),
        request: vi.fn().mockResolvedValue({}),
      };

      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockedToolRegistry = {
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {
          getPromptsByServer: vi.fn().mockReturnValue([]),
          registerPrompt: vi.fn(),
        } as unknown as PromptRegistry,
        resourceRegistry: {
          getResourcesByServer: vi.fn().mockReturnValue([]),
          registerResource: vi.fn(),
          removeResourcesByServer: vi.fn(),
          setResourcesForServer: vi.fn(),
        } as unknown as ResourceRegistry,
      });

      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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: {
          getToolsByServer: vi.fn().mockReturnValue([]),
          registerTool: vi.fn(),
          sortTools: vi.fn(),
        } as unknown as ToolRegistry,
        promptRegistry: {
          getPromptsByServer: vi.fn().mockReturnValue([]),
          registerPrompt: vi.fn(),
        } as unknown as PromptRegistry,
        resourceRegistry: {
          getResourcesByServer: vi.fn().mockReturnValue([]),
          registerResource: vi.fn(),
          removeResourcesByServer: vi.fn(),
          setResourcesForServer: vi.fn(),
        } as unknown as ResourceRegistry,
      });

      // Should be called for ProgressNotificationSchema, even if no other capabilities
      expect(mockedClient.setNotificationHandler).toHaveBeenCalled();
      const progressCall = mockedClient.setNotificationHandler.mock.calls.find(
        (call) => call[0] === ProgressNotificationSchema,
      );
      expect(progressCall).toBeDefined();
    });

    it('should set up notification handler even if listChanged is false (robustness)', async () => {
      // Setup mocks
      const mockedClient = {
        connect: vi.fn(),
        getServerCapabilities: vi
          .fn()
          .mockReturnValue({ tools: { listChanged: false } }),
        setNotificationHandler: vi.fn(),
        request: vi.fn().mockResolvedValue({}),
        registerCapabilities: vi.fn().mockResolvedValue({}),
        setRequestHandler: vi.fn().mockResolvedValue({}),
      };

      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: {
          getToolsByServer: vi.fn().mockReturnValue([]),
          registerTool: vi.fn(),
          sortTools: vi.fn(),
        } as unknown as ToolRegistry,
        promptRegistry: {
          getPromptsByServer: vi.fn().mockReturnValue([]),
          registerPrompt: vi.fn(),
        } as unknown as PromptRegistry,
        resourceRegistry: {
          getResourcesByServer: vi.fn().mockReturnValue([]),
          registerResource: vi.fn(),
          removeResourcesByServer: vi.fn(),
          setResourcesForServer: vi.fn(),
        } as unknown as ResourceRegistry,
      });

      const toolUpdateCall =
        mockedClient.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      expect(toolUpdateCall).toBeDefined();
    });

    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: [] }),
        listResources: vi.fn().mockResolvedValue({ resources: [] }),
        request: vi.fn().mockResolvedValue({}),
        registerCapabilities: vi.fn().mockResolvedValue({}),
        setRequestHandler: vi.fn().mockResolvedValue({}),
      };

      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockedToolRegistry = {
        removeMcpToolsByServer: vi.fn(),
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;

      const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);

      // Initialize client with onContextUpdated callback
      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
        onContextUpdatedSpy,
      );

      // 1. Connect (sets up listener)
      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {} as PromptRegistry,
        resourceRegistry: {
          removeMcpResourcesByServer: vi.fn(),
          registerResource: vi.fn(),
        } as unknown as ResourceRegistry,
      });

      // 2. Extract the callback passed to setNotificationHandler for tools
      const toolUpdateCall =
        mockedClient.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      const notificationCallback = toolUpdateCall![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(onContextUpdatedSpy).toHaveBeenCalled();

      // It should emit feedback event
      expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith(
        'info',
        'Tools updated for server: test-server',
        undefined,
        '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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockedToolRegistry = {
        removeMcpToolsByServer: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {} as PromptRegistry,
        resourceRegistry: {} as ResourceRegistry,
      });

      const toolUpdateCall =
        mockedClient.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      const notificationCallback = toolUpdateCall![1];

      // Trigger notification - should fail internally but catch the error
      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: string) => ({
        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 as unknown as ClientLib.Client)
        .mockReturnValueOnce(mockClientB as unknown as ClientLib.Client);

      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockedToolRegistry = {
        removeMcpToolsByServer: vi.fn(),
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;

      const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);

      const clientA = new McpClient(
        'server-A',
        { command: 'cmd-a' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
        onContextUpdatedSpy,
      );

      const clientB = new McpClient(
        'server-B',
        { command: 'cmd-b' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
        onContextUpdatedSpy,
      );

      await clientA.connect();
      // INJECTED REGISTRIES
      (clientA as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {} as PromptRegistry,
        resourceRegistry: {} as ResourceRegistry,
      });
      await clientB.connect();
      // INJECTED REGISTRIES
      (clientB as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {} as PromptRegistry,
        resourceRegistry: {} as ResourceRegistry,
      });

      const toolUpdateCallA =
        mockClientA.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      const handlerA = toolUpdateCallA![1];

      const toolUpdateCallB =
        mockClientB.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      const handlerB = toolUpdateCallB![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(onContextUpdatedSpy).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'));
                },
                { once: true },
              );
            }),
        ),
        listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
        listResources: vi.fn().mockResolvedValue({ resources: [] }),
        request: vi.fn().mockResolvedValue({}),
        registerCapabilities: vi.fn().mockResolvedValue({}),
        setRequestHandler: vi.fn().mockResolvedValue({}),
      };

      vi.mocked(ClientLib.Client).mockReturnValue(
        mockedClient as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockedToolRegistry = {
        removeMcpToolsByServer: vi.fn(),
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;

      const client = new McpClient(
        'test-server',
        // Set a very short timeout
        { command: 'test-command', timeout: 50 },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
      );

      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {
          getPromptsByServer: vi.fn().mockReturnValue([]),
          registerPrompt: vi.fn(),
          removePromptsByServer: vi.fn(),
        } as unknown as PromptRegistry,
        resourceRegistry: {
          getResourcesByServer: vi.fn().mockReturnValue([]),
          registerResource: vi.fn(),
          removeResourcesByServer: vi.fn(),
          setResourcesForServer: vi.fn(),
        } as unknown as ResourceRegistry,
      });

      const toolUpdateCall =
        mockedClient.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      const notificationCallback = toolUpdateCall![1];

      const refreshPromise = notificationCallback();

      // Advance timers to trigger the timeout (11 minutes to cover even the default timeout)
      await vi.advanceTimersByTimeAsync(11 * 60 * 1000);
      await refreshPromise;

      expect(mockedClient.listTools).toHaveBeenCalledWith(
        expect.anything(),
        expect.objectContaining({
          signal: expect.any(AbortSignal),
        }),
      );

      expect(mockedToolRegistry.registerTool).not.toHaveBeenCalled();
    });

    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 as unknown as ClientLib.Client,
      );
      vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
        {} as SdkClientStdioLib.StdioClientTransport,
      );

      const mockedToolRegistry = {
        removeMcpToolsByServer: vi.fn(),
        registerTool: vi.fn(),
        sortTools: vi.fn(),
        getToolsByServer: vi.fn().mockReturnValue([]),
        getMessageBus: vi.fn().mockReturnValue(undefined),
      } as unknown as ToolRegistry;

      const onContextUpdatedSpy = vi.fn().mockResolvedValue(undefined);

      const client = new McpClient(
        'test-server',
        { command: 'test-command' },
        workspaceContext,
        MOCK_CONTEXT,
        false,
        '0.0.1',
        onContextUpdatedSpy,
      );

      await client.connect();
      // INJECTED REGISTRIES
      (client as any).registeredRegistries?.add({
        toolRegistry: mockedToolRegistry,
        promptRegistry: {
          getPromptsByServer: vi.fn().mockReturnValue([]),
          registerPrompt: vi.fn(),
          removePromptsByServer: vi.fn(),
        } as unknown as PromptRegistry,
        resourceRegistry: {
          getResourcesByServer: vi.fn().mockReturnValue([]),
          registerResource: vi.fn(),
          removeResourcesByServer: vi.fn(),
          setResourcesForServer: vi.fn(),
        } as unknown as ResourceRegistry,
      });

      const toolUpdateCall =
        mockedClient.setNotificationHandler.mock.calls.find(
          (call) => call[0] === ToolListChangedNotificationSchema,
        );
      const notificationCallback = toolUpdateCall![1];

      vi.useFakeTimers();
      const refreshPromise = notificationCallback();
      await vi.advanceTimersByTimeAsync(500);
      await refreshPromise;

      expect(onContextUpdatedSpy).toHaveBeenCalledWith(expect.any(AbortSignal));

      // Verify the signal passed was not aborted (happy path)
      const signal = onContextUpdatedSpy.mock.calls[0][0];
      expect(signal.aborted).toBe(false);
    });
  });

  describe('appendMcpServerCommand', () => {
    it('should do nothing if no MCP servers or command are configured', () => {
      const out = populateMcpServerCommand({}, undefined);
      expect(out).toEqual({});
    });

    it('should discover tools via mcpServerCommand', () => {
      const commandString = 'command --arg1 value1';
      const out = populateMcpServerCommand({}, commandString);
      expect(out).toEqual({
        mcp: {
          command: 'command',
          args: ['--arg1', 'value1'],
        },
      });
    });

    it('should handle error if mcpServerCommand parsing fails', () => {
      expect(() => populateMcpServerCommand({}, 'derp && herp')).toThrowError();
    });
  });

  describe('createTransport', () => {
    describe('should connect via httpUrl', () => {
      it('without headers', async () => {
        const transport = await createTransport(
          'test-server',
          {
            httpUrl: 'http://test-server',
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: { headers: {} },
        });
      });

      it('with headers', async () => {
        const transport = await createTransport(
          'test-server',
          {
            httpUrl: 'http://test-server',
            headers: { Authorization: 'derp' },
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: {
            headers: { Authorization: 'derp' },
          },
        });
      });

      it('wraps fetch to convert GET 404 to 405 for POST-only servers (e.g. n8n)', async () => {
        const mockFetch = vi
          .fn()
          .mockResolvedValue(
            new Response(null, { status: 404, statusText: 'Not Found' }),
          );
        vi.stubGlobal('fetch', mockFetch);

        try {
          const transport = await createTransport(
            'test-server',
            { httpUrl: 'http://test-server' },
            false,
            MOCK_CONTEXT,
          );

          const wrappedFetch = (
            transport as unknown as {
              _fetch: (
                url: URL | string,
                init?: RequestInit,
              ) => Promise<Response>;
            }
          )._fetch;

          // GET 404 → 405: server doesn't support optional SSE GET stream
          const getRes = await wrappedFetch('http://test-server', {
            method: 'GET',
          });
          expect(getRes.status).toBe(405);
          expect(getRes.statusText).toBe('Method Not Allowed');

          // POST 404 → unchanged: real "not found" errors must still propagate
          const postRes = await wrappedFetch('http://test-server', {
            method: 'POST',
          });
          expect(postRes.status).toBe(404);
        } finally {
          vi.unstubAllGlobals();
        }
      });
    });

    describe('should connect via url', () => {
      it('without headers', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
          },
          false,
          MOCK_CONTEXT,
        );
        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: { headers: {} },
        });
      });

      it('with headers', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
            headers: { Authorization: 'derp' },
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: {
            headers: { Authorization: 'derp' },
          },
        });
      });

      it('with type="http" creates StreamableHTTPClientTransport', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
            type: 'http',
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: { headers: {} },
        });
      });

      it('with type="sse" creates SSEClientTransport', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
            type: 'sse',
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(SSEClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: { headers: {} },
        });
      });

      it('without type defaults to StreamableHTTPClientTransport', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: { headers: {} },
        });
      });

      it('with type="http" and headers applies headers correctly', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
            type: 'http',
            headers: { Authorization: 'Bearer token' },
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: {
            headers: { Authorization: 'Bearer token' },
          },
        });
      });

      it('with type="sse" and headers applies headers correctly', async () => {
        const transport = await createTransport(
          'test-server',
          {
            url: 'http://test-server',
            type: 'sse',
            headers: { 'X-API-Key': 'key123' },
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(SSEClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server'),
          _requestInit: {
            headers: { 'X-API-Key': 'key123' },
          },
        });
      });

      it('httpUrl takes priority over url when both are present', async () => {
        const transport = await createTransport(
          'test-server',
          {
            httpUrl: 'http://test-server-http',
            url: 'http://test-server-url',
          },
          false,
          MOCK_CONTEXT,
        );

        // httpUrl should take priority and create HTTP transport
        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(transport).toMatchObject({
          _url: new URL('http://test-server-http'),
          _requestInit: { headers: {} },
        });
      });
    });

    it('should connect via command', async () => {
      const mockedTransport = vi
        .spyOn(SdkClientStdioLib, 'StdioClientTransport')
        .mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

      await createTransport(
        'test-server',
        {
          command: 'test-command',
          args: ['--foo', 'bar'],
          env: { FOO: 'bar' },
          cwd: 'test/cwd',
        },
        false,
        MOCK_CONTEXT,
      );

      expect(mockedTransport).toHaveBeenCalledWith({
        command: 'test-command',
        args: ['--foo', 'bar'],
        cwd: 'test/cwd',
        env: expect.objectContaining({ FOO: 'bar' }),
        stderr: 'pipe',
      });
    });

    it('sets an env variable GEMINI_CLI=1 for stdio MCP servers', async () => {
      const mockedTransport = vi
        .spyOn(SdkClientStdioLib, 'StdioClientTransport')
        .mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

      await createTransport(
        'test-server',
        {
          command: 'test-command',
          args: ['--foo', 'bar'],
          env: {},
          cwd: 'test/cwd',
        },
        false,
        MOCK_CONTEXT,
      );

      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({} as SdkClientStdioLib.StdioClientTransport);

      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,
        MOCK_CONTEXT,
      );

      const callArgs = mockedTransport.mock.calls[0][0];
      expect(callArgs.env).toBeDefined();
      expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBeUndefined();
    });

    it('should include extension settings with defined values in environment', async () => {
      const mockedTransport = vi
        .spyOn(SdkClientStdioLib, 'StdioClientTransport')
        .mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

      await createTransport(
        'test-server',
        {
          command: 'test-command',
          extension: {
            name: 'test-ext',
            resolvedSettings: [
              {
                envVar: 'GEMINI_CLI_EXT_VAR',
                value: 'defined-value',
                sensitive: false,
                name: 'ext-setting',
              },
            ],
            version: '',
            isActive: false,
            path: '',
            contextFiles: [],
            id: '',
          },
        },
        false,
        MOCK_CONTEXT,
      );

      const callArgs = mockedTransport.mock.calls[0][0];
      expect(callArgs.env).toBeDefined();
      expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBe('defined-value');
    });

    it('should resolve environment variables in mcpServerConfig.env using extension settings', async () => {
      const mockedTransport = vi
        .spyOn(SdkClientStdioLib, 'StdioClientTransport')
        .mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

      await createTransport(
        'test-server',
        {
          command: 'test-command',
          env: {
            RESOLVED_VAR: '$GEMINI_CLI_EXT_VAR',
          },
          extension: {
            name: 'test-ext',
            resolvedSettings: [
              {
                envVar: 'GEMINI_CLI_EXT_VAR',
                value: 'ext-value',
                sensitive: false,
                name: 'ext-setting',
              },
            ],
            version: '',
            isActive: false,
            path: '',
            contextFiles: [],
            id: '',
          },
        },
        false,
        MOCK_CONTEXT,
      );

      const callArgs = mockedTransport.mock.calls[0][0];
      expect(callArgs.env).toBeDefined();
      expect(callArgs.env!['GEMINI_CLI_EXT_VAR']).toBe('ext-value');
      expect(callArgs.env!['RESOLVED_VAR']).toBe('ext-value');
    });
    it('should expand environment variables in mcpServerConfig.env and not redact them', async () => {
      const mockedTransport = vi
        .spyOn(SdkClientStdioLib, 'StdioClientTransport')
        .mockReturnValue({} as SdkClientStdioLib.StdioClientTransport);

      const originalEnv = process.env;
      process.env = {
        ...originalEnv,
        GEMINI_TEST_VAR: 'expanded-value',
      };

      try {
        await createTransport(
          'test-server',
          {
            command: 'test-command',
            env: {
              TEST_EXPANDED: 'Value is $GEMINI_TEST_VAR',
              SECRET_KEY: 'intentional-secret-123',
            },
          },
          false,
          MOCK_CONTEXT,
        );

        const callArgs = mockedTransport.mock.calls[0][0];
        expect(callArgs.env).toBeDefined();
        expect(callArgs.env!['TEST_EXPANDED']).toBe('Value is expanded-value');
        expect(callArgs.env!['SECRET_KEY']).toBe('intentional-secret-123');
      } finally {
        process.env = originalEnv;
      }
    });

    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',
            authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
            oauth: {
              scopes: ['scope1'],
            },
            headers: {
              'X-Goog-User-Project': 'myproject',
            },
          },
          false,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        const testableTransport = transport as unknown as TestableTransport;
        const authProvider = testableTransport._authProvider;
        expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
        const googUserProject =
          testableTransport._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,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        expect(mockGetRequestHeaders).toHaveBeenCalled();
        const testableTransport = transport as unknown as TestableTransport;
        const headers = testableTransport._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,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
        const testableTransport = transport as unknown as TestableTransport;
        const headers = testableTransport._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,
          MOCK_CONTEXT,
        );

        expect(transport).toBeInstanceOf(SSEClientTransport);
        const testableTransport = transport as unknown as TestableTransport;
        const authProvider = testableTransport._authProvider;
        expect(authProvider).toBeInstanceOf(GoogleCredentialProvider);
      });

      it('should throw an error if no URL is provided with GoogleCredentialProvider', async () => {
        await expect(
          createTransport(
            'test-server',
            {
              authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
              oauth: {
                scopes: ['scope1'],
              },
            },
            false,
            MOCK_CONTEXT,
          ),
        ).rejects.toThrow(
          'URL must be provided in the config for Google Credentials provider',
        );
      });
    });
  });
  describe('isEnabled', () => {
    const funcDecl = { name: 'myTool' };
    const serverName = 'myServer';

    it('should return true if no include or exclude lists are provided', () => {
      const mcpServerConfig = {};
      expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(true);
    });

    it('should return false if the tool is in the exclude list', () => {
      const mcpServerConfig = { excludeTools: ['myTool'] };
      expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(false);
    });

    it('should return true if the tool is in the include list', () => {
      const mcpServerConfig = { includeTools: ['myTool'] };
      expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(true);
    });

    it('should return true if the tool is in the include list with parentheses', () => {
      const mcpServerConfig = { includeTools: ['myTool()'] };
      expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(true);
    });

    it('should return false if the include list exists but does not contain the tool', () => {
      const mcpServerConfig = { includeTools: ['anotherTool'] };
      expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(false);
    });

    it('should return false if the tool is in both the include and exclude lists', () => {
      const mcpServerConfig = {
        includeTools: ['myTool'],
        excludeTools: ['myTool'],
      };
      expect(isEnabled(funcDecl, serverName, mcpServerConfig)).toBe(false);
    });

    it('should return false if the function declaration has no name', () => {
      const namelessFuncDecl = {};
      const mcpServerConfig = {};
      expect(isEnabled(namelessFuncDecl, serverName, mcpServerConfig)).toBe(
        false,
      );
    });
  });

  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(hasNetworkTransport(config)).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 empty config object', () => {
      const config = {};
      expect(hasNetworkTransport(config)).toBe(false);
    });
  });
});

describe('connectToMcpServer with OAuth', () => {
  let mockedClient: ClientLib.Client;
  let workspaceContext: WorkspaceContext;
  let testWorkspace: string;
  let mockAuthProvider: MCPOAuthProvider;
  let mockTokenStorage: MCPOAuthTokenStorage;

  beforeEach(() => {
    mockedClient = {
      connect: vi.fn(),
      close: vi.fn(),
      registerCapabilities: vi.fn(),
      setRequestHandler: vi.fn(),
      onclose: vi.fn(),
      notification: vi.fn(),
    } as unknown as ClientLib.Client;
    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' }),
    } as unknown as MCPOAuthTokenStorage;
    vi.mocked(MCPOAuthTokenStorage).mockReturnValue(mockTokenStorage);
    mockAuthProvider = {
      authenticate: vi.fn().mockResolvedValue(undefined),
      getValidToken: vi.fn().mockResolvedValue('test-access-token'),
      tokenStorage: mockTokenStorage,
    } as unknown as MCPOAuthProvider;
    vi.mocked(MCPOAuthProvider).mockReturnValue(mockAuthProvider);
  });

  afterEach(async () => {
    vi.useRealTimers();
    await cleanupTmpDir(testWorkspace);
    workspaceContext = null as unknown as WorkspaceContext;
    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 StreamableHTTPError(
        401,
        `Unauthorized\nwww-authenticate: ${wwwAuthHeader}`,
      ),
    );

    vi.mocked(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
      authorizationUrl: authUrl,
      tokenUrl,
      scopes: ['test-scope'],
    });

    // We need this to be typed to dig into its private state.
    let capturedTransport: TestableTransport | undefined;
    vi.mocked(mockedClient.connect).mockImplementationOnce(
      async (transport) => {
        capturedTransport = transport as unknown as TestableTransport;
        return Promise.resolve();
      },
    );

    const client = await connectToMcpServer(
      '0.0.1',
      'test-server',
      { httpUrl: serverUrl, oauth: { enabled: true } },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    expect(client).toBe(mockedClient);
    expect(mockedClient.connect).toHaveBeenCalledTimes(2);
    expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();

    const authHeader = (capturedTransport as TestableTransport)._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 StreamableHTTPError(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 typed to dig into its private state.
    let capturedTransport: TestableTransport | undefined;
    vi.mocked(mockedClient.connect).mockImplementationOnce(
      async (transport) => {
        capturedTransport = transport as unknown as TestableTransport;
        return Promise.resolve();
      },
    );

    const client = await connectToMcpServer(
      '0.0.1',
      'test-server',
      { httpUrl: serverUrl, oauth: { enabled: true } },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    expect(client).toBe(mockedClient);
    expect(mockedClient.connect).toHaveBeenCalledTimes(2);
    expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
    expect(OAuthUtils.discoverOAuthConfig).toHaveBeenCalledWith(serverUrl);

    const authHeader = (capturedTransport as TestableTransport)._requestInit
      ?.headers?.['Authorization'];
    expect(authHeader).toBe('Bearer test-access-token-from-discovery');
  });

  it('should use discoverOAuthFromWWWAuthenticate when it succeeds and skip discoverOAuthConfig', async () => {
    const serverUrl = 'http://test-server.com/mcp';
    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 StreamableHTTPError(
        401,
        `Unauthorized\nwww-authenticate: ${wwwAuthHeader}`,
      ),
    );

    vi.mocked(OAuthUtils.discoverOAuthFromWWWAuthenticate).mockResolvedValue({
      authorizationUrl: authUrl,
      tokenUrl,
      scopes: ['read'],
    });

    vi.mocked(mockedClient.connect).mockResolvedValueOnce(undefined);

    const client = await connectToMcpServer(
      '0.0.1',
      'test-server',
      { httpUrl: serverUrl, oauth: { enabled: true } },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    expect(client).toBe(mockedClient);
    expect(OAuthUtils.discoverOAuthFromWWWAuthenticate).toHaveBeenCalledWith(
      wwwAuthHeader,
      serverUrl,
    );
    expect(OAuthUtils.discoverOAuthConfig).not.toHaveBeenCalled();
    expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
  });

  it('should fall back to extractBaseUrl + discoverOAuthConfig when discoverOAuthFromWWWAuthenticate returns null', async () => {
    const serverUrl = 'http://test-server.com/mcp';
    const baseUrl = 'http://test-server.com';
    const authUrl = 'http://auth.example.com/auth';
    const tokenUrl = 'http://auth.example.com/token';
    const wwwAuthHeader = `Bearer realm="test"`;

    vi.mocked(mockedClient.connect).mockRejectedValueOnce(
      new StreamableHTTPError(
        401,
        `Unauthorized\nwww-authenticate: ${wwwAuthHeader}`,
      ),
    );

    vi.mocked(OAuthUtils.discoverOAuthFromWWWAuthenticate).mockResolvedValue(
      null,
    );
    vi.mocked(OAuthUtils.extractBaseUrl).mockReturnValue(baseUrl);
    vi.mocked(OAuthUtils.discoverOAuthConfig).mockResolvedValue({
      authorizationUrl: authUrl,
      tokenUrl,
      scopes: ['read'],
    });

    vi.mocked(mockedClient.connect).mockResolvedValueOnce(undefined);

    const client = await connectToMcpServer(
      '0.0.1',
      'test-server',
      { httpUrl: serverUrl, oauth: { enabled: true } },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    expect(client).toBe(mockedClient);
    expect(OAuthUtils.discoverOAuthFromWWWAuthenticate).toHaveBeenCalledWith(
      wwwAuthHeader,
      serverUrl,
    );
    expect(OAuthUtils.extractBaseUrl).toHaveBeenCalledWith(serverUrl);
    expect(OAuthUtils.discoverOAuthConfig).toHaveBeenCalledWith(baseUrl);
    expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
  });
});

describe('connectToMcpServer - HTTP→SSE fallback', () => {
  let mockedClient: ClientLib.Client;
  let workspaceContext: WorkspaceContext;
  let testWorkspace: string;

  beforeEach(() => {
    mockedClient = {
      connect: vi.fn(),
      close: vi.fn(),
      registerCapabilities: vi.fn(),
      setRequestHandler: vi.fn(),
      onclose: vi.fn(),
      notification: vi.fn(),
    } as unknown as ClientLib.Client;
    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(async () => {
    vi.useRealTimers();
    await cleanupTmpDir(testWorkspace);
    workspaceContext = null as unknown as WorkspaceContext;
    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(
        '0.0.1',
        'test-server',
        { url: 'http://test-server', type: 'http' },
        false,
        workspaceContext,
        MOCK_CONTEXT,
      ),
    ).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(
        '0.0.1',
        'test-server',
        { url: 'http://test-server', type: 'sse' },
        false,
        workspaceContext,
        MOCK_CONTEXT,
      ),
    ).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(
      '0.0.1',
      'test-server',
      { url: 'http://test-server' },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    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(
        '0.0.1',
        'test-server',
        { url: 'http://test-server' },
        false,
        workspaceContext,
        MOCK_CONTEXT,
      ),
    ).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(
      '0.0.1',
      'test-server',
      { url: 'http://test-server' },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    expect(client).toBe(mockedClient);
    expect(mockedClient.connect).toHaveBeenCalledTimes(2);
  });
});

describe('connectToMcpServer - OAuth with transport fallback', () => {
  let mockedClient: ClientLib.Client;
  let workspaceContext: WorkspaceContext;
  let testWorkspace: string;
  let mockAuthProvider: MCPOAuthProvider;
  let mockTokenStorage: MCPOAuthTokenStorage;

  beforeEach(() => {
    mockedClient = {
      connect: vi.fn(),
      close: vi.fn(),
      registerCapabilities: vi.fn(),
      setRequestHandler: vi.fn(),
      onclose: vi.fn(),
      notification: vi.fn(),
    } as unknown as ClientLib.Client;
    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(() => {});

    // 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' }),
    } as unknown as MCPOAuthTokenStorage;
    vi.mocked(MCPOAuthTokenStorage).mockReturnValue(mockTokenStorage);

    mockAuthProvider = {
      authenticate: vi.fn().mockResolvedValue(undefined),
      getValidToken: vi.fn().mockResolvedValue('test-access-token'),
      tokenStorage: mockTokenStorage,
    } as unknown as MCPOAuthProvider;
    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(async () => {
    vi.useRealTimers();
    await cleanupTmpDir(testWorkspace);
    workspaceContext = null as unknown as WorkspaceContext;
    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
    vi.mocked(mockedClient.connect)
      .mockRejectedValueOnce(new StreamableHTTPError(404, 'Not Found'))
      .mockRejectedValueOnce(new StreamableHTTPError(401, 'Unauthorized'))
      .mockResolvedValueOnce(undefined);

    const client = await connectToMcpServer(
      '0.0.1',
      'test-server',
      { url: 'http://test-server', oauth: { enabled: true } },
      false,
      workspaceContext,
      MOCK_CONTEXT,
    );

    expect(client).toBe(mockedClient);
    expect(mockedClient.connect).toHaveBeenCalledTimes(3);
    expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce();
  });
});
