promptfoo--promptfoo/test/providers/mcp/client.test.ts

import path from 'path';

import { beforeEach, describe, expect, it, vi } from 'vitest';

const mockGetEnvInt = vi.hoisted(() => vi.fn().mockReturnValue(undefined));
vi.mock('../../../src/envars', async () => ({
  ...(await vi.importActual<typeof import('../../../src/envars')>('../../../src/envars')),
  getEnvInt: (...args: unknown[]) => mockGetEnvInt(...args),
}));

const mockGetOAuthTokenWithExpiry = vi.hoisted(() =>
  vi.fn().mockResolvedValue({
    accessToken: 'mock-oauth-token',
    expiresAt: Date.now() + 3600000, // 1 hour from now
  }),
);
vi.mock('../../../src/providers/mcp/util', async () => ({
  ...(await vi.importActual<typeof import('../../../src/providers/mcp/util')>(
    '../../../src/providers/mcp/util',
  )),
  getOAuthTokenWithExpiry: (...args: unknown[]) => mockGetOAuthTokenWithExpiry(...args),
}));

const mcpMocks = vi.hoisted(() => {
  const mockClient = {
    _clientInfo: {},
    _capabilities: {},
    registerCapabilities: vi.fn(),
    assertCapability: vi.fn(),
    connect: vi.fn(),
    ping: vi.fn().mockResolvedValue({}),
    listTools: vi.fn().mockResolvedValue({
      tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
    }),
    callTool: vi.fn(),
    close: vi.fn().mockResolvedValue(undefined),
  };

  const mockStdioTransport = {
    close: vi.fn().mockResolvedValue(undefined),
    connect: vi.fn(),
    start: vi.fn(),
    send: vi.fn(),
  };

  const mockStreamableHTTPTransport = {
    close: vi.fn().mockResolvedValue(undefined),
    connect: vi.fn(),
    start: vi.fn(),
    send: vi.fn(),
  };

  const mockSSETransport = {
    close: vi.fn().mockResolvedValue(undefined),
    connect: vi.fn(),
    start: vi.fn(),
    send: vi.fn(),
  };

  const MockClient = vi.fn(function MockClient() {
    return mockClient;
  });

  const MockStdioTransport = vi.fn(function MockStdioTransport() {
    return mockStdioTransport;
  });

  const MockStreamableHTTPTransport = vi.fn(function MockStreamableHTTPTransport() {
    return mockStreamableHTTPTransport;
  });

  const MockSSETransport = vi.fn(function MockSSETransport() {
    return mockSSETransport;
  });

  return {
    mockClient,
    mockStdioTransport,
    mockStreamableHTTPTransport,
    mockSSETransport,
    MockClient,
    MockSSETransport,
    MockStdioTransport,
    MockStreamableHTTPTransport,
  };
});

const { mockClient, mockStdioTransport, mockStreamableHTTPTransport } = mcpMocks;

// Mock the modules before importing them
vi.mock('@modelcontextprotocol/sdk/client/index.js', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    Client: mcpMocks.MockClient,
  };
});

vi.mock('@modelcontextprotocol/sdk/client/stdio.js', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    StdioClientTransport: mcpMocks.MockStdioTransport,
  };
});

vi.mock('@modelcontextprotocol/sdk/client/streamableHttp.js', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    StreamableHTTPClientTransport: mcpMocks.MockStreamableHTTPTransport,
  };
});

vi.mock('@modelcontextprotocol/sdk/client/sse.js', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    SSEClientTransport: mcpMocks.MockSSETransport,
  };
});

// Import the mocked modules after mocking
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
import cliState from '../../../src/cliState';
import { MCPClient } from '../../../src/providers/mcp/client';
import { createDeferred } from '../../util/utils';

function createMockClient(callTool = vi.fn()) {
  return {
    _clientInfo: {},
    _capabilities: {},
    registerCapabilities: vi.fn(),
    assertCapability: vi.fn(),
    connect: vi.fn().mockResolvedValue(undefined),
    ping: vi.fn().mockResolvedValue({}),
    listTools: vi.fn().mockResolvedValue({
      tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
    }),
    callTool,
    close: vi.fn().mockResolvedValue(undefined),
  };
}

describe('MCPClient', () => {
  let mcpClient: MCPClient;

  beforeEach(() => {
    vi.clearAllMocks();
    mockGetEnvInt.mockReset();
    mockGetEnvInt.mockReturnValue(undefined);
    // Reset the OAuth token mock to return a valid token by default
    mockGetOAuthTokenWithExpiry.mockReset();
    mockGetOAuthTokenWithExpiry.mockResolvedValue({
      accessToken: 'mock-oauth-token',
      expiresAt: Date.now() + 3600000, // 1 hour from now
    });
    cliState.basePath = undefined;
  });

  describe('initialize', () => {
    it('should not initialize if disabled', async () => {
      mcpClient = new MCPClient({ enabled: false });
      await mcpClient.initialize();
      expect(mcpClient.hasInitialized).toBe(false);
    });

    it('should initialize with single server config', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();

      expect(StdioClientTransport).toHaveBeenCalledWith({
        command: 'npm',
        args: ['start'],
        env: process.env as Record<string, string>,
      });
      expect(mockClient.connect).toHaveBeenCalledWith(mockStdioTransport, undefined);
      await mcpClient.cleanup();
      expect(mcpClient.hasInitialized).toBe(false);
    });

    it('should initialize with multiple servers', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        servers: [
          {
            name: 'server1',
            command: 'npm',
            args: ['start'],
          },
          {
            name: 'server2',
            path: 'script.js',
          },
        ],
      });

      await mcpClient.initialize();

      expect(StdioClientTransport).toHaveBeenCalledTimes(2);
      expect(mockClient.connect).toHaveBeenCalledTimes(2);
    });

    it('should resolve local server paths relative to the config base path', async () => {
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      cliState.basePath = '/tmp/simple-mcp';

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          path: './example-server.js',
        },
      });

      await mcpClient.initialize();

      expect(StdioClientTransport).toHaveBeenCalledWith({
        command: process.execPath,
        args: [path.resolve(cliState.basePath, './example-server.js')],
        env: process.env as Record<string, string>,
      });
    });

    it('should throw error for unsupported file type', async () => {
      mcpClient = new MCPClient({
        enabled: true,
        server: {
          path: 'script.txt',
        },
      });

      await expect(mcpClient.initialize()).rejects.toThrow(
        'Local server must be a .js or .py file',
      );
    });

    it('should initialize with remote server using StreamableHTTPClientTransport', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
        },
      });

      await mcpClient.initialize();

      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(expect.any(URL), undefined);
      expect(mockClient.connect).toHaveBeenCalledWith(mockStreamableHTTPTransport, undefined);
    });

    it('should initialize with remote server using StreamableHTTPClientTransport with headers', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      const customHeaders = {
        'X-Custom-Header': 'custom-value',
        Authorization: 'Bearer test-token',
      };

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          headers: customHeaders,
        },
      });

      await mcpClient.initialize();

      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining(customHeaders),
          }),
        }),
      );
      expect(mockClient.connect).toHaveBeenCalledWith(mockStreamableHTTPTransport, undefined);
    });

    it('should fall back to SSEClientTransport if StreamableHTTPClientTransport fails', async () => {
      // Reset mocks for this test
      mockClient.connect
        .mockImplementationOnce(function () {
          throw new Error('Connection failed');
        })
        .mockResolvedValueOnce(undefined);

      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
        },
      });

      await mcpClient.initialize();

      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(expect.any(URL), undefined);
      expect(SSEClientTransport).toHaveBeenCalledWith(expect.any(URL), undefined);
      expect(mockClient.connect).toHaveBeenCalledTimes(2);
    });

    it('should fall back to SSEClientTransport with headers if StreamableHTTPClientTransport fails', async () => {
      // Reset mocks for this test
      mockClient.connect
        .mockImplementationOnce(function () {
          throw new Error('Connection failed');
        })
        .mockResolvedValueOnce(undefined);

      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      const customHeaders = {
        'X-Custom-Header': 'custom-value',
        Authorization: 'Bearer test-token',
      };

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          headers: customHeaders,
        },
      });

      await mcpClient.initialize();

      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining(customHeaders),
          }),
        }),
      );
      expect(SSEClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining(customHeaders),
          }),
        }),
      );
      expect(mockClient.connect).toHaveBeenCalledTimes(2);
    });

    it('should filter tools according to config.tools', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [
          { name: 'tool1', description: 'desc1', inputSchema: {} },
          { name: 'tool2', description: 'desc2', inputSchema: {} },
        ],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
        tools: ['tool2'],
      });

      await mcpClient.initialize();

      // getAllTools should return only tool2
      expect(mcpClient.getAllTools()).toEqual([
        { name: 'tool2', description: 'desc2', inputSchema: {} },
      ]);
    });

    it('should exclude tools according to config.exclude_tools', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [
          { name: 'tool1', description: 'desc1', inputSchema: {} },
          { name: 'tool2', description: 'desc2', inputSchema: {} },
        ],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
        exclude_tools: ['tool1'],
      });

      await mcpClient.initialize();

      expect(mcpClient.getAllTools()).toEqual([
        { name: 'tool2', description: 'desc2', inputSchema: {} },
      ]);
    });

    it('should initialize with correct client metadata including name, version, and description', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();

      expect(Client).toHaveBeenCalledWith({
        name: 'promptfoo-MCP',
        version: '1.0.0',
        description: 'Promptfoo MCP client for connecting to MCP servers during LLM evaluations',
      });
    });

    it('should provide a descriptive client description for MCP server identification', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();

      // Verify the description is provided and meaningful
      const clientCall = vi.mocked(Client).mock.calls[0][0];
      expect(clientCall).toHaveProperty('description');
      const description = clientCall.description as string;
      expect(typeof description).toBe('string');
      expect(description.length).toBeGreaterThan(0);
      // Case-insensitive check for key terms
      expect(description.toLowerCase()).toContain('promptfoo');
      expect(description).toContain('MCP');
    });

    it('should pass timeout to listTools when configured', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        timeout: 900000, // 15 minutes
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();

      expect(mockClient.listTools).toHaveBeenCalledWith(undefined, { timeout: 900000 });
    });

    it('should pass timeout options to connect()', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        timeout: 300000,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();

      expect(mockClient.connect).toHaveBeenCalledWith(expect.anything(), { timeout: 300000 });
    });

    it('should ping server when pingOnConnect is true', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.ping.mockResolvedValueOnce({});
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        pingOnConnect: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();

      expect(mockClient.ping).toHaveBeenCalled();
    });

    it('should fail initialization if ping fails', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.ping.mockRejectedValueOnce(new Error('Server not responding'));

      mcpClient = new MCPClient({
        enabled: true,
        pingOnConnect: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await expect(mcpClient.initialize()).rejects.toThrow('ping failed');
    });

    it('should pass resetTimeoutOnProgress option', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'result' });

      mcpClient = new MCPClient({
        enabled: true,
        timeout: 300000,
        resetTimeoutOnProgress: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      await mcpClient.callTool('tool1', {});

      expect(mockClient.callTool).toHaveBeenCalledWith(
        { name: 'tool1', arguments: {} },
        undefined,
        { timeout: 300000, resetTimeoutOnProgress: true },
      );
    });

    it('should pass maxTotalTimeout option', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'result' });

      mcpClient = new MCPClient({
        enabled: true,
        timeout: 300000,
        resetTimeoutOnProgress: true,
        maxTotalTimeout: 900000,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      await mcpClient.callTool('tool1', {});

      expect(mockClient.callTool).toHaveBeenCalledWith(
        { name: 'tool1', arguments: {} },
        undefined,
        { timeout: 300000, resetTimeoutOnProgress: true, maxTotalTimeout: 900000 },
      );
    });
  });

  describe('callTool', () => {
    it('should call tool successfully', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'result' });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', { arg: 'value' });

      expect(result).toEqual({ content: 'result', raw: { content: 'result' } });
      expect(mockClient.callTool).toHaveBeenCalledWith(
        {
          name: 'tool1',
          arguments: { arg: 'value' },
        },
        undefined,
        undefined,
      );
    });

    it('should pass timeout option when configured', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'result' });

      mcpClient = new MCPClient({
        enabled: true,
        timeout: 900000, // 15 minutes
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', { arg: 'value' });

      expect(result).toEqual({ content: 'result', raw: { content: 'result' } });
      expect(mockClient.callTool).toHaveBeenCalledWith(
        {
          name: 'tool1',
          arguments: { arg: 'value' },
        },
        undefined,
        { timeout: 900000 },
      );
    });

    it('should use MCP_REQUEST_TIMEOUT_MS env var as fallback when no config timeout', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'result' });

      // Mock env var to return a timeout value
      mockGetEnvInt.mockReturnValue(600000); // 10 minutes

      mcpClient = new MCPClient({
        enabled: true,
        // No timeout in config
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', { arg: 'value' });

      expect(result).toEqual({ content: 'result', raw: { content: 'result' } });
      expect(mockClient.callTool).toHaveBeenCalledWith(
        {
          name: 'tool1',
          arguments: { arg: 'value' },
        },
        undefined,
        { timeout: 600000 },
      );

      // Reset mock
      mockGetEnvInt.mockReturnValue(undefined);
    });

    it('should prefer config timeout over env var', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'result' });

      // Mock env var to return a different timeout value
      mockGetEnvInt.mockReturnValue(600000); // 10 minutes

      mcpClient = new MCPClient({
        enabled: true,
        timeout: 900000, // 15 minutes - should take precedence
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      await mcpClient.callTool('tool1', { arg: 'value' });

      expect(mockClient.callTool).toHaveBeenCalledWith(
        {
          name: 'tool1',
          arguments: { arg: 'value' },
        },
        undefined,
        { timeout: 900000 }, // Config timeout takes precedence
      );

      // Reset mock
      mockGetEnvInt.mockReturnValue(undefined);
    });

    it('should handle tool error', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockRejectedValueOnce(new Error('Tool error'));

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
        debug: true,
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', {});

      expect(result).toEqual({
        content: '',
        error: 'Tool error',
      });
    });

    it('should surface MCP tool error results', async () => {
      // Reset mocks for this test
      const errorContent = [{ type: 'text', text: 'Invalid arguments' }];
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({
        content: errorContent,
        isError: true,
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', {});

      expect(result).toEqual({
        content: JSON.stringify(errorContent),
        isError: true,
        raw: {
          content: errorContent,
          isError: true,
        },
      });
    });

    it('should throw error for unknown tool', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      await expect(mcpClient.callTool('unknown', {})).rejects.toThrow('Tool unknown not found');
    });

    it('should support tool content as Buffer', async () => {
      // Reset mocks for this test
      const contentBuffer = Buffer.from('buffered-result');
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: contentBuffer });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', {});

      expect(result).toEqual({ content: 'buffered-result', raw: { content: contentBuffer } });
    });

    it('should return empty string if result content is falsy', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({});

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', {});

      expect(result).toEqual({ content: '', raw: {} });
    });

    it('should parse JSON-stringified content correctly', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: '"Hello World"' });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', {});

      expect(result).toEqual({ content: 'Hello World', raw: { content: '"Hello World"' } });
    });

    it('should handle non-JSON string content correctly', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValueOnce({ content: 'Plain text response' });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      const result = await mcpClient.callTool('tool1', {});

      expect(result).toEqual({
        content: 'Plain text response',
        raw: { content: 'Plain text response' },
      });
    });
  });

  describe('cleanup', () => {
    it('should cleanup all clients', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      await mcpClient.cleanup();

      expect(mockClient.close).toHaveBeenCalledWith();
    });

    it('should handle cleanup errors', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.close.mockRejectedValueOnce(new Error('Cleanup error'));

      mcpClient = new MCPClient({
        enabled: true,
        server: { command: 'npm', args: ['start'] },
        debug: true,
      });

      await mcpClient.initialize();
      await mcpClient.cleanup();

      expect(mockClient.close).toHaveBeenCalledWith();
    });
  });

  describe('getAllTools', () => {
    it('should return all tools from all servers', async () => {
      // Reset mocks for this test
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [
          { name: 'tool1', description: 'desc1', inputSchema: {} },
          { name: 'tool2', description: 'desc2', inputSchema: {} },
        ],
      });

      mcpClient = new MCPClient({
        enabled: true,
        servers: [
          {
            name: 'server1',
            command: 'npm',
            args: ['start'],
          },
          {
            name: 'server2',
            path: 'script.js',
          },
        ],
      });
      await mcpClient.initialize();
      const allTools = mcpClient.getAllTools();
      expect(Array.isArray(allTools)).toBe(true);
      expect(allTools.length).toBeGreaterThan(0);
    });

    it('should return empty array if no tools', () => {
      mcpClient = new MCPClient({ enabled: true });
      // force tools to be empty
      expect(mcpClient.getAllTools()).toEqual([]);
    });

    it('should skip stale tool entries when another server still has the requested tool', async () => {
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValue({ content: 'result' });

      mcpClient = new MCPClient({
        enabled: true,
        servers: [
          {
            name: 'server1',
            command: 'npm',
            args: ['start'],
          },
          {
            name: 'server2',
            command: 'node',
            args: ['server.js'],
          },
        ],
      });

      await mcpClient.initialize();
      (
        mcpClient as unknown as {
          clients: Map<string, unknown>;
        }
      ).clients.delete('server1');

      await expect(mcpClient.callTool('tool1', {})).resolves.toEqual({
        content: 'result',
        raw: { content: 'result' },
      });
      expect(mockClient.callTool).toHaveBeenCalledTimes(1);
    });

    it('should skip a server whose proactive OAuth refresh fails when another server still has the requested tool', async () => {
      const server1Client = createMockClient(vi.fn().mockResolvedValue({ content: 'server1' }));
      const server2Client = createMockClient(
        vi.fn().mockResolvedValue({ content: 'server2 result' }),
      );
      vi.mocked(Client)
        .mockImplementationOnce(function MockServer1Client() {
          return server1Client as unknown as Client;
        })
        .mockImplementationOnce(function MockServer2Client() {
          return server2Client as unknown as Client;
        });

      mockGetOAuthTokenWithExpiry
        .mockResolvedValueOnce({
          accessToken: 'stale-token',
          expiresAt: Date.now() - 1000,
        })
        .mockRejectedValueOnce(new Error('refresh failed'));

      mcpClient = new MCPClient({
        enabled: true,
        servers: [
          {
            name: 'server1',
            url: 'http://localhost:3000',
            auth: {
              type: 'oauth',
              grantType: 'client_credentials',
              clientId: 'test-client',
              clientSecret: 'test-secret',
              tokenUrl: 'https://auth.example.com/token',
            },
          },
          {
            name: 'server2',
            command: 'npm',
            args: ['start'],
          },
        ],
      });

      await mcpClient.initialize();

      await expect(mcpClient.callTool('tool1', {})).resolves.toEqual({
        content: 'server2 result',
        raw: { content: 'server2 result' },
      });
      expect(server1Client.callTool).not.toHaveBeenCalled();
      expect(server2Client.callTool).toHaveBeenCalledTimes(1);
    });

    it('should report when a known tool only exists on disconnected servers', async () => {
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          name: 'server1',
          command: 'npm',
          args: ['start'],
        },
      });

      await mcpClient.initialize();
      (
        mcpClient as unknown as {
          clients: Map<string, unknown>;
        }
      ).clients.delete('server1');

      await expect(mcpClient.callTool('tool1', {})).rejects.toThrow(
        'Tool tool1 is known but MCP server is disconnected: server1',
      );
    });
  });

  describe('OAuth authentication', () => {
    it('should use static headers for OAuth with tokenUrl configured', async () => {
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      // When tokenUrl is configured, we use static Authorization header
      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              Authorization: 'Bearer mock-oauth-token',
            }),
          }),
        }),
      );
      // Verify authProvider is NOT used when tokenUrl is configured
      const callArgs = vi.mocked(StreamableHTTPClientTransport).mock.calls[0];
      expect(callArgs[1]).not.toHaveProperty('authProvider');
    });

    it('should use static headers for OAuth password grant with tokenUrl', async () => {
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'password',
            tokenUrl: 'https://auth.example.com/token',
            username: 'testuser',
            password: 'testpass',
          },
        },
      });

      await mcpClient.initialize();

      // When tokenUrl is configured, we use static Authorization header
      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              Authorization: 'Bearer mock-oauth-token',
            }),
          }),
        }),
      );
    });

    it('should NOT use authProvider for bearer auth type', async () => {
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'bearer',
            token: 'static-token',
          },
        },
      });

      await mcpClient.initialize();

      // Should have headers but NOT authProvider
      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              Authorization: 'Bearer static-token',
            }),
          }),
        }),
      );
      // Verify authProvider is not in the options
      const callArgs = vi.mocked(StreamableHTTPClientTransport).mock.calls[0];
      expect(callArgs[1]).not.toHaveProperty('authProvider');
    });

    it('should NOT use authProvider for basic auth type', async () => {
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'basic',
            username: 'user',
            password: 'pass',
          },
        },
      });

      await mcpClient.initialize();

      // Should have headers but NOT authProvider
      const expectedAuth = 'Basic ' + Buffer.from('user:pass').toString('base64');
      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              Authorization: expectedAuth,
            }),
          }),
        }),
      );
      // Verify authProvider is not in the options
      const callArgs = vi.mocked(StreamableHTTPClientTransport).mock.calls[0];
      expect(callArgs[1]).not.toHaveProperty('authProvider');
    });

    it('should combine OAuth static headers with custom headers when tokenUrl configured', async () => {
      mockClient.connect.mockResolvedValueOnce(undefined);
      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      const customHeaders = {
        'X-Custom-Header': 'custom-value',
      };

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          headers: customHeaders,
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      // Should have both Authorization header and custom headers
      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              ...customHeaders,
              Authorization: 'Bearer mock-oauth-token',
            }),
          }),
        }),
      );
      // Verify authProvider is NOT used when tokenUrl is configured
      const callArgs = vi.mocked(StreamableHTTPClientTransport).mock.calls[0];
      expect(callArgs[1]).not.toHaveProperty('authProvider');
    });

    it('should fall back to SSEClientTransport with static headers if StreamableHTTPClientTransport fails', async () => {
      mockClient.connect
        .mockImplementationOnce(function () {
          throw new Error('Connection failed');
        })
        .mockResolvedValueOnce(undefined);

      mockClient.listTools.mockResolvedValueOnce({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      // Both transports should receive static Authorization headers when tokenUrl is configured
      expect(StreamableHTTPClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              Authorization: 'Bearer mock-oauth-token',
            }),
          }),
        }),
      );
      expect(SSEClientTransport).toHaveBeenCalledWith(
        expect.any(URL),
        expect.objectContaining({
          requestInit: expect.objectContaining({
            headers: expect.objectContaining({
              Authorization: 'Bearer mock-oauth-token',
            }),
          }),
        }),
      );
    });

    it('should proactively refresh token before callTool if close to expiration', async () => {
      // First call with valid token
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValue({ content: 'result' });

      // Set token to expire soon (within buffer)
      mockGetOAuthTokenWithExpiry.mockResolvedValueOnce({
        accessToken: 'initial-token',
        expiresAt: Date.now() + 30000, // 30 seconds, within 60s buffer
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      // Now set up a new token for the refresh
      mockGetOAuthTokenWithExpiry.mockResolvedValueOnce({
        accessToken: 'refreshed-token',
        expiresAt: Date.now() + 3600000, // 1 hour
      });

      // Call tool - should trigger proactive refresh
      await mcpClient.callTool('tool1', {});

      // Should have called getOAuthTokenWithExpiry twice (initial + refresh)
      expect(mockGetOAuthTokenWithExpiry).toHaveBeenCalledTimes(2);
    });

    it('should call the reconnected client after a proactive token refresh', async () => {
      const oldClient = createMockClient(vi.fn());
      const refreshedClient = createMockClient(vi.fn().mockResolvedValue({ content: 'result' }));
      vi.mocked(Client)
        .mockImplementationOnce(function () {
          return oldClient as unknown as Client;
        })
        .mockImplementationOnce(function () {
          return refreshedClient as unknown as Client;
        });

      mockGetOAuthTokenWithExpiry
        .mockResolvedValueOnce({
          accessToken: 'initial-token',
          expiresAt: Date.now() + 30_000,
        })
        .mockResolvedValueOnce({
          accessToken: 'refreshed-token',
          expiresAt: Date.now() + 3_600_000,
        });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      await expect(mcpClient.callTool('tool1', {})).resolves.toEqual({
        content: 'result',
        raw: { content: 'result' },
      });
      expect(oldClient.callTool).not.toHaveBeenCalled();
      expect(refreshedClient.callTool).toHaveBeenCalledTimes(1);
    });

    it('should wait for an in-progress reactive reconnect when token metadata is already refreshed', async () => {
      const reconnectStarted = createDeferred<void>();
      const reconnectContinue = createDeferred<void>();
      const oldClient = createMockClient(
        vi.fn().mockRejectedValueOnce(new Error('401 Unauthorized')),
      );
      const refreshedClient = createMockClient(
        vi.fn().mockResolvedValue({ content: 'refreshed-result' }),
      );
      refreshedClient.connect.mockImplementationOnce(async () => {
        reconnectStarted.resolve(undefined);
        await reconnectContinue.promise;
      });

      vi.mocked(Client)
        .mockImplementationOnce(function () {
          return oldClient as unknown as Client;
        })
        .mockImplementationOnce(function () {
          return refreshedClient as unknown as Client;
        });

      mockGetOAuthTokenWithExpiry
        .mockResolvedValueOnce({
          accessToken: 'initial-token',
          expiresAt: Date.now() + 3_600_000,
        })
        .mockResolvedValueOnce({
          accessToken: 'refreshed-token',
          expiresAt: Date.now() + 3_600_000,
        });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      const call1 = mcpClient.callTool('tool1', {});
      await reconnectStarted.promise;
      const call2 = mcpClient.callTool('tool1', {});

      reconnectContinue.resolve(undefined);

      await expect(Promise.all([call1, call2])).resolves.toEqual([
        { content: 'refreshed-result', raw: { content: 'refreshed-result' } },
        { content: 'refreshed-result', raw: { content: 'refreshed-result' } },
      ]);
      expect(oldClient.callTool).toHaveBeenCalledTimes(1);
      expect(refreshedClient.callTool).toHaveBeenCalledTimes(2);
    });

    it('should deduplicate concurrent proactive token refreshes', async () => {
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValue({ content: 'result' });

      mockGetOAuthTokenWithExpiry.mockResolvedValueOnce({
        accessToken: 'initial-token',
        expiresAt: Date.now() + 30_000,
      });

      const refreshStarted = createDeferred<void>();
      const refreshToken = createDeferred<{ accessToken: string; expiresAt: number }>();
      mockGetOAuthTokenWithExpiry.mockImplementationOnce(async () => {
        refreshStarted.resolve(undefined);
        return refreshToken.promise;
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      const call1 = mcpClient.callTool('tool1', {});
      await refreshStarted.promise;
      const call2 = mcpClient.callTool('tool1', {});

      refreshToken.resolve({
        accessToken: 'refreshed-token',
        expiresAt: Date.now() + 3_600_000,
      });

      await Promise.all([call1, call2]);

      expect(mockGetOAuthTokenWithExpiry).toHaveBeenCalledTimes(2);
      expect(mockClient.callTool).toHaveBeenCalledTimes(2);
    });

    it('should deduplicate retry refreshes after a failed in-progress proactive refresh', async () => {
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValue({ content: 'result' });

      mockGetOAuthTokenWithExpiry.mockResolvedValueOnce({
        accessToken: 'initial-token',
        expiresAt: Date.now() + 30_000,
      });

      const failedRefreshStarted = createDeferred<void>();
      const failedRefreshToken = createDeferred<{ accessToken: string; expiresAt: number }>();
      mockGetOAuthTokenWithExpiry
        .mockImplementationOnce(async () => {
          failedRefreshStarted.resolve(undefined);
          return failedRefreshToken.promise;
        })
        .mockResolvedValueOnce({
          accessToken: 'refreshed-token',
          expiresAt: Date.now() + 3_600_000,
        });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      const call1 = mcpClient.callTool('tool1', {}).catch(() => {
        // Expected to fail on the first refresh attempt.
      });
      await failedRefreshStarted.promise;
      const call2 = mcpClient.callTool('tool1', {});
      const call3 = mcpClient.callTool('tool1', {});

      failedRefreshToken.reject(new Error('token refresh failed'));

      await Promise.allSettled([call1, call2, call3]);

      expect(mockGetOAuthTokenWithExpiry).toHaveBeenCalledTimes(3);
      expect(mockClient.callTool).toHaveBeenCalledTimes(2);
    });

    it('should deduplicate concurrent reactive token refreshes after auth errors', async () => {
      const oldClient = createMockClient(
        vi
          .fn()
          .mockRejectedValueOnce(new Error('401 Unauthorized'))
          .mockRejectedValueOnce(new Error('401 Unauthorized')),
      );
      const refreshedClient = createMockClient(vi.fn().mockResolvedValue({ content: 'result' }));
      vi.mocked(Client)
        .mockImplementationOnce(function () {
          return oldClient as unknown as Client;
        })
        .mockImplementationOnce(function () {
          return refreshedClient as unknown as Client;
        });

      mockGetOAuthTokenWithExpiry.mockResolvedValueOnce({
        accessToken: 'initial-token',
        expiresAt: Date.now() + 3_600_000,
      });

      const refreshStarted = createDeferred<void>();
      const refreshToken = createDeferred<{ accessToken: string; expiresAt: number }>();
      mockGetOAuthTokenWithExpiry.mockImplementationOnce(async () => {
        refreshStarted.resolve(undefined);
        return refreshToken.promise;
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      const call1 = mcpClient.callTool('tool1', {});
      const call2 = mcpClient.callTool('tool1', {});

      await refreshStarted.promise;
      refreshToken.resolve({
        accessToken: 'reactive-refresh-token',
        expiresAt: Date.now() + 3_600_000,
      });

      await Promise.all([call1, call2]);

      expect(mockGetOAuthTokenWithExpiry).toHaveBeenCalledTimes(2);
      expect(oldClient.callTool).toHaveBeenCalledTimes(2);
      expect(refreshedClient.callTool).toHaveBeenCalledTimes(2);
    });

    it('should not refresh token if still valid', async () => {
      mockClient.connect.mockResolvedValue(undefined);
      mockClient.listTools.mockResolvedValue({
        tools: [{ name: 'tool1', description: 'desc1', inputSchema: {} }],
      });
      mockClient.callTool.mockResolvedValue({ content: 'result' });

      // Token valid for 1 hour (outside buffer)
      mockGetOAuthTokenWithExpiry.mockResolvedValueOnce({
        accessToken: 'valid-token',
        expiresAt: Date.now() + 3600000, // 1 hour
      });

      mcpClient = new MCPClient({
        enabled: true,
        server: {
          url: 'http://localhost:3000',
          auth: {
            type: 'oauth',
            grantType: 'client_credentials',
            clientId: 'test-client',
            clientSecret: 'test-secret',
            tokenUrl: 'https://auth.example.com/token',
          },
        },
      });

      await mcpClient.initialize();

      // Call tool - should NOT trigger refresh
      await mcpClient.callTool('tool1', {});

      // Should have called getOAuthTokenWithExpiry only once (initial)
      expect(mockGetOAuthTokenWithExpiry).toHaveBeenCalledTimes(1);
    });
  });
});