mcpctl/src/mcplocal/tests/acp-client.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { EventEmitter, Readable } from 'node:stream';
import { AcpClient } from '../src/providers/acp-client.js';
import type { AcpClientConfig } from '../src/providers/acp-client.js';

/**
 * Creates a mock child process that speaks ACP protocol.
 * Returns the mock process and helpers to send responses.
 */
function createMockProcess() {
  const stdin = {
    write: vi.fn(),
    writable: true,
  };

  const stdoutEmitter = new EventEmitter();
  const stdout = Object.assign(stdoutEmitter, {
    readable: true,
    // readline needs these
    [Symbol.asyncIterator]: undefined,
    pause: vi.fn(),
    resume: vi.fn(),
    isPaused: () => false,
    setEncoding: vi.fn(),
    read: vi.fn(),
    destroy: vi.fn(),
    pipe: vi.fn(),
    unpipe: vi.fn(),
    unshift: vi.fn(),
    wrap: vi.fn(),
  }) as unknown as Readable;

  const proc = Object.assign(new EventEmitter(), {
    stdin,
    stdout,
    stderr: new EventEmitter(),
    pid: 12345,
    killed: false,
    kill: vi.fn(function (this: { killed: boolean }) {
      this.killed = true;
    }),
  });

  /** Send a line of JSON from the "agent" to our client */
  function sendLine(data: unknown) {
    stdoutEmitter.emit('data', Buffer.from(JSON.stringify(data) + '\n'));
  }

  /** Send a JSON-RPC response */
  function sendResponse(id: number, result: unknown) {
    sendLine({ jsonrpc: '2.0', id, result });
  }

  /** Send a JSON-RPC error */
  function sendError(id: number, code: number, message: string) {
    sendLine({ jsonrpc: '2.0', id, error: { code, message } });
  }

  /** Send a session/update notification with agent_message_chunk */
  function sendChunk(sessionId: string, text: string) {
    sendLine({
      jsonrpc: '2.0',
      method: 'session/update',
      params: {
        sessionId,
        update: {
          sessionUpdate: 'agent_message_chunk',
          content: [{ type: 'text', text }],
        },
      },
    });
  }

  /** Send a session/request_permission request */
  function sendPermissionRequest(id: number, sessionId: string) {
    sendLine({
      jsonrpc: '2.0',
      id,
      method: 'session/request_permission',
      params: { sessionId },
    });
  }

  return { proc, stdin, stdout: stdoutEmitter, sendLine, sendResponse, sendError, sendChunk, sendPermissionRequest };
}

function createConfig(overrides?: Partial<AcpClientConfig>): AcpClientConfig {
  return {
    binaryPath: '/usr/bin/gemini',
    model: 'gemini-2.5-flash',
    requestTimeoutMs: 5000,
    initTimeoutMs: 5000,
    ...overrides,
  };
}

describe('AcpClient', () => {
  let client: AcpClient;
  let mock: ReturnType<typeof createMockProcess>;

  beforeEach(() => {
    mock = createMockProcess();
  });

  afterEach(() => {
    client?.dispose();
  });

  function createClient(configOverrides?: Partial<AcpClientConfig>) {
    const config = createConfig({
      spawn: (() => mock.proc) as unknown as AcpClientConfig['spawn'],
      ...configOverrides,
    });
    client = new AcpClient(config);
    return client;
  }

  /** Helper: auto-respond to the initialize + session/new handshake */
  function autoHandshake(sessionId = 'test-session') {
    mock.stdin.write.mockImplementation((data: string) => {
      const msg = JSON.parse(data.trim()) as { id: number; method: string };
      if (msg.method === 'initialize') {
        // Respond async to simulate real behavior
        setImmediate(() => mock.sendResponse(msg.id, {
          protocolVersion: 1,
          agentInfo: { name: 'gemini-cli', version: '1.0.0' },
        }));
      } else if (msg.method === 'session/new') {
        setImmediate(() => mock.sendResponse(msg.id, { sessionId }));
      }
    });
  }

  describe('ensureReady', () => {
    it('spawns process and completes ACP handshake', async () => {
      createClient();
      autoHandshake();

      await client.ensureReady();

      expect(client.isAlive).toBe(true);
      // Verify initialize was sent
      const calls = mock.stdin.write.mock.calls.map((c) => JSON.parse(c[0] as string));
      expect(calls[0].method).toBe('initialize');
      expect(calls[0].params.protocolVersion).toBe(1);
      expect(calls[0].params.clientInfo.name).toBe('mcpctl');
      // Verify session/new was sent
      expect(calls[1].method).toBe('session/new');
      expect(calls[1].params.cwd).toBe('/tmp');
      expect(calls[1].params.mcpServers).toEqual([]);
    });

    it('is idempotent when already ready', async () => {
      createClient();
      autoHandshake();

      await client.ensureReady();
      await client.ensureReady();

      // Should only have sent initialize + session/new once
      const calls = mock.stdin.write.mock.calls;
      expect(calls.length).toBe(2);
    });

    it('shares init promise for concurrent calls', async () => {
      createClient();
      autoHandshake();

      const p1 = client.ensureReady();
      const p2 = client.ensureReady();

      await Promise.all([p1, p2]);
      expect(mock.stdin.write.mock.calls.length).toBe(2);
    });
  });

  describe('prompt', () => {
    it('sends session/prompt and collects agent_message_chunk text', async () => {
      createClient();
      const sessionId = 'sess-123';
      autoHandshake(sessionId);

      await client.ensureReady();

      // Now set up the prompt response handler
      mock.stdin.write.mockImplementation((data: string) => {
        const msg = JSON.parse(data.trim()) as { id: number; method: string };
        if (msg.method === 'session/prompt') {
          setImmediate(() => {
            mock.sendChunk(sessionId, 'Hello ');
            mock.sendChunk(sessionId, 'world!');
            mock.sendResponse(msg.id, { stopReason: 'end_turn' });
          });
        }
      });

      const result = await client.prompt('Say hello');
      expect(result).toBe('Hello world!');
    });

    it('handles multi-block content in a single chunk', async () => {
      createClient();
      autoHandshake('sess-1');
      await client.ensureReady();

      mock.stdin.write.mockImplementation((data: string) => {
        const msg = JSON.parse(data.trim()) as { id: number; method: string };
        if (msg.method === 'session/prompt') {
          setImmediate(() => {
            mock.sendLine({
              jsonrpc: '2.0',
              method: 'session/update',
              params: {
                sessionId: 'sess-1',
                update: {
                  sessionUpdate: 'agent_message_chunk',
                  content: [
                    { type: 'text', text: 'Part A' },
                    { type: 'text', text: ' Part B' },
                  ],
                },
              },
            });
            mock.sendResponse(msg.id, { stopReason: 'end_turn' });
          });
        }
      });

      const result = await client.prompt('test');
      expect(result).toBe('Part A Part B');
    });

    it('calls ensureReady automatically (lazy init)', async () => {
      createClient();
      autoHandshake('sess-auto');

      // After handshake, handle prompts
      const originalWrite = mock.stdin.write;
      let handshakeDone = false;
      mock.stdin.write.mockImplementation((data: string) => {
        const msg = JSON.parse(data.trim()) as { id: number; method: string };
        if (msg.method === 'initialize') {
          setImmediate(() => mock.sendResponse(msg.id, { protocolVersion: 1 }));
        } else if (msg.method === 'session/new') {
          setImmediate(() => {
            mock.sendResponse(msg.id, { sessionId: 'sess-auto' });
            handshakeDone = true;
          });
        } else if (msg.method === 'session/prompt') {
          setImmediate(() => {
            mock.sendChunk('sess-auto', 'ok');
            mock.sendResponse(msg.id, { stopReason: 'end_turn' });
          });
        }
      });

      // Call prompt directly without ensureReady
      const result = await client.prompt('test');
      expect(result).toBe('ok');
    });
  });

  describe('auto-restart', () => {
    it('restarts after process exit', async () => {
      createClient();
      autoHandshake('sess-1');
      await client.ensureReady();
      expect(client.isAlive).toBe(true);

      // Simulate process exit
      mock.proc.killed = true;
      mock.proc.emit('exit', 1);
      expect(client.isAlive).toBe(false);

      // Create a new mock for the respawned process
      mock = createMockProcess();
      // Update the spawn function to return new mock
      (client as unknown as { config: { spawn: unknown } }).config.spawn = () => mock.proc;
      autoHandshake('sess-2');

      await client.ensureReady();
      expect(client.isAlive).toBe(true);
    });
  });

  describe('timeout', () => {
    it('kills process and rejects on request timeout', async () => {
      createClient({ requestTimeoutMs: 50 });
      autoHandshake('sess-1');
      await client.ensureReady();

      // Don't respond to the prompt — let it timeout
      mock.stdin.write.mockImplementation(() => {});

      await expect(client.prompt('test')).rejects.toThrow('timed out');
      expect(client.isAlive).toBe(false);
    });

    it('rejects on init timeout', async () => {
      createClient({ initTimeoutMs: 50 });
      // Don't respond to initialize
      mock.stdin.write.mockImplementation(() => {});

      await expect(client.ensureReady()).rejects.toThrow('timed out');
    });
  });

  describe('error handling', () => {
    it('rejects on ACP error response', async () => {
      createClient();
      mock.stdin.write.mockImplementation((data: string) => {
        const msg = JSON.parse(data.trim()) as { id: number; method: string };
        setImmediate(() => mock.sendError(msg.id, -32603, 'Internal error'));
      });

      await expect(client.ensureReady()).rejects.toThrow('ACP error -32603: Internal error');
    });

    it('rejects pending requests on process crash', async () => {
      createClient();
      autoHandshake('sess-1');
      await client.ensureReady();

      // Override write so prompt sends but gets no response; then crash the process
      mock.stdin.write.mockImplementation(() => {
        // After the prompt is sent, simulate a process crash
        setImmediate(() => {
          mock.proc.killed = true;
          mock.proc.emit('exit', 1);
        });
      });

      const promptPromise = client.prompt('test');
      await expect(promptPromise).rejects.toThrow('process exited');
    });
  });

  describe('permission requests', () => {
    it('rejects session/request_permission from agent', async () => {
      createClient();
      autoHandshake('sess-1');
      await client.ensureReady();

      mock.stdin.write.mockImplementation((data: string) => {
        const msg = JSON.parse(data.trim()) as { id: number; method: string };
        if (msg.method === 'session/prompt') {
          setImmediate(() => {
            // Agent asks for permission first
            mock.sendPermissionRequest(100, 'sess-1');
            // Then provides the actual response
            mock.sendChunk('sess-1', 'done');
            mock.sendResponse(msg.id, { stopReason: 'end_turn' });
          });
        }
      });

      const result = await client.prompt('test');
      expect(result).toBe('done');

      // Verify we sent a rejection for the permission request
      const writes = mock.stdin.write.mock.calls.map((c) => {
        try { return JSON.parse(c[0] as string); } catch { return null; }
      }).filter(Boolean);
      const rejection = writes.find((w: Record<string, unknown>) => w.id === 100);
      expect(rejection).toBeTruthy();
      expect((rejection as { result: { outcome: { outcome: string } } }).result.outcome.outcome).toBe('cancelled');
    });
  });

  describe('dispose', () => {
    it('kills process and rejects pending', async () => {
      createClient();
      autoHandshake('sess-1');
      await client.ensureReady();

      // Override write so prompt is sent but gets no response; then dispose
      mock.stdin.write.mockImplementation(() => {
        setImmediate(() => client.dispose());
      });

      const promptPromise = client.prompt('test');
      await expect(promptPromise).rejects.toThrow('disposed');
      expect(mock.proc.kill).toHaveBeenCalledWith('SIGTERM');
    });

    it('is safe to call multiple times', () => {
      createClient();
      client.dispose();
      client.dispose();
      // No error thrown
    });
  });

  describe('isAlive', () => {
    it('returns false before init', () => {
      createClient();
      expect(client.isAlive).toBe(false);
    });

    it('returns true after successful init', async () => {
      createClient();
      autoHandshake();
      await client.ensureReady();
      expect(client.isAlive).toBe(true);
    });

    it('returns false after dispose', async () => {
      createClient();
      autoHandshake();
      await client.ensureReady();
      client.dispose();
      expect(client.isAlive).toBe(false);
    });
  });
});