stack/apps/api/src/mcp/stdio-transport.ts

import { spawn, type ChildProcess } from "node:child_process";
import { Logger } from "@nestjs/common";
import type { McpRequest, McpResponse } from "./interfaces";

/**
 * STDIO transport for MCP server communication
 * Spawns a child process and communicates via stdin/stdout using JSON-RPC 2.0
 */
export class StdioTransport {
  private readonly logger = new Logger(StdioTransport.name);
  private process?: ChildProcess;
  private pendingRequests = new Map<
    string | number,
    { resolve: (value: McpResponse) => void; reject: (error: Error) => void }
  >();
  private buffer = "";

  constructor(
    private readonly command: string,
    private readonly args?: string[],
    private readonly env?: Record<string, string>
  ) {}

  /**
   * Start the child process
   */
  async start(): Promise<void> {
    if (this.isRunning()) {
      return;
    }

    return new Promise((resolve, reject) => {
      try {
        this.process = spawn(this.command, this.args ?? [], {
          env: { ...process.env, ...this.env },
          stdio: ["pipe", "pipe", "pipe"],
        });

        this.process.stdout?.on("data", (data: Buffer) => {
          this.handleStdout(data);
        });

        this.process.stderr?.on("data", (data: Buffer) => {
          this.logger.warn(`MCP stderr: ${data.toString()}`);
        });

        this.process.on("error", (error) => {
          this.handleProcessError(error);
          reject(error);
        });

        this.process.on("exit", (code) => {
          this.handleProcessExit(code);
        });

        // Resolve immediately after spawn
        resolve();
      } catch (error: unknown) {
        reject(error instanceof Error ? error : new Error(String(error)));
      }
    });
  }

  /**
   * Send a request and wait for response
   */
  async send(request: McpRequest): Promise<McpResponse> {
    if (!this.isRunning()) {
      throw new Error("Process not running");
    }

    return new Promise((resolve, reject) => {
      this.pendingRequests.set(request.id, { resolve, reject });

      const message = JSON.stringify(request) + "\n";
      this.process?.stdin?.write(message, (error) => {
        if (error) {
          this.pendingRequests.delete(request.id);
          reject(error);
        }
      });
    });
  }

  /**
   * Stop the child process
   */
  async stop(): Promise<void> {
    if (!this.isRunning()) {
      return;
    }

    return new Promise((resolve) => {
      if (!this.process) {
        resolve();
        return;
      }

      this.process.once("exit", () => {
        delete this.process;
        resolve();
      });

      // Reject all pending requests
      this.rejectAllPending(new Error("Process stopped"));

      this.process.kill();
    });
  }

  /**
   * Check if process is running
   */
  isRunning(): boolean {
    return this.process !== undefined && !this.process.killed;
  }

  /**
   * Handle stdout data
   */
  private handleStdout(data: Buffer): void {
    this.buffer += data.toString();

    // Process complete JSON messages (delimited by newlines)
    let newlineIndex: number;
    while ((newlineIndex = this.buffer.indexOf("\n")) !== -1) {
      const message = this.buffer.substring(0, newlineIndex);
      this.buffer = this.buffer.substring(newlineIndex + 1);

      if (message.trim()) {
        try {
          const response = JSON.parse(message) as McpResponse;
          this.handleResponse(response);
        } catch (error) {
          this.logger.error("Failed to parse MCP response:", error);
        }
      }
    }
  }

  /**
   * Handle parsed response
   */
  private handleResponse(response: McpResponse): void {
    const pending = this.pendingRequests.get(response.id);
    if (pending) {
      this.pendingRequests.delete(response.id);
      pending.resolve(response);
    }
  }

  /**
   * Handle process error
   */
  private handleProcessError(error: Error): void {
    this.rejectAllPending(error);
    delete this.process;
  }

  /**
   * Handle process exit
   */
  private handleProcessExit(code: number | null): void {
    const exitCode = code !== null ? String(code) : "unknown";
    this.rejectAllPending(new Error(`Process exited with code ${exitCode}`));
    delete this.process;
  }

  /**
   * Reject all pending requests
   */
  private rejectAllPending(error: Error): void {
    for (const pending of this.pendingRequests.values()) {
      pending.reject(error);
    }
    this.pendingRequests.clear();
  }
}