mcpctl/src/pi-ext/mcpctl-pi.ts

/**
 * mcpctl pi extension — native integration (no MCP client, no Claude).
 *
 * Talks JSON-RPC directly to mcplocal's project MCP endpoint and registers
 * each project tool as a native pi tool.
 *
 * KEY ARCHITECTURAL POINT: mcplocal's project *gate* state is bound to a
 * `mcp-session-id`. A gated project only exposes `begin_session` until it is
 * called, after which the full tool set becomes callable — but only within
 * THAT session. So we keep ONE persistent `McpHttpSession` per active project
 * and route every tools/call through it, re-listing after each call to pick up
 * the post-begin_session tool additions.
 *
 * Features:
 *   - project gating auto-ungates when begin_session is called
 *   - switching the active project from the pi TUI (/mcpctl)
 *   - a status line
 *
 * Load standalone:  pi -e ./src/pi-ext/mcpctl-pi.ts
 * or via settings:  "extensions": ["/abs/path/to/mcpctl-pi.ts"]
 *
 * Only imports pi-bundled packages — no @mcpctl/*, no ~/.claude.
 *
 * RUNTIME IMPORTS ARE LOAD-BEARING: pi resolves an extension's bare specifiers
 * through a fixed alias table in its own loader, and that table differs between
 * pi distributions — `@earendil-works/*` exists only in the newer packages,
 * while `@mariozechner/*` installs alias only the old names. `typebox` is the
 * one specifier every published pi aliases, so it is the ONLY runtime import
 * allowed here. Anything else must be `import type` (erased before jiti runs)
 * or inlined — see `stringEnum` below.
 */
import { Type, type TSchema } from "typebox";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
  McpHttpSession,
  type ToolInfo,
  type McpHttpOptions,
} from "./mcp-http.js";
import { readFile, writeFile, mkdir } from "node:fs/promises";
import { join, dirname } from "node:path";
import { homedir } from "node:os";

// ── config / state ───────────────────────────────────────────────────────────
const McpctlDir = join(homedir(), ".mcpctl");
const ConfigPath = join(McpctlDir, "config.json");
const CredPath = join(McpctlDir, "credentials");
const PiStatePath = join(McpctlDir, "pi-state.json");
const ProjectMarker = ".mcpctl-project";

interface PiState {
  project?: string;
}

async function readJson<T>(path: string): Promise<T | null> {
  try {
    return JSON.parse(await readFile(path, "utf-8")) as T;
  } catch {
    return null;
  }
}

async function readConfig(): Promise<{ mcplocalUrl: string }> {
  const cfg = await readJson<{ mcplocalUrl?: string }>(ConfigPath);
  return { mcplocalUrl: cfg?.mcplocalUrl ?? "http://localhost:3200" };
}

async function readCreds(): Promise<{ token?: string; mcpdUrl?: string } | null> {
  try {
    const raw = await readFile(CredPath, "utf-8");
    const parsed = JSON.parse(raw) as { token?: string; mcpdUrl?: string };
    if (!parsed.token) return null;
    return parsed;
  } catch {
    return null;
  }
}

async function readPiState(): Promise<PiState> {
  return (await readJson<PiState>(PiStatePath)) ?? {};
}

async function writePiState(state: PiState): Promise<void> {
  await mkdir(McpctlDir, { recursive: true });
  await writeFile(PiStatePath, JSON.stringify(state, null, 2) + "\n", "utf-8");
}

/** Walk up directories looking for a `.mcpctl-project` marker. */
async function findProjectMarker(start = process.cwd()): Promise<string | null> {
  const home = homedir();
  let dir = start;
  for (let i = 0; i < 50; i++) {
    try {
      const raw = await readFile(join(dir, ProjectMarker), "utf-8");
      const project = raw.split("\n")[0]?.trim() ?? "";
      if (project) return project;
    } catch {
      // continue walking
    }
    if (dir === home || dir === "/") break;
    const parent = dirname(dir);
    if (parent === dir) break;
    dir = parent;
  }
  return null;
}

/** Request the list of projects from mcpd (via mcplocal API passthrough). */
async function listProjects(mcplocalUrl: string, token?: string): Promise<string[]> {
  const url = `${mcplocalUrl}/api/v1/projects`;
  const headers: Record<string, string> = { Accept: "application/json" };
  if (token) headers.Authorization = `Bearer ${token}`;
  const res = await fetch(url, { headers });
  if (!res.ok) throw new Error(`list projects: HTTP ${res.status}`);
  const body = (await res.json()) as Array<{ name: string }> | { data?: Array<{ name: string }> };
  const arr = Array.isArray(body) ? body : (body.data ?? []);
  return arr.map((p) => p.name);
}

// ── JSON Schema → TypeBox ────────────────────────────────────────────────────

/**
 * `{ type: "string", enum: [...] }` rather than a union of literals: Google's
 * API (and other providers that reject anyOf/const) only accept the flat form.
 *
 * Inlined from pi-ai's `StringEnum` on purpose — importing it dragged in
 * `@earendil-works/pi-ai`, which older pi installs cannot resolve, and the
 * whole extension then failed to load. See the import note at the top.
 */
function stringEnum(values: string[], description?: string): TSchema {
  return Type.Unsafe<string>({
    type: "string",
    enum: values,
    ...(description ? { description } : {}),
  });
}

function convertSchema(inputSchema: unknown): TSchema {
  if (!inputSchema || typeof inputSchema !== "object") {
    return Type.Object({});
  }
  const s = inputSchema as {
    type?: string;
    properties?: Record<string, unknown>;
    required?: string[];
    items?: unknown;
    enum?: unknown[];
  };

  const required = new Set(Array.isArray(s.required) ? s.required : []);
  const props = s.properties ?? {};
  const fields: Record<string, TSchema> = {};

  for (const [name, raw] of Object.entries(props)) {
    const schema = convertProp(raw);
    fields[name] = required.has(name) ? schema : Type.Optional(schema);
  }
  return Type.Object(fields);
}

function convertProp(raw: unknown): TSchema {
  if (!raw || typeof raw !== "object") return Type.Unknown();
  const s = raw as {
    type?: string | string[];
    enum?: unknown[];
    description?: string;
    items?: unknown;
    properties?: Record<string, unknown>;
  };

  const desc = typeof s.description === "string" ? s.description : undefined;
  const enumVals = Array.isArray(s.enum) && s.enum.length > 0 ? s.enum : undefined;

  if (enumVals && enumVals.every((v) => typeof v === "string")) {
    return stringEnum(enumVals as string[], desc);
  }
  if (enumVals && enumVals.every((v) => typeof v === "number")) {
    const literals = enumVals.map((v) => Type.Literal(v));
    return desc ? Type.Union(literals, { description: desc }) : Type.Union(literals);
  }

  let type = s.type;
  if (Array.isArray(type)) type = type.find((t) => t !== "null");
  type = type ?? (s.properties ? "object" : s.items ? "array" : "string");
  if (type === "integer") type = "number";

  switch (type) {
    case "string":
      return desc ? Type.String({ description: desc }) : Type.String();
    case "number":
      return desc ? Type.Number({ description: desc }) : Type.Number();
    case "boolean":
      return desc ? Type.Boolean({ description: desc }) : Type.Boolean();
    case "array": {
      const inner = s.items ? convertProp(s.items) : Type.Unknown();
      return desc ? Type.Array(inner, { description: desc }) : Type.Array(inner);
    }
    case "object":
      return convertSchema(s as unknown);
    default:
      return desc ? Type.Unknown({ description: desc }) : Type.Unknown();
  }
}

/**
 * Order and filter the project list for the picker.
 *
 * Ordering: the active project first (most likely pick), then alphabetical.
 * Filtering: space-separated terms, ALL of which must appear as
 * case-insensitive substrings — so `home auto` finds `homeautomation`. A blank
 * query keeps everything.
 *
 * Exported so the behaviour is unit-tested rather than eyeballed through a TUI.
 */
export function filterProjects(projects: string[], query: string, active: string | null): string[] {
  const ordered = [...projects].sort((a, b) => {
    if (a === active) return -1;
    if (b === active) return 1;
    return a.localeCompare(b);
  });
  const terms = query.toLowerCase().split(/\s+/).filter((t) => t.length > 0);
  if (terms.length === 0) return ordered;
  return ordered.filter((p) => {
    const name = p.toLowerCase();
    return terms.every((t) => name.includes(t));
  });
}

/**
 * The message injected into the conversation after a project switch, telling
 * the model which mcpctl tools are live now.
 *
 * Exported so its wording is unit-tested; it is the only thing standing between
 * the model and a stale tool list it will otherwise keep calling.
 */
export function toolChangeAnnouncement(project: string, tools: string[]): string {
  const gate = tools.find((n) => n.endsWith("_begin_session"));
  const parts = [
    `[mcpctl] The active project is now '${project}'.`,
    "mcpctl tool names listed earlier in this conversation belong to the previously active project and are no longer callable — ignore them.",
    tools.length > 0
      ? `Currently available mcpctl tools: ${tools.join(", ")}.`
      : "No mcpctl tools are currently available for this project.",
  ];
  if (gate !== undefined) {
    parts.push(`This project is gated: call ${gate} first and its remaining tools become available.`);
  }
  return parts.join(" ");
}

/** Sanitize a name for use as a pi tool name segment ([a-z0-9_]). */
function safeSegment(name: string): string {
  return name.toLowerCase().replace(/[^a-z0-9_]+/g, "_").replace(/^_+|_+$/g, "") || "x";
}

function slugify(name: string): string {
  return safeSegment(name);
}

/** pi tool name for a project + MCP tool (e.g. `mc_docmost_docmost_create_page`). */
function piToolName(project: string, mcpTool: string): string {
  const proj = slugify(project);
  const tool = mcpTool.replace(/^[a-z0-9_-]+\//i, (m) => slugify(m.slice(0, -1)) + "_") || mcpTool;
  return `mc_${proj}_${safeSegment(tool)}`;
}

// ── the extension ────────────────────────────────────────────────────────────
export default function (pi: ExtensionAPI) {
  let mcplocalUrl = "http://localhost:3200";
  let token: string | undefined;
  let activeProject: string | null = null;
  /** Persistent session for the active project — MUST be reused to keep gate state. */
  let session: McpHttpSession | null = null;
  /** Registered pi tool name → the MCP tool it forwards to, for the active project. */
  const projectToolNames = new Map<string, string>();
  let activeCtx: ExtensionContext | null = null;

  function endpointUrl(): string {
    return `${mcplocalUrl.replace(/\/$/, "")}/projects/${encodeURIComponent(activeProject ?? "")}/mcp`;
  }

  /** Build (or reuse) the persistent session for the active project. */
  async function ensureSession(): Promise<McpHttpSession> {
    if (session) return session;
    const opts: McpHttpOptions = { bearer: token, timeoutMs: 120_000 };
    const sess = new McpHttpSession(endpointUrl(), opts);
    await sess.initialize();
    await sess.sendNotification("notifications/initialized");
    session = sess;
    return sess;
  }

  async function closeSession(): Promise<void> {
    const s = session;
    session = null;
    if (s) await s.close();
  }

  /**
   * (Re)discover the active project's current tool set and register any new pi
   * tools. Called on load and after every tool call (to pick up ungating).
   */
  async function reconcileTools(project: string): Promise<{ tools: string[] }> {
    const sess = await ensureSession();
    const { tools } = await sess.listTools();
    const names: string[] = [];
    for (const t of tools) {
      const piName = uniqueToolName(project, t.name);
      if (piName === null) continue; // already registered for this same MCP tool
      pi.registerTool(createTool(piName, t, project));
      projectToolNames.set(piName, t.name);
      names.push(piName);
    }
    // Names already registered for tools still in the list stay active.
    for (const [piName, mcpName] of projectToolNames) {
      if (!names.includes(piName) && tools.some((t) => t.name === mcpName)) names.push(piName);
    }
    // Activate built-ins/other extensions + all tools discovered in this project.
    const active = pi.getActiveTools().filter((n) => !n.startsWith("mc_"));
    pi.setActiveTools([...new Set([...active, ...names])]);
    return { tools: names };
  }

  /**
   * The pi tool name to register `mcpTool` under, or null if it is already
   * registered. Sanitising MCP names down to `[a-z0-9_]` can collide (e.g.
   * `docs.search` and `docs-search`); without a suffix the second tool would
   * silently never be registered while still being reported as available, and
   * its calls would be forwarded to the first tool instead.
   */
  function uniqueToolName(project: string, mcpTool: string): string | null {
    const base = piToolName(project, mcpTool);
    if (projectToolNames.get(base) === mcpTool) return null;
    if (!projectToolNames.has(base)) return base;
    for (let i = 2; i < 100; i++) {
      const candidate = `${base}_${String(i)}`;
      if (projectToolNames.get(candidate) === mcpTool) return null;
      if (!projectToolNames.has(candidate)) return candidate;
    }
    return null;
  }

  function createTool(piName: string, t: ToolInfo, project: string) {
    const mcpName = t.name;
    return {
      name: piName,
      label: `${project} :: ${mcpName}`,
      description: t.description ?? `mcpctl tool ${mcpName} on project ${project}`,
      parameters: convertSchema(t.inputSchema),
      async execute(_toolCallId: string, params: Record<string, unknown>) {
        const sess = await ensureSession();
        const res = await sess.callTool(mcpName, params ?? {});
        // Re-list after every call — picks up post-begin_session tool additions.
        try {
          await reconcileTools(project);
        } catch {
          // non-fatal
        }
        const text = (res.content ?? [])
          .filter((c) => c.type === "text")
          .map((c) => c.text ?? "")
          .join("\n");
        if (res.isError) {
          throw new Error(text || `${mcpName} returned an error`);
        }
        return {
          // `type` must be the literal "text" — pi's AgentToolResult content is
          // a TextContent | ImageContent union, not { type: string }.
          content: [{ type: "text" as const, text: text || "(no text content)" }],
          details: { mcpTool: mcpName, project, raw: res },
        };
      },
    };
  }

  function setStatus(msg?: string): void {
    try {
      activeCtx?.ui?.setStatus?.("mcpctl", msg ?? (activeProject ? `project: ${activeProject}` : "no project"));
    } catch {
      /* non-fatal */
    }
  }

  // ── /mcpctl command ──
  /**
   * Menu labels → actions. `ctx.ui.select` takes plain strings and returns the
   * chosen string, so the mapping has to live here rather than in the options.
   */
  const MENU: ReadonlyArray<readonly [label: string, action: string]> = [
    ["Status", "status"],
    ["Switch project", "switch"],
    ["Refresh tools", "refresh"],
    ["Sync skills into pi's skills dir", "sync"],
    ["Close", "close"],
  ];

  pi.registerCommand("mcpctl", {
    description: "mcpctl integration: status, switch project, refresh tools",
    handler: async (_args, ctx) => {
      if (!ctx.hasUI) {
        // A command handler resolves to void — pi ignores a returned string, so
        // headless status has to go through notify.
        ctx.ui.notify(
          `mcpctl: active project = ${activeProject ?? "none"}; mcplocal = ${mcplocalUrl}`,
          "info",
        );
        return;
      }
      const picked = await ctx.ui.select("mcpctl", MENU.map(([label]) => label));
      if (!picked) return;
      const choice = MENU.find(([label]) => label === picked)?.[1];
      if (!choice || choice === "close") return;

      if (choice === "status") {
        const cs = await readConfig();
        const creds = await readCreds();
        const info =
          `mcplocal: ${cs.mcplocalUrl}\n` +
          `user: ${creds?.mcpdUrl ?? "not logged in"}\n` +
          `active project: ${activeProject ?? "none"}\n` +
          `tools registered: ${projectToolNames.size}\n` +
          `  ${[...projectToolNames].join("\n  ")}\n`;
        await ctx.ui.editor("mcpctl status", info);
        return;
      }

      if (choice === "switch") {
        await switchProject(ctx);
        return;
      }

      if (choice === "refresh") {
        if (!activeProject) {
          ctx.ui.notify("No active project — switch project first", "warning");
          return;
        }
        try {
          const r = await reconcileTools(activeProject);
          ctx.ui.notify(`Refreshed: ${r.tools.length} tools for '${activeProject}'`, "info");
          setStatus();
        } catch (e) {
          ctx.ui.notify(`Refresh failed: ${(e as Error).message}`, "error");
        }
        return;
      }

      if (choice === "sync") {
        try {
          // `--agent pi` is what keeps skills in ~/.pi/agent/skills; without it
          // this syncs into ~/.claude/skills, which the whole point of this
          // integration is to avoid depending on.
          const args = ["skills", "sync", "--agent", "pi", "--quiet"];
          if (activeProject) args.push("-p", activeProject);
          // execFile, not execSync+stdio:"inherit": a shell would need the
          // project name quoted, and inherited stdio writes raw output over
          // pi's TUI. Capture instead and report through notify.
          const { execFile } = await import("node:child_process");
          const stderr = await new Promise<string>((resolve, reject) => {
            execFile("mcpctl", args, { timeout: 120_000 }, (err, _stdout, errOut) => {
              if (err) reject(new Error((errOut || err.message).trim()));
              else resolve(errOut);
            });
          });
          const detail = stderr.trim().split("\n").pop();
          ctx.ui.notify(
            `mcpctl skills synced into pi's skills dir${detail ? ` — ${detail}` : ""}. Run /reload to pick them up.`,
            "info",
          );
        } catch (e) {
          ctx.ui.notify(`skills sync failed: ${(e as Error).message}`, "error");
        }
      }
    },
  });

  /** Above this many projects, arrowing through the list stops being usable. */
  const FILTER_THRESHOLD = 20;
  /**
   * Hard cap on rendered rows. pi's ExtensionSelectorComponent.updateList()
   * renders *every* option with no windowing, so an uncapped list scrolls the
   * whole screen away. prime-agent's selector windows to ~20 itself; this
   * matches that height so both hosts feel the same.
   */
  const MAX_SHOWN = 20;

  /**
   * Choose a project, asking for a filter first when the list is long.
   *
   * pi's selector is a plain arrow-key list — `ExtensionUIDialogOptions` has no
   * search and `ExtensionSelectorComponent` ignores typed characters — so the
   * filtering has to happen before the list is handed over. Real installs run
   * to hundreds of projects (smoke-test leftovers included), where scrolling is
   * hopeless.
   *
   * Terms are space-separated and ALL must match, case-insensitively, as
   * substrings: `home auto` finds `homeautomation`. Blank shows everything.
   */
  async function pickProject(ctx: ExtensionContext, projects: string[]): Promise<string | undefined> {
    let candidates = filterProjects(projects, "", activeProject);
    if (candidates.length > FILTER_THRESHOLD) {
      const query = await ctx.ui.input(
        `Filter ${String(candidates.length)} projects (blank = all, Esc = cancel)`,
        "e.g. home auto",
      );
      if (query === undefined) return undefined; // cancelled
      candidates = filterProjects(projects, query, activeProject);
      if (candidates.length === 0) {
        ctx.ui.notify(`No project matches '${query}'`, "warning");
        return undefined;
      }
    }

    const truncated = candidates.length > MAX_SHOWN;
    const shown = truncated ? candidates.slice(0, MAX_SHOWN) : candidates;
    const title = truncated
      // Say what was dropped: a silently capped list reads as "that's all of them".
      ? `Switch to project (${String(MAX_SHOWN)} of ${String(candidates.length)} matches — narrow the filter)`
      : `Switch to project (${String(shown.length)})`;
    return ctx.ui.select(title, shown);
  }

  async function switchProject(ctx: ExtensionContext): Promise<void> {
    let projects: string[] = [];
    try {
      projects = await listProjects(mcplocalUrl, token);
    } catch (e) {
      ctx.ui.notify(`Failed to list projects: ${(e as Error).message}`, "error");
      return;
    }
    if (projects.length === 0) {
      ctx.ui.notify("No projects returned by mcpd", "warning");
      return;
    }
    const picked = await pickProject(ctx, projects);
    if (!picked) return;

    // Tear down old project's session + active tools.
    await closeSession();
    const active = pi.getActiveTools().filter((n) => !n.startsWith("mc_"));
    pi.setActiveTools([...active]);
    projectToolNames.clear();

    activeProject = picked;
    await writePiState({ project: picked });
    try {
      const r = await reconcileTools(picked);
      const gate = r.tools.find((n) => n.endsWith("_begin_session"));
      ctx.ui.notify(
        `Switched to '${picked}': ${String(r.tools.length)} tool(s) ready${gate ? ` — gated, call ${gate} to unlock the rest` : ""}`,
        "info",
      );
      setStatus();
      announceToolChange(picked, r.tools);
    } catch (e) {
      ctx.ui.notify(`Could not load tools for '${picked}': ${(e as Error).message}`, "error");
    }
  }

  /**
   * Tell the *model* that the tool set changed.
   *
   * pi cannot unregister a tool — `registerTool` only ever adds to the
   * extension's tool map — so the previous project's `mc_*` tools stay
   * registered and merely go inactive. Meanwhile the conversation still
   * contains the old project's tool list, so the model keeps calling names that
   * now answer "Tool ... not found" and concludes it has lost MCP access
   * entirely, which is what happens in practice on the second switch.
   *
   * A custom message is converted to a user-role message and does reach the
   * LLM (unlike `appendEntry`, which is explicitly excluded from context), so
   * this corrects the stale context instead of leaving the model to guess.
   * `display: false` keeps it out of the transcript — the notify above is what
   * the human reads.
   */
  function announceToolChange(project: string, tools: string[]): void {
    try {
      pi.sendMessage(
        { customType: "mcpctl-project-switch", content: toolChangeAnnouncement(project, tools), display: false },
        { triggerTurn: false, deliverAs: "nextTurn" },
      );
    } catch {
      // Older pi without sendMessage — the notify above still informs the user.
    }
  }

  // ── session lifecycle ──
  pi.on("session_start", async (_event, ctx) => {
    activeCtx = ctx;
    try { ctx.ui?.setStatus?.("mcpctl", "connecting..."); } catch { /* ignore */ }
    const cfg = await readConfig();
    mcplocalUrl = cfg.mcplocalUrl;
    const creds = await readCreds();
    token = creds?.token ?? undefined;

    const piState = await readPiState();
    const project = piState.project ?? (await findProjectMarker());
    if (!project) {
      try { ctx.ui?.setStatus?.("mcpctl", "no project (use /mcpctl)"); } catch { /* ignore */ }
      return;
    }
    activeProject = project;
    try {
      const r = await reconcileTools(project);
      if (ctx.hasUI) ctx.ui.notify(`mcpctl: project '${project}' — ${r.tools.length} tool(s) ready`, "info");
      setStatus();
    } catch (e) {
      if (ctx.hasUI) ctx.ui.notify(`mcpctl: could not load tools — ${(e as Error).message}`, "error");
    }
  });

  pi.on("session_shutdown", () => {
    void closeSession();
  });
}