@pi-archimedes/mcp 2.3.0

package/src/index.ts

@@ -0,0 +1,146 @@
+import type { ExtensionAPI, Theme } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import { loadMcpConfig, loadServerDefs, loadAllServerDefs } from "./config.js";
+import { registerMcpCommand } from "./commands.js";
+import { LifecycleManager } from "./lifecycle.js";
+import { getCachedTools, getCachedPrompts } from "./metadata-cache.js";
+import { ServerManager } from "./server-manager.js";
+import type { McpConfig, ServerDef } from "./types.js";
+import { renderProxyCall, renderProxyResult, type RenderContext } from "./renderer.js";
+import { buildProxyToolExecute, buildSessionStartHandler } from "./proxy-tool.js";
+
+// ── Module-level server manager (survives session restarts) ─────────────────
+
+let manager = new ServerManager();
+let _idleTimeoutMinutes = 10; // updated in session_start from config
+
+// ── Test seams ──────────────────────────────────────────────────────────────
+// Unit tests swap the module-level manager and the config loaders so the
+// proxy's execute() can be exercised without touching real files or
+// spawning real server processes.
+let seamDefs: (() => Record<string, ServerDef>) | null = null;
+let seamAllDefs: (() => Record<string, ServerDef>) | null = null;
+let seamConfig: (() => McpConfig) | null = null;
+
+/** Test-only: replace the module-level manager and/or config loaders. */
+export function setIndexSeamsForTest(seams: {
+  manager?: ServerManager;
+  loadServerDefs?: () => Record<string, ServerDef>;
+  loadAllServerDefs?: () => Record<string, ServerDef>;
+  loadMcpConfig?: () => McpConfig;
+} | null): void {
+  if (seams) {
+    if (seams.manager) manager = seams.manager;
+    seamDefs = seams.loadServerDefs ?? null;
+    seamAllDefs = seams.loadAllServerDefs ?? null;
+    seamConfig = seams.loadMcpConfig ?? null;
+  } else {
+    seamDefs = null;
+    seamAllDefs = null;
+    seamConfig = null;
+    // Fully restore initial state: discard any swapped-in manager (its fake
+    // clients/state must not leak into subsequent tests) and rebind the
+    // lifecycle to the fresh module-level manager so both agree again.
+    manager = new ServerManager();
+    lifecycle.setManager(manager);
+  }
+}
+
+/** Config-loader indirection: seam overrides win, real loaders otherwise */
+const loadDefs = (): Record<string, ServerDef> => seamDefs?.() ?? loadServerDefs();
+const loadAllDefs = (): Record<string, ServerDef> => seamAllDefs?.() ?? loadAllServerDefs();
+const loadCfg = (): McpConfig => seamConfig?.() ?? loadMcpConfig();
+
+// Health-check interval: reconnects downed keep-alive servers and closes
+// idle non-keep-alive servers past their idleTimeout. unref'd, so it never
+// keeps the process alive; stopped on session_shutdown.
+const lifecycle = new LifecycleManager(manager, loadServerDefs, () => _idleTimeoutMinutes);
+
+// ── Tool registration ───────────────────────────────────────────────────────
+
+export function registerMcp(pi: ExtensionAPI): void {
+  // The /mcp command family (plan-027 Task 2): /mcp [status|tools|prompts|
+  // reconnect|enable|disable|logout|auth|panel|setup]. The command registry
+  // contains ONLY "mcp" — the standalone /mcp-auth + /mcp-logout commands
+  // are retired (their bodies live in commands-auth.ts, called via
+  // /mcp auth / /mcp logout).
+  registerMcpCommand(pi, {
+    getManager: () => manager,
+    // loadAllServerDefs includes disabled servers — they still have status.
+    // seamAllDefs indirection: real loader in production, stable fixture in
+    // tests (the same seam-discipline as loadDefs/loadCfg).
+    getServerDefs: () => loadAllDefs(),
+    getCachedTools,
+    getCachedPrompts,
+  });
+
+  // session_shutdown must be registered at the TOP LEVEL of registerMcp,
+  // NOT inside session_start — per AGENTS.md rule (prevents handler accumulation on /reload)
+  pi.on("session_shutdown", async () => {
+    lifecycle.stop();
+    await manager.closeAll();
+  });
+
+  pi.on(
+    "session_start",
+    buildSessionStartHandler({
+      pi,
+      getManager: () => manager,
+      loadDefs,
+      loadCfg,
+      setIdleTimeout: (m) => { _idleTimeoutMinutes = m; },
+      startLifecycle: () => lifecycle.start(),
+    }),
+  );
+
+  pi.registerTool({
+    name: "mcp",
+    label: "MCP",
+    description: [
+      "Gateway to MCP (Model Context Protocol) servers. Use this tool to discover and call tools from connected MCP servers.",
+      "",
+      "Workflow:",
+      "1. Search: mcp({ search: 'keyword' }) — find available tools",
+      "2. Describe: mcp({ describe: 'tool_name' }) — see full parameters",
+      "3. Call: mcp({ tool: 'tool_name', args: { ... } }) — execute the tool",
+      "",
+      "Other actions:",
+      "  Status:  mcp({}) or mcp({ action: 'status' }) — list all servers and their connection status",
+      "  List:    mcp({ server: 'name' }) — list all tools on a specific server",
+      "  Connect: mcp({ connect: 'name' }) — eagerly connect to a server",
+      "",
+      "Use 'server' to disambiguate when two servers export a tool with the same name.",
+    ].join("\n"),
+
+    parameters: Type.Object({
+      tool: Type.Optional(Type.String({ description: "Tool name to call" })),
+      args: Type.Optional(
+        Type.Union([
+          Type.String({ description: "Tool arguments as a JSON string" }),
+          Type.Object({}, { additionalProperties: true, description: "Tool arguments as an object" }),
+        ]),
+      ),
+      search: Type.Optional(Type.String({ description: "Search tools by name/description keyword" })),
+      describe: Type.Optional(Type.String({ description: "Tool name to show full parameter schema for" })),
+      connect: Type.Optional(Type.String({ description: "Server name to eagerly connect" })),
+      server: Type.Optional(Type.String({ description: "Filter to a specific server (for list, search, or disambiguating calls)" })),
+      action: Type.Optional(Type.String({ description: "Action string (e.g. 'status')" })),
+    }),
+
+    renderCall(args: unknown, theme: Theme, context: unknown) {
+      return renderProxyCall(args as Record<string, unknown>, theme, context as RenderContext);
+    },
+
+    renderResult(result: unknown, options: unknown, theme: Theme, context: unknown) {
+      const typedResult = result as { content: Array<{ type: string; text?: string }>; details?: Record<string, unknown> };
+      const typedOptions = options as { expanded?: boolean; isPartial?: boolean };
+      return renderProxyResult(typedResult, typedOptions, theme, context as RenderContext);
+    },
+
+    execute: buildProxyToolExecute({
+      getManager: () => manager,
+      loadDefs,
+      loadCfg,
+    }),
+  });
+}

package/src/lifecycle.test.ts

@@ -0,0 +1,274 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { ServerClient } from "./server-client.js";
+import type { ServerManager } from "./server-manager.js";
+import { LifecycleManager } from "./lifecycle.js";
+import type { ServerDef } from "./types.js";
+
+/** Interval the lifecycle manager ticks at in production. */
+const TICK_MS = 30_000;
+
+interface FakeClient {
+  name: string;
+  status: "disconnected" | "connecting" | "connected" | "error" | "needs-auth";
+  connect: ReturnType<typeof vi.fn>;
+  close: ReturnType<typeof vi.fn>;
+}
+
+function makeClient(
+  name: string,
+  status: FakeClient["status"] = "disconnected",
+): FakeClient {
+  return {
+    name,
+    status,
+    connect: vi.fn().mockResolvedValue(undefined),
+    close: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+function stdioDef(extra?: Partial<ServerDef>): ServerDef {
+  return { type: "stdio", command: "true", ...extra } as ServerDef;
+}
+
+function makeLifecycle(
+  clients: FakeClient[],
+  defs: Record<string, ServerDef>,
+  globalIdleMin: number,
+  isIdleImpl: (name: string, timeoutMs: number) => boolean,
+) {
+  const isIdle = vi.fn(isIdleImpl);
+  const manager = {
+    getClients: () => clients as unknown as ServerClient[],
+    isIdle,
+  } as unknown as ServerManager;
+  const lifecycle = new LifecycleManager(manager, () => defs, () => globalIdleMin);
+  return { lifecycle, isIdle };
+}
+
+beforeEach(() => {
+  vi.useFakeTimers();
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+describe("LifecycleManager", () => {
+  it("reconnects a disconnected keep-alive server on tick", () => {
+    const client = makeClient("ka");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { ka: stdioDef({ lifecycle: "keep-alive" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(client.connect).toHaveBeenCalledTimes(1);
+    expect(client.close).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("reconnects a disconnected lazy-keep-alive server on tick", () => {
+    const client = makeClient("lka");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { lka: stdioDef({ lifecycle: "lazy-keep-alive" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(client.connect).toHaveBeenCalledTimes(1);
+    lifecycle.stop();
+  });
+
+  it("does not reconnect or close a disconnected non-keep-alive server", () => {
+    const client = makeClient("lazy");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { lazy: stdioDef({ lifecycle: "lazy" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS * 2);
+
+    expect(client.connect).not.toHaveBeenCalled();
+    expect(client.close).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("reconnects a keep-alive server in error status on tick", () => {
+    const client = makeClient("err", "error");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { err: stdioDef({ lifecycle: "keep-alive" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(client.connect).toHaveBeenCalledTimes(1);
+    expect(client.close).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("does not touch a needs-auth keep-alive server", () => {
+    const client = makeClient("auth", "needs-auth");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { auth: stdioDef({ lifecycle: "keep-alive" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(client.connect).not.toHaveBeenCalled();
+    expect(client.close).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("closes a lazy server past its per-server idle timeout", () => {
+    const client = makeClient("lazy");
+    const { lifecycle, isIdle } = makeLifecycle(
+      [client],
+      { lazy: stdioDef({ lifecycle: "lazy", idleTimeout: 5 }) },
+      10,
+      () => true,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(isIdle).toHaveBeenCalledWith("lazy", 5 * 60_000);
+    expect(client.close).toHaveBeenCalledTimes(1);
+    expect(client.connect).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("leaves a lazy server alone when it is not idle yet", () => {
+    const client = makeClient("lazy");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { lazy: stdioDef({ lifecycle: "lazy", idleTimeout: 5 }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS * 2);
+
+    expect(client.close).not.toHaveBeenCalled();
+    expect(client.connect).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("falls back to the global idle timeout when no per-server timeout is set", () => {
+    const client = makeClient("lazy");
+    const { lifecycle, isIdle } = makeLifecycle(
+      [client],
+      { lazy: stdioDef({ lifecycle: "lazy" }) },
+      10,
+      () => true,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(isIdle).toHaveBeenCalledWith("lazy", 10 * 60_000);
+    expect(client.close).toHaveBeenCalledTimes(1);
+    lifecycle.stop();
+  });
+
+  it("idleTimeout 0 disables idle shutdown even when idle", () => {
+    const client = makeClient("lazy");
+    const { lifecycle, isIdle } = makeLifecycle(
+      [client],
+      { lazy: stdioDef({ lifecycle: "lazy", idleTimeout: 0 }) },
+      10,
+      () => true,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(isIdle).not.toHaveBeenCalled();
+    expect(client.close).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("skips servers whose definition was removed", () => {
+    const client = makeClient("gone");
+    const { lifecycle } = makeLifecycle([client], {}, 10, () => true);
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(client.connect).not.toHaveBeenCalled();
+    expect(client.close).not.toHaveBeenCalled();
+    lifecycle.stop();
+  });
+
+  it("stop() clears the timer so no further ticks run", () => {
+    const client = makeClient("ka");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { ka: stdioDef({ lifecycle: "keep-alive" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+    expect(client.connect).toHaveBeenCalledTimes(1);
+
+    lifecycle.stop();
+    vi.advanceTimersByTime(TICK_MS * 4);
+    expect(client.connect).toHaveBeenCalledTimes(1);
+  });
+
+  it("start() is idempotent — calling it twice does not double the ticks", () => {
+    const client = makeClient("ka");
+    const { lifecycle } = makeLifecycle(
+      [client],
+      { ka: stdioDef({ lifecycle: "keep-alive" }) },
+      10,
+      () => false,
+    );
+
+    lifecycle.start();
+    lifecycle.start();
+    vi.advanceTimersByTime(TICK_MS);
+
+    expect(client.connect).toHaveBeenCalledTimes(1);
+    lifecycle.stop();
+  });
+
+  it("unrefs the interval so it does not keep the process alive", () => {
+    vi.useRealTimers();
+    const unref = vi.fn();
+    const spy = vi
+      .spyOn(globalThis, "setInterval")
+      .mockImplementation(() => ({ unref }) as unknown as NodeJS.Timeout);
+
+    const lifecycle = new LifecycleManager(
+      { getClients: () => [], isIdle: () => false } as unknown as ServerManager,
+      () => ({}),
+      () => 10,
+    );
+    lifecycle.start();
+    expect(unref).toHaveBeenCalledTimes(1);
+    lifecycle.stop();
+    spy.mockRestore();
+  });
+});

package/src/lifecycle.ts

@@ -0,0 +1,77 @@
+import type { ServerManager } from "./server-manager.js";
+import type { ServerDef } from "./types.js";
+
+/** Health-check tick interval in production. */
+const TICK_INTERVAL_MS = 30_000;
+
+/**
+ * Periodic resource-lifecycle management for MCP server clients:
+ *
+ * - keep-alive / lazy-keep-alive servers that end up disconnected or in an
+ *   error state are reconnected on the next tick.
+ * - Non-keep-alive servers are closed once they have been idle past their
+ *   `idleTimeout` (per-server, falling back to the global default; 0 disables).
+ *
+ * The interval is `unref`'d so it never keeps the process alive, and `stop()`
+ * is idempotent so repeated session restarts cannot accumulate timers.
+ */
+export class LifecycleManager {
+  private timer: NodeJS.Timeout | null = null;
+
+  constructor(
+    private manager: ServerManager,
+    private getDefs: () => Record<string, ServerDef>,
+    private getGlobalIdleMinutes: () => number,
+    private intervalMs: number = TICK_INTERVAL_MS,
+  ) {}
+
+  /**
+   * Replace the managed ServerManager. Used by the index test-seam reset so
+   * the lifecycle stays bound to the module-level manager after a swap.
+   */
+  setManager(manager: ServerManager): void {
+    this.manager = manager;
+  }
+
+  start(): void {
+    if (this.timer) return;
+    this.timer = setInterval(() => this.tick(), this.intervalMs);
+    this.timer.unref?.(); // don't keep the process alive
+  }
+
+  stop(): void {
+    if (this.timer) {
+      clearInterval(this.timer);
+      this.timer = null;
+    }
+  }
+
+  /**
+   * One pass over all clients. Public so tests can drive a single tick
+   * deterministically without waiting for (or faking) the interval.
+   */
+  tick(): void {
+    const defs = this.getDefs();
+    for (const client of this.manager.getClients()) {
+      const def = defs[client.name];
+      if (!def) continue;
+      const lifecycle = def.lifecycle ?? "lazy";
+      if (lifecycle === "keep-alive" || lifecycle === "lazy-keep-alive") {
+        // Reconnect servers that fell out — including ones stuck in "error"
+        // after a failed connect (needs-auth/connecting/connected are left
+        // alone: 401 retries loop without OAuth, in-flight connects race).
+        if (client.status === "disconnected" || client.status === "error") {
+          // ADR 0004: deliberately NOT recorded — a 30s-tick reconnect is not
+          // a settle point; the live client wins within the session, and
+          // writing the ledger on every tick would only churn the cache file.
+          void client.connect().catch(() => {});
+        }
+      } else {
+        const idleMin = def.idleTimeout ?? this.getGlobalIdleMinutes();
+        if (idleMin > 0 && this.manager.isIdle(client.name, idleMin * 60_000)) {
+          void client.close();
+        }
+      }
+    }
+  }
+}