mcpctl/src/cli/tests/commands/mcp.test.ts

import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import http from 'node:http';
import { Readable, Writable } from 'node:stream';
import { runMcpBridge } from '../../src/commands/mcp.js';

// ---- Mock MCP server (simulates mcplocal project endpoint) ----

interface RecordedRequest {
  method: string;
  url: string;
  headers: http.IncomingHttpHeaders;
  body: string;
}

let mockServer: http.Server;
let mockPort: number;
const recorded: RecordedRequest[] = [];
let sessionCounter = 0;

function makeInitializeResponse(id: number | string) {
  return JSON.stringify({
    jsonrpc: '2.0',
    id,
    result: {
      protocolVersion: '2024-11-05',
      capabilities: { tools: {} },
      serverInfo: { name: 'test-server', version: '1.0.0' },
    },
  });
}

function makeToolsListResponse(id: number | string) {
  return JSON.stringify({
    jsonrpc: '2.0',
    id,
    result: {
      tools: [
        { name: 'grafana/query', description: 'Query Grafana', inputSchema: { type: 'object', properties: {} } },
      ],
    },
  });
}

function makeToolCallResponse(id: number | string) {
  return JSON.stringify({
    jsonrpc: '2.0',
    id,
    result: {
      content: [{ type: 'text', text: 'tool result' }],
    },
  });
}

beforeAll(async () => {
  mockServer = http.createServer((req, res) => {
    const chunks: Buffer[] = [];
    req.on('data', (c: Buffer) => chunks.push(c));
    req.on('end', () => {
      const body = Buffer.concat(chunks).toString('utf-8');
      recorded.push({ method: req.method ?? '', url: req.url ?? '', headers: req.headers, body });

      if (req.method === 'DELETE') {
        res.writeHead(200);
        res.end();
        return;
      }

      if (req.method === 'POST' && req.url?.startsWith('/projects/')) {
        let sessionId = req.headers['mcp-session-id'] as string | undefined;

        // Assign session ID on first request
        if (!sessionId) {
          sessionCounter++;
          sessionId = `session-${sessionCounter}`;
        }
        res.setHeader('mcp-session-id', sessionId);

        // Parse JSON-RPC and respond based on method
        try {
          const rpc = JSON.parse(body) as { id: number | string; method: string };
          let responseBody: string;

          switch (rpc.method) {
            case 'initialize':
              responseBody = makeInitializeResponse(rpc.id);
              break;
            case 'tools/list':
              responseBody = makeToolsListResponse(rpc.id);
              break;
            case 'tools/call':
              responseBody = makeToolCallResponse(rpc.id);
              break;
            default:
              responseBody = JSON.stringify({ jsonrpc: '2.0', id: rpc.id, error: { code: -32601, message: 'Method not found' } });
          }

          res.writeHead(200, { 'Content-Type': 'application/json' });
          res.end(responseBody);
        } catch {
          res.writeHead(400, { 'Content-Type': 'application/json' });
          res.end(JSON.stringify({ error: 'Invalid JSON' }));
        }
        return;
      }

      res.writeHead(404);
      res.end();
    });
  });

  await new Promise<void>((resolve) => {
    mockServer.listen(0, () => {
      const addr = mockServer.address();
      if (addr && typeof addr === 'object') {
        mockPort = addr.port;
      }
      resolve();
    });
  });
});

afterAll(() => {
  mockServer.close();
});

// ---- Helper to run bridge with mock streams ----

function createMockStreams() {
  const stdoutChunks: string[] = [];
  const stderrChunks: string[] = [];

  const stdout = new Writable({
    write(chunk: Buffer, _encoding, callback) {
      stdoutChunks.push(chunk.toString());
      callback();
    },
  });

  const stderr = new Writable({
    write(chunk: Buffer, _encoding, callback) {
      stderrChunks.push(chunk.toString());
      callback();
    },
  });

  return { stdout, stderr, stdoutChunks, stderrChunks };
}

function pushAndEnd(stdin: Readable, lines: string[]) {
  for (const line of lines) {
    stdin.push(line + '\n');
  }
  stdin.push(null); // EOF
}

// ---- Tests ----

describe('MCP STDIO Bridge', () => {
  beforeAll(() => {
    recorded.length = 0;
    sessionCounter = 0;
  });

  it('forwards initialize request and returns response', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout, stdoutChunks } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });

    pushAndEnd(stdin, [initMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    // Verify request was made to correct URL
    expect(recorded.some((r) => r.url === '/projects/test-project/mcp' && r.method === 'POST')).toBe(true);

    // Verify response on stdout
    const output = stdoutChunks.join('');
    const parsed = JSON.parse(output.trim());
    expect(parsed.result.serverInfo.name).toBe('test-server');
    expect(parsed.result.protocolVersion).toBe('2024-11-05');
  });

  it('sends session ID on subsequent requests', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout, stdoutChunks } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });
    const toolsListMsg = JSON.stringify({ jsonrpc: '2.0', id: 2, method: 'tools/list', params: {} });

    pushAndEnd(stdin, [initMsg, toolsListMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    // First POST should NOT have mcp-session-id header
    const firstPost = recorded.find((r) => r.method === 'POST' && r.body.includes('initialize'));
    expect(firstPost).toBeDefined();
    expect(firstPost!.headers['mcp-session-id']).toBeUndefined();

    // Second POST SHOULD have mcp-session-id header
    const secondPost = recorded.find((r) => r.method === 'POST' && r.body.includes('tools/list'));
    expect(secondPost).toBeDefined();
    expect(secondPost!.headers['mcp-session-id']).toMatch(/^session-/);

    // Verify tools/list response
    const lines = stdoutChunks.join('').trim().split('\n');
    expect(lines.length).toBe(2);
    const toolsResponse = JSON.parse(lines[1]);
    expect(toolsResponse.result.tools[0].name).toBe('grafana/query');
  });

  it('forwards tools/call and returns result', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout, stdoutChunks } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });
    const callMsg = JSON.stringify({
      jsonrpc: '2.0', id: 2, method: 'tools/call',
      params: { name: 'grafana/query', arguments: { query: 'test' } },
    });

    pushAndEnd(stdin, [initMsg, callMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    const lines = stdoutChunks.join('').trim().split('\n');
    expect(lines.length).toBe(2);
    const callResponse = JSON.parse(lines[1]);
    expect(callResponse.result.content[0].text).toBe('tool result');
  });

  it('forwards Authorization header when token provided', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });

    pushAndEnd(stdin, [initMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      token: 'my-secret-token',
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    const post = recorded.find((r) => r.method === 'POST');
    expect(post).toBeDefined();
    expect(post!.headers['authorization']).toBe('Bearer my-secret-token');
  });

  it('does not send Authorization header when no token', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });

    pushAndEnd(stdin, [initMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    const post = recorded.find((r) => r.method === 'POST');
    expect(post).toBeDefined();
    expect(post!.headers['authorization']).toBeUndefined();
  });

  it('sends DELETE to clean up session on stdin EOF', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });

    pushAndEnd(stdin, [initMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    // Should have a DELETE request for session cleanup
    const deleteReq = recorded.find((r) => r.method === 'DELETE');
    expect(deleteReq).toBeDefined();
    expect(deleteReq!.headers['mcp-session-id']).toMatch(/^session-/);
  });

  it('does not send DELETE if no session was established', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout } = createMockStreams();

    // Push EOF immediately with no messages
    stdin.push(null);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    expect(recorded.filter((r) => r.method === 'DELETE')).toHaveLength(0);
  });

  it('writes errors to stderr, not stdout', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout, stdoutChunks, stderr, stderrChunks } = createMockStreams();

    // Send to a non-existent port to trigger connection error
    const badMsg = JSON.stringify({ jsonrpc: '2.0', id: 1, method: 'initialize', params: {} });
    pushAndEnd(stdin, [badMsg]);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: 'http://localhost:1', // will fail to connect
      stdin, stdout, stderr,
    });

    // Error should be on stderr
    expect(stderrChunks.join('')).toContain('MCP bridge error');
    // stdout should be empty (no corrupted output)
    expect(stdoutChunks.join('')).toBe('');
  });

  it('skips blank lines in stdin', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout, stdoutChunks } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });

    pushAndEnd(stdin, ['', '  ', initMsg, '']);

    await runMcpBridge({
      projectName: 'test-project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr: new Writable({ write(_, __, cb) { cb(); } }),
    });

    // Only one POST (for the actual message)
    const posts = recorded.filter((r) => r.method === 'POST');
    expect(posts).toHaveLength(1);

    // One response line
    const lines = stdoutChunks.join('').trim().split('\n');
    expect(lines).toHaveLength(1);
  });

  it('URL-encodes project name', async () => {
    recorded.length = 0;
    const stdin = new Readable({ read() {} });
    const { stdout } = createMockStreams();
    const { stderr } = createMockStreams();

    const initMsg = JSON.stringify({
      jsonrpc: '2.0', id: 1, method: 'initialize',
      params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1.0' } },
    });

    pushAndEnd(stdin, [initMsg]);

    await runMcpBridge({
      projectName: 'my project',
      mcplocalUrl: `http://localhost:${mockPort}`,
      stdin, stdout, stderr,
    });

    const post = recorded.find((r) => r.method === 'POST');
    expect(post?.url).toBe('/projects/my%20project/mcp');
  });
});