mcpctl/src/mcpd/tests/persistent-stdio.test.ts

import { describe, it, expect, vi } from 'vitest';
import { PassThrough } from 'node:stream';
import { PersistentStdioClient } from '../src/services/transport/persistent-stdio.js';
import type { InteractiveExec, McpOrchestrator } from '../src/services/orchestrator.js';

function makeFakeExec(): {
  iexec: InteractiveExec;
  written: string[];
  emit: (line: unknown) => void;
} {
  const stdout = new PassThrough();
  const written: string[] = [];
  const iexec: InteractiveExec = {
    stdout,
    write(data) { written.push(data); },
    close() { stdout.destroy(); },
  };
  const emit = (msg: unknown) => {
    stdout.write(JSON.stringify(msg) + '\n');
  };
  return { iexec, written, emit };
}

function makeOrchestrator(overrides: Partial<McpOrchestrator> = {}): McpOrchestrator {
  return {
    pullImage: vi.fn(),
    createContainer: vi.fn(),
    stopContainer: vi.fn(),
    removeContainer: vi.fn(),
    inspectContainer: vi.fn(),
    getContainerLogs: vi.fn(),
    execInContainer: vi.fn(),
    ping: vi.fn(),
    ...overrides,
  } as McpOrchestrator;
}

describe('PersistentStdioClient', () => {
  it('exec mode calls execInteractive with the command', async () => {
    const fake = makeFakeExec();
    const execInteractive = vi.fn(async () => fake.iexec);
    const orch = makeOrchestrator({ execInteractive });

    const client = new PersistentStdioClient(
      orch,
      'container-1',
      { kind: 'exec', command: ['node', 'index.js'] },
    );

    // Drive the handshake: respond to the first init request (id=1)
    // then to the subsequent tools/list request (id=2).
    const sendPromise = client.send('tools/list');
    await new Promise((r) => setTimeout(r, 10));

    const init = JSON.parse(fake.written[0]!);
    expect(init.method).toBe('initialize');
    fake.emit({ jsonrpc: '2.0', id: init.id, result: { capabilities: {} } });
    await new Promise((r) => setTimeout(r, 150));

    // Second written msg is notifications/initialized; third is tools/list
    const toolsReq = JSON.parse(fake.written[2]!);
    expect(toolsReq.method).toBe('tools/list');
    fake.emit({ jsonrpc: '2.0', id: toolsReq.id, result: { tools: [] } });

    const res = await sendPromise;
    expect(res.result).toEqual({ tools: [] });
    expect(execInteractive).toHaveBeenCalledWith('container-1', ['node', 'index.js']);
    client.close();
  });

  it('attach mode calls attachInteractive and never execInteractive', async () => {
    const fake = makeFakeExec();
    const attachInteractive = vi.fn(async () => fake.iexec);
    const execInteractive = vi.fn();
    const orch = makeOrchestrator({ attachInteractive, execInteractive });

    const client = new PersistentStdioClient(
      orch,
      'container-gitea',
      { kind: 'attach' },
    );

    const sendPromise = client.send('tools/list');
    await new Promise((r) => setTimeout(r, 10));

    const init = JSON.parse(fake.written[0]!);
    fake.emit({ jsonrpc: '2.0', id: init.id, result: { capabilities: {} } });
    await new Promise((r) => setTimeout(r, 150));

    const req = JSON.parse(fake.written[2]!);
    fake.emit({ jsonrpc: '2.0', id: req.id, result: { tools: [{ name: 'list_repos' }] } });

    const res = await sendPromise;
    expect((res.result as { tools: unknown[] }).tools).toHaveLength(1);
    expect(attachInteractive).toHaveBeenCalledWith('container-gitea');
    expect(execInteractive).not.toHaveBeenCalled();
    client.close();
  });

  it('attach mode throws if orchestrator does not support attach', async () => {
    const orch = makeOrchestrator({}); // no attachInteractive
    const client = new PersistentStdioClient(orch, 'c', { kind: 'attach' });
    await expect(client.send('tools/list')).rejects.toThrow(/attach/i);
  });

  it('exec mode throws if orchestrator does not support execInteractive', async () => {
    const orch = makeOrchestrator({}); // no execInteractive
    const client = new PersistentStdioClient(orch, 'c', { kind: 'exec', command: ['x'] });
    await expect(client.send('tools/list')).rejects.toThrow(/interactive exec/i);
  });
});

// Drive a fake session through the init handshake so send() can round-trip.
async function completeHandshake(fake: ReturnType<typeof makeFakeExec>): Promise<void> {
  await new Promise((r) => setTimeout(r, 10));
  const init = JSON.parse(fake.written[0]!);
  fake.emit({ jsonrpc: '2.0', id: init.id, result: { capabilities: {} } });
  await new Promise((r) => setTimeout(r, 150));
}

describe('PersistentStdioClient teardown & reconnect (mcpctl#114)', () => {
  it("stdout 'end' rejects in-flight requests and the next send() redials", async () => {
    const fakes = [makeFakeExec(), makeFakeExec()];
    let call = 0;
    const execInteractive = vi.fn(async () => fakes[call++]!.iexec);
    const client = new PersistentStdioClient(
      makeOrchestrator({ execInteractive }),
      'c1',
      { kind: 'exec', command: ['node'] },
    );

    const first = client.send('tools/list');
    await completeHandshake(fakes[0]!);
    // Kill the session while tools/list is in flight.
    fakes[0]!.iexec.stdout.end();
    await expect(first).rejects.toThrow('STDIO process exited');
    expect(client.isConnected).toBe(false);

    // Next send must redial through a fresh interactive session.
    const second = client.send('tools/list');
    await completeHandshake(fakes[1]!);
    const req = JSON.parse(fakes[1]!.written[2]!);
    fakes[1]!.emit({ jsonrpc: '2.0', id: req.id, result: { tools: [] } });
    const res = await second;
    expect(res.result).toEqual({ tools: [] });
    expect(execInteractive).toHaveBeenCalledTimes(2);
    client.close();
  });

  it("stdout 'error' tears down with the error message", async () => {
    const fake = makeFakeExec();
    const client = new PersistentStdioClient(
      makeOrchestrator({ execInteractive: vi.fn(async () => fake.iexec) }),
      'c1',
      { kind: 'exec', command: ['node'] },
    );
    const pending = client.send('tools/list');
    await completeHandshake(fake);
    fake.iexec.stdout.emit('error', new Error('socket hang up'));
    await expect(pending).rejects.toThrow('STDIO stream error: socket hang up');
    expect(client.isConnected).toBe(false);
    client.close();
  });

  it("a LATE 'end' from the previous session does not clobber the reconnected one", async () => {
    const fakes = [makeFakeExec(), makeFakeExec()];
    let call = 0;
    const client = new PersistentStdioClient(
      makeOrchestrator({ attachInteractive: vi.fn(async () => fakes[call++]!.iexec) }),
      'c1',
      { kind: 'attach' },
    );

    const first = client.send('tools/list');
    await completeHandshake(fakes[0]!);
    fakes[0]!.iexec.stdout.end();
    await expect(first).rejects.toThrow('STDIO process exited');

    // Reconnect on a fresh session.
    const second = client.send('tools/list');
    await completeHandshake(fakes[1]!);

    // The old session's stream fires a stale 'close' AFTER the reconnect —
    // the identity guard must ignore it.
    fakes[0]!.iexec.stdout.emit('close');
    expect(client.isConnected).toBe(true);

    const req = JSON.parse(fakes[1]!.written[2]!);
    fakes[1]!.emit({ jsonrpc: '2.0', id: req.id, result: { tools: [{ name: 't' }] } });
    const res = await second;
    expect((res.result as { tools: unknown[] }).tools).toHaveLength(1);
    client.close();
  });

  it('attach mode also redials after teardown', async () => {
    const fakes = [makeFakeExec(), makeFakeExec()];
    let call = 0;
    const attachInteractive = vi.fn(async () => fakes[call++]!.iexec);
    const client = new PersistentStdioClient(
      makeOrchestrator({ attachInteractive }),
      'c-gitea',
      { kind: 'attach' },
    );

    const first = client.send('tools/list');
    await completeHandshake(fakes[0]!);
    fakes[0]!.iexec.stdout.emit('close');
    await expect(first).rejects.toThrow('STDIO stream closed');

    const second = client.send('tools/list');
    await completeHandshake(fakes[1]!);
    const req = JSON.parse(fakes[1]!.written[2]!);
    fakes[1]!.emit({ jsonrpc: '2.0', id: req.id, result: { tools: [] } });
    await expect(second).resolves.toMatchObject({ result: { tools: [] } });
    expect(attachInteractive).toHaveBeenCalledTimes(2);
    client.close();
  });
});