@sammorrowdrums/mcpi-ext 0.1.0

package/src/index.test.ts

@@ -0,0 +1,8 @@
+import { describe, it, expect } from "vitest";
+
+describe("mcpi-ext", () => {
+  it("exports a default function", async () => {
+    const mod = await import("./index.js");
+    expect(typeof mod.default).toBe("function");
+  });
+});

package/src/index.ts

@@ -0,0 +1,158 @@
+import type {
+  BeforeAgentStartEvent,
+  ExtensionAPI,
+  ExtensionContext,
+  SessionStartEvent,
+  ToolCallEvent,
+} from "@sammorrowdrums/mcpi";
+import { CodeModeManager } from "./code-mode/index.js";
+import { dockerE2ETool } from "./docker-e2e.js";
+import { McpClientManager, loadMcpConfig } from "./mcp/index.js";
+import {
+  SkillRegistry,
+  createLoadSkillTool,
+  discoverSkillsFromServer,
+  formatMcpSkillsForPrompt,
+  registerMcpToolProxies,
+} from "./skills/index.js";
+import { ToolCliRpcServer, formatToolCliForPrompt } from "./tool-cli/index.js";
+
+export default function (pi: ExtensionAPI) {
+  pi.registerTool(dockerE2ETool);
+
+  pi.registerFlag("mcp-config", {
+    description: "Path to MCP server configuration JSON file",
+    type: "string",
+  });
+
+  const mcpManager = new McpClientManager();
+  const skillRegistry = new SkillRegistry();
+  const rpcServer = new ToolCliRpcServer(mcpManager);
+  const codeModeManager = new CodeModeManager();
+  const enabledTools = new Set<string>();
+  const gatedToolNames = new Set<string>();
+
+  // Register the load_skill tool so the model can activate MCP skills
+  pi.registerTool(createLoadSkillTool({ registry: skillRegistry, mcpManager, enabledTools }));
+
+  pi.on("session_start", async (_event: SessionStartEvent, ctx: ExtensionContext) => {
+    if (ctx.hasUI) {
+      ctx.ui.notify("mcpi-ext loaded", "info");
+    }
+
+    const configPath = pi.getFlag("mcp-config") as string | undefined;
+    const log = (msg: string) => {
+      if (ctx.hasUI) {
+        ctx.ui.notify(msg, "info");
+      } else {
+        console.error(msg);
+      }
+    };
+    try {
+      const config = await loadMcpConfig(configPath);
+      const serverCount = Object.keys(config.mcpServers).length;
+      if (serverCount > 0) {
+        await mcpManager.connectAll(config, log);
+        const tools = mcpManager.getTools();
+        log(
+          `MCP: ${mcpManager.getConnectedServers().length} server(s), ${tools.length} tool(s) discovered`,
+        );
+
+        // Pre-register all MCP tools as deferred Pi tool proxies
+        // (in tools array for dispatch but excluded from system prompt)
+        const allToolNames = tools.map((t) => t.name);
+        registerMcpToolProxies(allToolNames, mcpManager, pi);
+
+        // Discover skills from all connected servers
+        for (const serverName of mcpManager.getConnectedServers()) {
+          const client = mcpManager.getClient(serverName);
+          if (!client) continue;
+          try {
+            const skills = await discoverSkillsFromServer(client, serverName, log);
+            skillRegistry.registerAll(skills);
+          } catch (err) {
+            log(
+              `[skills] Failed to discover skills from "${serverName}": ${(err as Error).message}`,
+            );
+          }
+        }
+
+        if (skillRegistry.size > 0) {
+          for (const skill of skillRegistry.getAll()) {
+            for (const t of skill.allowedTools) {
+              gatedToolNames.add(t);
+            }
+          }
+          log(
+            `MCP: ${skillRegistry.size} skill(s) discovered, ${gatedToolNames.size} tool(s) deferred`,
+          );
+        }
+
+        // Start the tool-cli RPC server for progressive tool discovery
+        try {
+          await rpcServer.start(log);
+        } catch (err) {
+          log(`[tool-cli] Failed to start RPC server: ${(err as Error).message}`);
+        }
+
+        // Initialize code mode (Tier 3) for read-only tools with structured output
+        codeModeManager.initialize(mcpManager);
+        if (codeModeManager.isActive) {
+          const { codeSearch, codeExecute } = codeModeManager.createTools();
+          pi.registerTool(codeSearch);
+          pi.registerTool(codeExecute);
+          log(
+            `MCP: Code mode active (${codeModeManager.getEligibleTools().length} eligible tool(s))`,
+          );
+        }
+      }
+    } catch (err) {
+      const msg = `MCP config error: ${(err as Error).message}`;
+      if (ctx.hasUI) {
+        ctx.ui.notify(msg, "warning");
+      } else {
+        console.error(msg);
+      }
+    }
+  });
+
+  // Inject MCP skills, tool-cli advice, and code mode type hints into the system prompt
+  pi.on("before_agent_start", async (event: BeforeAgentStartEvent) => {
+    let extra = "";
+
+    const skills = skillRegistry.getAll();
+    if (skills.length > 0) {
+      extra += formatMcpSkillsForPrompt(skills);
+    }
+
+    const serverCount = mcpManager.getConnectedServers().length;
+    extra += formatToolCliForPrompt(serverCount);
+
+    if (codeModeManager.isActive) {
+      extra += codeModeManager.formatSystemPromptSection();
+    }
+
+    if (extra.length === 0) return;
+    return { systemPrompt: event.systemPrompt + extra };
+  });
+
+  // Block deferred MCP tools until their skill is loaded
+  pi.on("tool_call", async (event: ToolCallEvent) => {
+    const name = "toolName" in event ? event.toolName : undefined;
+    if (!name || !gatedToolNames.has(name) || enabledTools.has(name)) return;
+    const relevantSkills = skillRegistry
+      .getAll()
+      .filter((s) => s.allowedTools.includes(name))
+      .map((s) => s.name);
+    return {
+      block: true,
+      reason: `Tool "${name}" requires loading a skill first. Call load_skill with one of: ${relevantSkills.join(", ")}`,
+    };
+  });
+
+  pi.on("session_shutdown", async () => {
+    await rpcServer.stop();
+    await mcpManager.disconnectAll();
+    skillRegistry.clear();
+  });
+}

package/src/mcp/client-manager.test.ts

@@ -0,0 +1,250 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { McpClientManager } from "./client-manager.js";
+
+// Mock the MCP SDK transports and client
+vi.mock("@modelcontextprotocol/sdk/client/index.js", () => {
+  class MockClient {
+    private _onToolsChanged?: (error: Error | null, tools: unknown[] | null) => void;
+
+    constructor(
+      _info: unknown,
+      options?: {
+        listChanged?: {
+          tools?: { onChanged?: (error: Error | null, tools: unknown[] | null) => void };
+        };
+      },
+    ) {
+      this._onToolsChanged = options?.listChanged?.tools?.onChanged;
+    }
+    async connect(_transport: unknown): Promise<void> {
+      // no-op for tests
+    }
+    async close(): Promise<void> {
+      // no-op for tests
+    }
+    async listTools(): Promise<{ tools: unknown[] }> {
+      return {
+        tools: [
+          {
+            name: "mock_tool",
+            description: "A mock tool for testing",
+            inputSchema: { type: "object", properties: { arg: { type: "string" } } },
+            annotations: { readOnlyHint: true },
+          },
+          {
+            name: "mock_tool_2",
+            description: "Another mock tool",
+            inputSchema: { type: "object", properties: {} },
+          },
+        ],
+      };
+    }
+
+    /** Test helper to simulate tools/list_changed notification. */
+    _simulateToolsChanged(tools: unknown[]): void {
+      this._onToolsChanged?.(null, tools);
+    }
+  }
+  return { Client: MockClient };
+});
+
+vi.mock("@modelcontextprotocol/sdk/client/stdio.js", () => ({
+  // eslint-disable-next-line @typescript-eslint/no-extraneous-class
+  StdioClientTransport: class {
+    // eslint-disable-next-line @typescript-eslint/no-useless-constructor
+    constructor(_params: unknown) {
+      // no-op mock
+    }
+  },
+}));
+
+vi.mock("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
+  // eslint-disable-next-line @typescript-eslint/no-extraneous-class
+  StreamableHTTPClientTransport: class {
+    // eslint-disable-next-line @typescript-eslint/no-useless-constructor
+    constructor(_url: URL, _opts?: unknown) {
+      // no-op mock
+    }
+  },
+}));
+
+describe("McpClientManager", () => {
+  let manager: McpClientManager;
+
+  beforeEach(() => {
+    manager = new McpClientManager();
+  });
+
+  it("connects to a stdio server and discovers tools", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        "test-stdio": {
+          type: "stdio",
+          command: "node",
+          args: ["test-server.js"],
+        },
+      },
+    });
+
+    expect(manager.getConnectedServers()).toEqual(["test-stdio"]);
+
+    const tools = manager.getTools();
+    expect(tools).toHaveLength(2);
+    expect(tools[0].name).toBe("mock_tool");
+    expect(tools[0].serverName).toBe("test-stdio");
+    expect(tools[0].description).toBe("A mock tool for testing");
+    expect(tools[0].annotations).toEqual({ readOnlyHint: true });
+    expect(tools[1].name).toBe("mock_tool_2");
+  });
+
+  it("connects to a remote server and discovers tools", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        "test-remote": {
+          type: "remote",
+          url: "https://example.com/mcp",
+          headers: { Authorization: "Bearer test" },
+        },
+      },
+    });
+
+    expect(manager.getConnectedServers()).toEqual(["test-remote"]);
+    expect(manager.getTools()).toHaveLength(2);
+  });
+
+  it("connects to multiple servers and aggregates tools", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        server1: { type: "stdio", command: "node", args: ["s1.js"] },
+        server2: { type: "remote", url: "https://example.com/mcp" },
+      },
+    });
+
+    expect(manager.getConnectedServers()).toHaveLength(2);
+    // 2 tools per server
+    expect(manager.getTools()).toHaveLength(4);
+  });
+
+  it("getToolsForServer returns tools for a specific server", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        srv: { type: "stdio", command: "echo" },
+      },
+    });
+
+    expect(manager.getToolsForServer("srv")).toHaveLength(2);
+    expect(manager.getToolsForServer("nonexistent")).toEqual([]);
+  });
+
+  it("getClient returns the client for a connected server", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        srv: { type: "stdio", command: "echo" },
+      },
+    });
+
+    expect(manager.getClient("srv")).toBeDefined();
+    expect(manager.getClient("nonexistent")).toBeUndefined();
+  });
+
+  it("disconnectOne removes a server", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        srv: { type: "stdio", command: "echo" },
+      },
+    });
+
+    expect(manager.getConnectedServers()).toEqual(["srv"]);
+    await manager.disconnectOne("srv");
+    expect(manager.getConnectedServers()).toEqual([]);
+    expect(manager.getTools()).toEqual([]);
+  });
+
+  it("disconnectAll removes all servers", async () => {
+    await manager.connectAll({
+      mcpServers: {
+        s1: { type: "stdio", command: "echo" },
+        s2: { type: "remote", url: "https://example.com/mcp" },
+      },
+    });
+
+    expect(manager.getConnectedServers()).toHaveLength(2);
+    await manager.disconnectAll();
+    expect(manager.getConnectedServers()).toEqual([]);
+  });
+
+  it("handles empty config gracefully", async () => {
+    await manager.connectAll({ mcpServers: {} });
+    expect(manager.getConnectedServers()).toEqual([]);
+    expect(manager.getTools()).toEqual([]);
+  });
+
+  it("reconnects when connecting to an already-connected server name", async () => {
+    await manager.connectOne("srv", { type: "stdio", command: "node" });
+    expect(manager.getConnectedServers()).toEqual(["srv"]);
+
+    // Connect again with same name — should replace
+    await manager.connectOne("srv", { type: "remote", url: "https://example.com/mcp" });
+    expect(manager.getConnectedServers()).toEqual(["srv"]);
+    expect(manager.getTools()).toHaveLength(2);
+  });
+
+  it("logs failures but continues connecting other servers", async () => {
+    const logs: string[] = [];
+    // Patch Client.connect to fail for a specific transport
+    const { Client } = await import("@modelcontextprotocol/sdk/client/index.js");
+    const origConnect = Client.prototype.connect;
+    let callCount = 0;
+    Client.prototype.connect = async function (transport: unknown) {
+      callCount++;
+      if (callCount === 1) throw new Error("Connection refused");
+      return origConnect.call(
+        this,
+        transport as import("@modelcontextprotocol/sdk/shared/transport.js").Transport,
+      );
+    };
+
+    try {
+      await manager.connectAll(
+        {
+          mcpServers: {
+            failing: { type: "stdio", command: "bad-command" },
+            working: { type: "stdio", command: "good-command" },
+          },
+        },
+        (msg) => logs.push(msg),
+      );
+
+      // The working server should still be connected
+      expect(manager.getConnectedServers()).toContain("working");
+      expect(logs.some((l) => l.includes("Failed to connect") && l.includes("failing"))).toBe(true);
+    } finally {
+      Client.prototype.connect = origConnect;
+    }
+  });
+
+  it("updates tools when list_changed notification fires", async () => {
+    const logs: string[] = [];
+    await manager.connectAll({ mcpServers: { srv: { type: "stdio", command: "echo" } } }, (msg) =>
+      logs.push(msg),
+    );
+
+    expect(manager.getToolsForServer("srv")).toHaveLength(2);
+
+    // Simulate a tools/list_changed notification via the mock
+    const client = manager.getClient("srv") as unknown as {
+      _simulateToolsChanged: (tools: unknown[]) => void;
+    };
+    client._simulateToolsChanged([
+      {
+        name: "new_tool",
+        description: "A dynamically added tool",
+        inputSchema: { type: "object", properties: {} },
+      },
+    ]);
+
+    expect(manager.getToolsForServer("srv")).toHaveLength(1);
+    expect(manager.getToolsForServer("srv")[0].name).toBe("new_tool");
+    expect(logs.some((l) => l.includes("Tools updated"))).toBe(true);
+  });
+});

package/src/mcp/client-manager.ts

@@ -0,0 +1,199 @@
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import type { McpConfig, ServerConfig } from "./config.js";
+
+/** An MCP tool as discovered from a server. */
+export interface McpTool {
+  /** Tool name as reported by the server. */
+  name: string;
+  description?: string;
+  /** JSON Schema for the tool's parameters. */
+  inputSchema: Record<string, unknown>;
+  /** JSON Schema for structured output, if declared. */
+  outputSchema?: Record<string, unknown>;
+  /** Tool annotations (readOnlyHint, destructiveHint, etc.). */
+  annotations?: Record<string, unknown>;
+  /** Which configured server this tool came from. */
+  serverName: string;
+}
+
+interface ManagedConnection {
+  client: Client;
+  transport: StdioClientTransport | StreamableHTTPClientTransport;
+  tools: McpTool[];
+}
+
+/**
+ * Manages connections to multiple MCP servers and aggregates their tools.
+ *
+ * Tools are discovered and stored internally. They are NOT registered with pi
+ * automatically — that responsibility belongs to the access tiers (Skills,
+ * Football, Code Mode).
+ */
+export class McpClientManager {
+  private connections = new Map<string, ManagedConnection>();
+
+  /**
+   * Connect to all servers defined in the config.
+   * Connections that fail are logged and skipped — partial success is fine.
+   */
+  async connectAll(config: McpConfig, log: (msg: string) => void = console.error): Promise<void> {
+    const entries = Object.entries(config.mcpServers);
+    const results = await Promise.allSettled(
+      entries.map(([name, serverConfig]) => this.connectOne(name, serverConfig, log)),
+    );
+
+    for (let i = 0; i < results.length; i++) {
+      const result = results[i];
+      if (result.status === "rejected") {
+        log(`[mcp] Failed to connect to "${entries[i][0]}": ${result.reason}`);
+      }
+    }
+  }
+
+  /** Connect to a single server by name. */
+  async connectOne(
+    name: string,
+    serverConfig: ServerConfig,
+    log: (msg: string) => void = console.error,
+  ): Promise<void> {
+    // Disconnect existing connection with this name if any
+    if (this.connections.has(name)) {
+      await this.disconnectOne(name);
+    }
+
+    const transport = createTransport(serverConfig);
+
+    const client = new Client(
+      { name: "mcpi-ext", version: "0.1.0" },
+      {
+        capabilities: {},
+        listChanged: {
+          tools: {
+            onChanged: (_error, tools) => {
+              if (tools) {
+                const conn = this.connections.get(name);
+                if (conn && conn.client === client) {
+                  conn.tools = toMcpTools(name, tools);
+                  log(`[mcp] Tools updated for "${name}" (${conn.tools.length} tools)`);
+                }
+              }
+            },
+          },
+        },
+      },
+    );
+
+    // Insert placeholder so listChanged notifications during connect aren't dropped
+    this.connections.set(name, { client, transport, tools: [] });
+
+    try {
+      await client.connect(transport);
+
+      const toolsResult = await client.listTools();
+      const tools = toMcpTools(name, toolsResult.tools);
+      const conn = this.connections.get(name);
+      if (conn && conn.client === client) {
+        conn.tools = tools;
+      }
+      log(`[mcp] Connected to "${name}" (${tools.length} tools)`);
+    } catch (error) {
+      // Clean up on failure so we don't leak a partially connected client
+      const conn = this.connections.get(name);
+      if (conn && conn.client === client) {
+        this.connections.delete(name);
+      }
+      try {
+        await client.close();
+      } catch {
+        // best-effort cleanup
+      }
+      throw error;
+    }
+  }
+
+  /** Disconnect a single server. */
+  async disconnectOne(name: string): Promise<void> {
+    const conn = this.connections.get(name);
+    if (!conn) return;
+    this.connections.delete(name);
+    try {
+      await conn.client.close();
+    } catch {
+      // best-effort
+    }
+  }
+
+  /** Disconnect all servers. */
+  async disconnectAll(): Promise<void> {
+    const names = [...this.connections.keys()];
+    await Promise.allSettled(names.map((n) => this.disconnectOne(n)));
+  }
+
+  /** Get all discovered tools across all connected servers. */
+  getTools(): McpTool[] {
+    const tools: McpTool[] = [];
+    for (const conn of this.connections.values()) {
+      tools.push(...conn.tools);
+    }
+    return tools;
+  }
+
+  /** Get tools from a specific server. */
+  getToolsForServer(name: string): McpTool[] {
+    return this.connections.get(name)?.tools ?? [];
+  }
+
+  /** Get the MCP Client for a specific server (needed to call tools). */
+  getClient(name: string): Client | undefined {
+    return this.connections.get(name)?.client;
+  }
+
+  /** List connected server names. */
+  getConnectedServers(): string[] {
+    return [...this.connections.keys()];
+  }
+}
+
+function createTransport(
+  config: ServerConfig,
+): StdioClientTransport | StreamableHTTPClientTransport {
+  if (config.type === "stdio") {
+    const env = config.env
+      ? Object.fromEntries(
+          Object.entries({ ...process.env, ...config.env }).filter(
+            (entry): entry is [string, string] => entry[1] !== undefined,
+          ),
+        )
+      : undefined;
+    return new StdioClientTransport({
+      command: config.command,
+      args: config.args,
+      env,
+      cwd: config.cwd,
+    });
+  }
+
+  const headers: Record<string, string> = {
+    ...(config.headers ?? {}),
+  };
+
+  return new StreamableHTTPClientTransport(new URL(config.url), {
+    requestInit: {
+      headers,
+    },
+  });
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function toMcpTools(serverName: string, tools: any[]): McpTool[] {
+  return tools.map((t) => ({
+    name: t.name,
+    description: t.description,
+    inputSchema: t.inputSchema as Record<string, unknown>,
+    outputSchema: t.outputSchema as Record<string, unknown> | undefined,
+    annotations: t.annotations as Record<string, unknown> | undefined,
+    serverName,
+  }));
+}