flynn/src/mcp/manager.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { McpManager } from './manager.js';
import { ToolRegistry } from '../tools/registry.js';

// Mock McpClient to avoid spawning real processes
vi.mock('./client.js', () => ({
  McpClient: vi.fn().mockImplementation((config) => ({
    serverName: config.name,
    status: 'disconnected',
    tools: [],
    error: undefined,
    connect: vi.fn(async function (this: { status: string; tools: { name: string; description: string; inputSchema: { type: string; properties: Record<string, unknown> } }[] }) {
      this.status = 'connected';
      // Simulate tool discovery
      this.tools = [
        {
          name: 'do_thing',
          description: 'Does a thing',
          inputSchema: { type: 'object', properties: { input: { type: 'string' } } },
        },
      ];
    }),
    disconnect: vi.fn(async function (this: { status: string; tools: unknown[] }) {
      this.status = 'disconnected';
      this.tools = [];
    }),
    callTool: vi.fn().mockResolvedValue({ content: 'result', isError: false }),
  })),
}));

describe('McpManager', () => {
  let registry: ToolRegistry;
  let manager: McpManager;

  beforeEach(() => {
    vi.clearAllMocks();
    registry = new ToolRegistry();
    manager = new McpManager(registry);
  });

  it('starts a server and registers its tools', async () => {
    await manager.startServer({
      name: 'test-server',
      command: 'test-cmd',
      args: [],
    });

    // Tool should be registered with mcp: prefix
    const tool = registry.get('mcp:test-server:do_thing');
    expect(tool).toBeDefined();
    expect(tool!.name).toBe('mcp:test-server:do_thing');
    expect(tool!.description).toContain('[MCP:test-server]');
  });

  it('startAll handles multiple servers', async () => {
    await manager.startAll([
      { name: 'server-a', command: 'cmd-a', args: [] },
      { name: 'server-b', command: 'cmd-b', args: [] },
    ]);

    expect(registry.get('mcp:server-a:do_thing')).toBeDefined();
    expect(registry.get('mcp:server-b:do_thing')).toBeDefined();
  });

  it('stopServer unregisters tools and disconnects', async () => {
    await manager.startServer({
      name: 'test-server',
      command: 'test-cmd',
      args: [],
    });

    expect(registry.get('mcp:test-server:do_thing')).toBeDefined();

    await manager.stopServer('test-server');

    expect(registry.get('mcp:test-server:do_thing')).toBeUndefined();
  });

  it('stopAll stops all servers', async () => {
    await manager.startAll([
      { name: 'server-a', command: 'cmd-a', args: [] },
      { name: 'server-b', command: 'cmd-b', args: [] },
    ]);

    await manager.stopAll();

    expect(registry.get('mcp:server-a:do_thing')).toBeUndefined();
    expect(registry.get('mcp:server-b:do_thing')).toBeUndefined();
  });

  it('stopServer is safe when server does not exist', async () => {
    await expect(manager.stopServer('nonexistent')).resolves.toBeUndefined();
  });

  it('startServer replaces existing server with same name', async () => {
    await manager.startServer({
      name: 'test-server',
      command: 'cmd-v1',
      args: [],
    });

    // Start again with same name — should replace
    await manager.startServer({
      name: 'test-server',
      command: 'cmd-v2',
      args: [],
    });

    // Tool should still be registered (re-registered)
    expect(registry.get('mcp:test-server:do_thing')).toBeDefined();
  });

  it('listServers returns state for all servers', async () => {
    await manager.startAll([
      { name: 'server-a', command: 'cmd-a', args: [] },
      { name: 'server-b', command: 'cmd-b', args: [] },
    ]);

    const servers = manager.listServers();
    expect(servers).toHaveLength(2);
    expect(servers[0].config.name).toBe('server-a');
    expect(servers[1].config.name).toBe('server-b');
  });

  it('getServerState returns undefined for unknown server', () => {
    expect(manager.getServerState('nonexistent')).toBeUndefined();
  });

  it('getServerState returns state for known server', async () => {
    await manager.startServer({
      name: 'test-server',
      command: 'test-cmd',
      args: ['--flag'],
    });

    const state = manager.getServerState('test-server');
    expect(state).toBeDefined();
    expect(state!.config.name).toBe('test-server');
    expect(state!.config.args).toEqual(['--flag']);
    expect(state!.tools).toHaveLength(1);
  });

  it('getRegisteredTools returns all MCP tools', async () => {
    await manager.startAll([
      { name: 'server-a', command: 'cmd-a', args: [] },
      { name: 'server-b', command: 'cmd-b', args: [] },
    ]);

    const tools = manager.getRegisteredTools();
    expect(tools).toHaveLength(2);
    expect(tools.map((t) => t.name)).toContain('mcp:server-a:do_thing');
    expect(tools.map((t) => t.name)).toContain('mcp:server-b:do_thing');
  });

  it('restartServer stops and restarts with same config', async () => {
    await manager.startServer({
      name: 'test-server',
      command: 'test-cmd',
      args: ['--arg1'],
    });

    await manager.restartServer('test-server');

    const state = manager.getServerState('test-server');
    expect(state).toBeDefined();
    expect(state!.config.args).toEqual(['--arg1']);
    expect(registry.get('mcp:test-server:do_thing')).toBeDefined();
  });

  it('restartServer throws for unknown server', async () => {
    await expect(manager.restartServer('nonexistent')).rejects.toThrow(
      "MCP server 'nonexistent' not found",
    );
  });
});