@aria-cli/tools 1.0.0

package/src/mcp/connection.ts

@@ -0,0 +1,568 @@
+/**
+ * MCP Server Connection
+ *
+ * Manages lifecycle and protocol methods for a single MCP server.
+ */
+
+import type { ChildProcess } from "child_process";
+import { spawn } from "child_process";
+import { EventEmitter } from "events";
+import { PassThrough, Writable } from "stream";
+import { JSONRPCClient } from "./jsonrpc.js";
+import type {
+  MCPServerConfig,
+  MCPCapabilities,
+  MCPInitializeResult,
+  MCPTool,
+  MCPToolCallResult,
+  MCPResource,
+  MCPResourceContent,
+  MCPPrompt,
+  MCPPromptMessage,
+} from "./types.js";
+import type { ToolResult } from "../types.js";
+
+function sanitizeGitEnv(
+  env: NodeJS.ProcessEnv | Record<string, string | undefined> = process.env,
+): Record<string, string> {
+  const sanitized: Record<string, string> = {};
+  for (const [key, value] of Object.entries(env)) {
+    if (value === undefined) continue;
+    if (key.toUpperCase().startsWith("GIT_")) continue;
+    sanitized[key] = value;
+  }
+  return sanitized;
+}
+
+export class MCPServerConnection extends EventEmitter {
+  private process?: ChildProcess;
+  private transport: "stdio" | "sse" = "stdio";
+  private rpc: JSONRPCClient;
+  private capabilities: MCPCapabilities = {};
+  private _initialized = false;
+  private sseAbortController?: AbortController;
+  private sseStreamPromise?: Promise<void>;
+  private ssePostUrl?: string;
+  private sseSessionId?: string;
+  private sseOutput?: PassThrough;
+
+  // Auto-reconnect state
+  private _disconnecting = false;
+  private _reconnecting = false;
+  private _reconnectAttempts = 0;
+
+  private static readonly RECONNECT_BASE_DELAY_MS = 1000;
+  private static readonly RECONNECT_MAX_DELAY_MS = 30000;
+  private static readonly RECONNECT_MAX_ATTEMPTS = 5;
+
+  constructor(private config: MCPServerConfig) {
+    super();
+    this.rpc = new JSONRPCClient();
+  }
+
+  get name(): string {
+    return this.config.name;
+  }
+
+  get initialized(): boolean {
+    return this._initialized;
+  }
+
+  /**
+   * Initialize connection to MCP server
+   */
+  async initialize(): Promise<MCPCapabilities> {
+    // Reset disconnecting flag on fresh initialization so auto-reconnect is active.
+    // During reconnect, attemptReconnect() checks _disconnecting before each attempt.
+    if (!this._reconnecting) {
+      this._disconnecting = false;
+      this._reconnectAttempts = 0;
+    }
+
+    this.transport = this.config.transport ?? "stdio";
+
+    try {
+      if (this.transport === "stdio") {
+        await this.initializeStdioTransport();
+      } else {
+        await this.initializeSSETransport();
+      }
+
+      // Forward notifications
+      this.rpc.on("notifications/tools/list_changed", () => this.emit("tools/list_changed"));
+      this.rpc.on("notifications/resources/list_changed", () =>
+        this.emit("resources/list_changed"),
+      );
+      this.rpc.on("notifications/prompts/list_changed", () => this.emit("prompts/list_changed"));
+
+      // MCP initialize handshake
+      const result = await this.rpc.request<MCPInitializeResult>("initialize", {
+        protocolVersion: "2024-11-05",
+        capabilities: {
+          roots: { listChanged: true },
+        },
+        clientInfo: {
+          name: "aria-cli",
+          version: "1.0.0",
+        },
+      });
+
+      this.capabilities = result.capabilities || {};
+      this._initialized = true;
+
+      // Send initialized notification
+      this.rpc.notify("notifications/initialized", {});
+
+      return this.capabilities;
+    } catch (error) {
+      await this.cleanupAfterInitializeFailure();
+      throw error;
+    }
+  }
+
+  /**
+   * Attempt to reconnect with exponential backoff.
+   * Called when the transport drops unexpectedly (not via explicit shutdown).
+   */
+  private async attemptReconnect(): Promise<void> {
+    if (this._disconnecting || this._reconnecting) {
+      return;
+    }
+
+    this._reconnecting = true;
+
+    while (
+      this._reconnectAttempts < MCPServerConnection.RECONNECT_MAX_ATTEMPTS &&
+      !this._disconnecting
+    ) {
+      this._reconnectAttempts++;
+
+      const delay = Math.min(
+        MCPServerConnection.RECONNECT_BASE_DELAY_MS * Math.pow(2, this._reconnectAttempts - 1),
+        MCPServerConnection.RECONNECT_MAX_DELAY_MS,
+      );
+
+      this.emit("log", {
+        level: "info",
+        message: `Reconnect attempt ${this._reconnectAttempts}/${MCPServerConnection.RECONNECT_MAX_ATTEMPTS} in ${delay}ms`,
+      });
+
+      await new Promise((resolve) => setTimeout(resolve, delay));
+
+      if (this._disconnecting) {
+        break;
+      }
+
+      try {
+        // Clean up stale RPC state before reconnecting
+        this.rpc.close();
+        this.rpc = new JSONRPCClient();
+
+        await this.initialize();
+
+        // Success — reset reconnect state
+        this._reconnectAttempts = 0;
+        this._reconnecting = false;
+        this.emit("reconnected");
+        return;
+      } catch (error) {
+        this.emit("log", {
+          level: "error",
+          message: `Reconnect attempt ${this._reconnectAttempts} failed: ${(error as Error).message}`,
+        });
+      }
+    }
+
+    // All attempts exhausted or disconnecting flag was set
+    this._reconnecting = false;
+    if (!this._disconnecting) {
+      this._initialized = false;
+      this.emit(
+        "error",
+        new Error(
+          `Failed to reconnect after ${MCPServerConnection.RECONNECT_MAX_ATTEMPTS} attempts`,
+        ),
+      );
+    }
+  }
+
+  private async initializeStdioTransport(): Promise<void> {
+    if (!this.config.command) {
+      throw new Error("Command required for stdio transport");
+    }
+
+    // Spawn server process
+    this.process = spawn(this.config.command, this.config.args || [], {
+      env: { ...sanitizeGitEnv(), ...sanitizeGitEnv(this.config.env ?? {}) },
+      stdio: ["pipe", "pipe", "pipe"],
+    });
+
+    if (!this.process.stdin || !this.process.stdout) {
+      throw new Error("Failed to create stdio streams");
+    }
+
+    this.rpc.connect(this.process.stdin, this.process.stdout);
+
+    // Handle stderr for logging
+    this.process.stderr?.on("data", (data) => {
+      this.emit("log", { level: "error", message: data.toString() });
+    });
+
+    this.process.on("exit", (code) => {
+      this._initialized = false;
+      this.rpc.close();
+      this.emit("exit", code);
+
+      // Attempt auto-reconnect on unexpected exit
+      if (!this._disconnecting) {
+        void this.attemptReconnect();
+      }
+    });
+  }
+
+  private async initializeSSETransport(): Promise<void> {
+    const url = this.config.url;
+    if (!url) {
+      throw new Error("URL required for sse transport");
+    }
+
+    const stdin = new Writable({
+      write: (chunk, _encoding, callback) => {
+        const payload = chunk.toString();
+        this.postSSEMessages(payload)
+          .then(() => callback())
+          .catch((err) => callback(err as Error));
+      },
+    });
+    const stdout = new PassThrough();
+    this.sseOutput = stdout;
+    this.rpc.connect(stdin, stdout);
+
+    this.sseAbortController = new AbortController();
+    const response = await fetch(url, {
+      method: "GET",
+      headers: { Accept: "text/event-stream" },
+      signal: this.sseAbortController.signal,
+    });
+    if (!response.ok || !response.body) {
+      const errorText = await response.text();
+      throw new Error(
+        `Failed to connect SSE transport: ${response.status} ${response.statusText} ${errorText}`,
+      );
+    }
+
+    this.sseSessionId = response.headers.get("mcp-session-id") ?? undefined;
+    this.ssePostUrl = url;
+    this.sseStreamPromise = this.consumeSSE(response.body, url);
+  }
+
+  private async postSSEMessages(payload: string): Promise<void> {
+    const trimmed = payload.trim();
+    if (!trimmed) {
+      return;
+    }
+
+    const endpoint = this.ssePostUrl || this.config.url;
+    if (!endpoint) {
+      throw new Error("SSE transport endpoint is not configured");
+    }
+
+    const messages = trimmed
+      .split("\n")
+      .map((line) => line.trim())
+      .filter((line) => line.length > 0);
+
+    for (const message of messages) {
+      const response = await fetch(endpoint, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          ...(this.sseSessionId ? { "mcp-session-id": this.sseSessionId } : {}),
+        },
+        body: message,
+      });
+
+      if (!response.ok) {
+        const errorText = await response.text();
+        throw new Error(
+          `SSE transport POST failed: ${response.status} ${response.statusText} ${errorText}`,
+        );
+      }
+
+      const responseText = (await response.text()).trim();
+      if (!responseText || !this.sseOutput) {
+        continue;
+      }
+
+      // Some servers return direct JSON-RPC responses on POST while others
+      // publish responses only via SSE. Forward direct JSON replies into RPC.
+      for (const line of responseText.split("\n")) {
+        const jsonLine = line.trim();
+        if (!jsonLine) {
+          continue;
+        }
+        if (jsonLine.startsWith("{") && jsonLine.endsWith("}")) {
+          this.sseOutput.write(`${jsonLine}\n`);
+        }
+      }
+    }
+  }
+
+  private async consumeSSE(body: ReadableStream<Uint8Array>, baseUrl: string): Promise<void> {
+    const reader = body.getReader();
+    const decoder = new TextDecoder();
+    let buffer = "";
+    let eventType = "";
+    let dataLines: string[] = [];
+
+    const dispatch = (): void => {
+      if (dataLines.length === 0) {
+        eventType = "";
+        return;
+      }
+
+      const data = dataLines.join("\n").trim();
+      const event = eventType || "message";
+      eventType = "";
+      dataLines = [];
+
+      if (!data) {
+        return;
+      }
+
+      if (event === "endpoint") {
+        this.ssePostUrl = new URL(data, baseUrl).toString();
+        return;
+      }
+
+      if (!this.sseOutput) {
+        return;
+      }
+
+      if (data.startsWith("{") && data.endsWith("}")) {
+        this.sseOutput.write(`${data}\n`);
+      }
+    };
+
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          dispatch();
+          break;
+        }
+
+        buffer += decoder.decode(value, { stream: true });
+        const lines = buffer.split(/\r?\n/);
+        buffer = lines.pop() || "";
+
+        for (const line of lines) {
+          if (line === "") {
+            dispatch();
+            continue;
+          }
+          if (line.startsWith("event:")) {
+            eventType = line.slice(6).trim();
+            continue;
+          }
+          if (line.startsWith("data:")) {
+            dataLines.push(line.slice(5).trim());
+          }
+        }
+      }
+    } catch (error) {
+      // Ignore abort errors from normal shutdown.
+      if (!(error instanceof DOMException && error.name === "AbortError")) {
+        this.emit("log", {
+          level: "error",
+          message: `SSE stream error: ${(error as Error).message}`,
+        });
+      }
+    } finally {
+      if (this._initialized) {
+        this._initialized = false;
+        this.rpc.close();
+        this.emit("exit", 0);
+
+        // Attempt auto-reconnect on unexpected SSE stream end
+        if (!this._disconnecting) {
+          void this.attemptReconnect();
+        }
+      }
+    }
+  }
+
+  private async cleanupAfterInitializeFailure(): Promise<void> {
+    if (this.transport === "stdio") {
+      if (this.process && !this.process.killed) {
+        try {
+          this.process.kill("SIGTERM");
+        } catch {
+          // Ignore kill failures during cleanup.
+        }
+      }
+      this.process = undefined;
+    } else {
+      this.sseAbortController?.abort();
+      try {
+        await this.sseStreamPromise;
+      } catch {
+        // Ignore stream errors during cleanup.
+      }
+      this.sseOutput?.end();
+      this.sseOutput = undefined;
+    }
+    this.rpc.close();
+  }
+
+  /**
+   * Gracefully shutdown the connection
+   */
+  async shutdown(): Promise<void> {
+    this._disconnecting = true;
+    if (!this._initialized) return;
+    this._initialized = false;
+    if (this.transport === "sse") {
+      try {
+        await this.rpc.request("shutdown", {});
+        this.rpc.notify("notifications/exit", {});
+      } catch {
+        // Best effort
+      }
+      this.sseAbortController?.abort();
+      try {
+        await this.sseStreamPromise;
+      } catch {
+        // Ignore stream shutdown errors
+      }
+      this.sseOutput?.end();
+      this.sseOutput = undefined;
+      this.rpc.close();
+      return;
+    }
+
+    if (this.process) {
+      // Send MCP shutdown/exit protocol before killing
+      try {
+        await this.rpc.request("shutdown", {});
+        this.rpc.notify("notifications/exit", {});
+        // Give process a moment to clean up
+        await new Promise((resolve) => setTimeout(resolve, 500));
+      } catch {
+        // Process may already be dead, proceed with kill
+      }
+      this.process.kill("SIGTERM");
+
+      // Await process exit before resolving
+      await new Promise<void>((resolve) => {
+        // Escalate to SIGKILL if process doesn't exit within 5s
+        const killTimeout = setTimeout(() => {
+          try {
+            this.process?.kill("SIGKILL");
+          } catch {
+            /* already dead */
+          }
+        }, 5000);
+        this.process!.once("exit", () => {
+          clearTimeout(killTimeout);
+          this.rpc.close();
+          resolve();
+        });
+      });
+    } else {
+      this.rpc.close();
+    }
+  }
+
+  // ========== Tools ==========
+
+  async listTools(): Promise<MCPTool[]> {
+    const result = await this.rpc.request<{ tools: MCPTool[] }>("tools/list", {});
+    return result.tools || [];
+  }
+
+  async callTool(name: string, args: unknown, abortSignal?: AbortSignal): Promise<ToolResult> {
+    // Early abort check before making the RPC call
+    if (abortSignal?.aborted) {
+      return { success: false, message: "Tool execution cancelled by user." };
+    }
+
+    try {
+      // Race the RPC call against the abort signal so user cancel
+      // doesn't block waiting for a slow MCP server response.
+      const rpcPromise = this.rpc.request<MCPToolCallResult>("tools/call", {
+        name,
+        arguments: args,
+      });
+
+      let result: MCPToolCallResult;
+      if (abortSignal) {
+        result = await new Promise<MCPToolCallResult>((resolve, reject) => {
+          const onAbort = () => reject(new DOMException("Aborted", "AbortError"));
+          abortSignal.addEventListener("abort", onAbort, { once: true });
+          rpcPromise.then(
+            (r) => {
+              abortSignal.removeEventListener("abort", onAbort);
+              resolve(r);
+            },
+            (e) => {
+              abortSignal.removeEventListener("abort", onAbort);
+              reject(e);
+            },
+          );
+        });
+      } else {
+        result = await rpcPromise;
+      }
+
+      const text =
+        result.content
+          ?.filter((c): c is { type: "text"; text: string } => c.type === "text")
+          .map((c) => c.text)
+          .join("\n") || "";
+
+      return {
+        success: !result.isError,
+        message: text,
+        data: result.content,
+      };
+    } catch (error) {
+      if (error instanceof DOMException && error.name === "AbortError") {
+        return { success: false, message: "Tool execution cancelled by user." };
+      }
+      return {
+        success: false,
+        message: (error as Error).message,
+      };
+    }
+  }
+
+  // ========== Resources ==========
+
+  async listResources(): Promise<MCPResource[]> {
+    const result = await this.rpc.request<{ resources: MCPResource[] }>("resources/list", {});
+    return result.resources || [];
+  }
+
+  async readResource(uri: string): Promise<MCPResourceContent[]> {
+    const result = await this.rpc.request<{ contents: MCPResourceContent[] }>("resources/read", {
+      uri,
+    });
+    return result.contents || [];
+  }
+
+  // ========== Prompts ==========
+
+  async listPrompts(): Promise<MCPPrompt[]> {
+    const result = await this.rpc.request<{ prompts: MCPPrompt[] }>("prompts/list", {});
+    return result.prompts || [];
+  }
+
+  async getPrompt(name: string, args?: Record<string, string>): Promise<MCPPromptMessage[]> {
+    const result = await this.rpc.request<{ messages: MCPPromptMessage[] }>("prompts/get", {
+      name,
+      arguments: args,
+    });
+    return result.messages || [];
+  }
+}

package/src/mcp/index.ts

@@ -0,0 +1,11 @@
+/**
+ * MCP (Model Context Protocol) Module
+ *
+ * Provides client implementation for connecting to MCP servers
+ * and discovering tools, resources, and prompts dynamically.
+ */
+
+export * from './types.js';
+export { JSONRPCClient } from './jsonrpc.js';
+export { MCPServerConnection } from './connection.js';
+export { MCPClient } from './client.js';