memtrace 0.8.9 → 0.8.11

package/pi-package/extension/index.ts

@@ -0,0 +1,764 @@
+// Memtrace Pi extension entry point.
+//
+// Makes Memtrace a first-class citizen of the Pi coding agent:
+//   • spawns `memtrace mcp` over stdio and speaks MCP JSON-RPC itself
+//   • registers every Memtrace tool as a NATIVE pi tool (memtrace_<name>)
+//     with real Typebox schemas discovered live from the server
+//   • ships /memtrace slash commands and a live status footer
+//   • injects a concise memtrace-first directive when the skill isn't loaded
+//
+// No build step — pi loads this via jiti. Only Node built-ins at runtime.
+
+import { writeFileSync, mkdirSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join, resolve, dirname, sep } from "node:path";
+
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import {
+  truncateHead,
+  formatSize,
+  DEFAULT_MAX_BYTES,
+  DEFAULT_MAX_LINES,
+} from "@earendil-works/pi-coding-agent";
+import type { TextContent, ImageContent } from "@earendil-works/pi-ai";
+
+import { MemtraceClient, type McpToolDef, type McpContentBlock } from "./mcp-client.ts";
+import { jsonSchemaToTypebox, TOOL_OVERRIDES, piToolName, memtraceName } from "./schema.ts";
+import { loadConfig, type MemtraceConfig } from "./config.ts";
+
+interface State {
+  cfg: MemtraceConfig;
+  client: MemtraceClient | null;
+  connecting: Promise<MemtraceClient> | null;
+  installed: boolean | null;
+  installedVersion: string | null;
+  registered: Set<string>;
+  cachedTools: McpToolDef[];
+  indexedRepos: Array<{ repo_id: string; path?: string; [k: string]: unknown }>;
+  cwdRepoId: string | null;
+  toolsEverRegistered: boolean;
+}
+
+const NO_TOOLS: McpToolDef[] = [];
+
+function createState(cwd: string): State {
+  return {
+    cfg: loadConfig(cwd),
+    client: null,
+    connecting: null,
+    installed: null,
+    installedVersion: null,
+    registered: new Set(),
+    cachedTools: NO_TOOLS,
+    indexedRepos: [],
+    cwdRepoId: null,
+    toolsEverRegistered: false,
+  };
+}
+
+let state: State = createState(process.cwd());
+
+// ---- helpers --------------------------------------------------------------
+
+function cachePath(): string {
+  return join(process.env.HOME ?? tmpdir(), ".pi/agent/memtrace-tools-cache.json");
+}
+
+function writeToolCache(tools: McpToolDef[]): void {
+  try {
+    const slim = tools.map((t) => ({ name: t.name, description: t.description, inputSchema: t.inputSchema }));
+    mkdirSync(dirname(cachePath()), { recursive: true });
+    writeFileSync(cachePath(), JSON.stringify(slim, null, 2));
+  } catch {
+    /* cache is best-effort */
+  }
+}
+
+function readToolCache(): McpToolDef[] {
+  try {
+    // Lazily require to avoid import cost when unused.
+    // eslint-disable-next-line @typescript-eslint/no-require-imports
+    const { readFileSync } = require("node:fs") as typeof import("node:fs");
+    const raw = readFileSync(cachePath(), "utf8");
+    const parsed = JSON.parse(raw);
+    return Array.isArray(parsed) ? (parsed as McpToolDef[]) : [];
+  } catch {
+    return [];
+  }
+}
+
+function looksLikeUnknownCommand(err: unknown, client: MemtraceClient | null): boolean {
+  const hay = `${(err as Error)?.message ?? ""}\n${client?.lastStderr ?? ""}`.toLowerCase();
+  return /unrecognized|unknown (sub)?command|no such|not found|invalid|unsubcommand|tip: a similar|usage: memtrace/.test(
+    hay,
+  );
+}
+
+async function checkInstalled(pi: ExtensionAPI): Promise<boolean> {
+  if (state.installed !== null) return state.installed;
+  try {
+    const res = await pi.exec(state.cfg.command, ["--version"], { timeout: 5000 });
+    state.installed = res.code === 0;
+    state.installedVersion = (res.stdout || res.stderr || "").trim().split("\n")[0] || null;
+  } catch {
+    state.installed = false;
+    state.installedVersion = null;
+  }
+  return state.installed;
+}
+
+/** Ensure the MCP server is connected. Idempotent. Throws if it cannot connect. */
+function ensureConnected(opts: { timeoutMs?: number; force?: boolean } = {}): Promise<MemtraceClient> {
+  if (state.connecting) return state.connecting;
+  if (state.client && state.client.connected && !opts.force) return Promise.resolve(state.client);
+  if (opts.force && state.client) {
+    state.client.dispose();
+    state.client = null;
+  }
+
+  const connect = async (): Promise<MemtraceClient> => {
+    const installed = state.installed ?? (await checkInstalled(getPi()));
+    if (!installed) {
+      throw new Error(
+        `memtrace binary "${state.cfg.command}" not found. Install it with: npm install -g memtrace`,
+      );
+    }
+    const make = (args: string[]) =>
+      new MemtraceClient({
+        command: state.cfg.command,
+        args,
+        env: state.cfg.env,
+        cwd: state.cfg.cwd,
+        requestTimeoutMs: state.cfg.requestTimeoutMs,
+        debug: process.env.MEMTRACE_PI_DEBUG === "1",
+      });
+
+    let client = make(state.cfg.args);
+    try {
+      await raceTimeout(client.start(), state.cfg.startupTimeoutMs, "connect");
+    } catch (e) {
+      client.dispose();
+      // Legacy fallback: a `serve` subcommand never shipped in memtrace, but
+      // older docs/configs may still reference it. Retry with the real `mcp`
+      // stdio subcommand before giving up.
+      if (state.cfg.args.join(" ") === "serve" && looksLikeUnknownCommand(e, client)) {
+        client = make(["mcp"]);
+        await raceTimeout(client.start(), state.cfg.startupTimeoutMs, "connect");
+      } else {
+        throw e;
+      }
+    }
+    state.client = client;
+    await refreshToolsAndCoverage(client);
+    return client;
+  };
+
+  const timeoutMs = opts.timeoutMs ?? state.cfg.requestTimeoutMs;
+  state.connecting = raceTimeout(connect(), timeoutMs, "ensureConnected").finally(() => {
+    state.connecting = null;
+  });
+  return state.connecting;
+}
+
+async function refreshToolsAndCoverage(client: MemtraceClient): Promise<void> {
+  try {
+    const tools = await client.listTools();
+    registerAllTools(tools);
+    writeToolCache(tools);
+  } catch (e) {
+    debug(`tools/list failed: ${(e as Error).message}`);
+  }
+  await refreshIndexCoverage();
+}
+
+function debug(msg: string): void {
+  if (process.env.MEMTRACE_PI_DEBUG === "1") console.error(`[memtrace-pi] ${msg}`);
+}
+
+function raceTimeout<T>(p: Promise<T>, ms: number, label: string): Promise<T> {
+  return new Promise<T>((resolveP, rejectP) => {
+    let done = false;
+    const t = setTimeout(() => {
+      if (!done) {
+        done = true;
+        rejectP(new Error(`memtrace ${label} timed out after ${ms}ms`));
+      }
+    }, ms);
+    p.then(
+      (v) => {
+        if (!done) {
+          done = true;
+          clearTimeout(t);
+          resolveP(v);
+        }
+      },
+      (e) => {
+        if (!done) {
+          done = true;
+          clearTimeout(t);
+          rejectP(e);
+        }
+      },
+    );
+  });
+}
+
+// ---- tool registration ----------------------------------------------------
+
+function registerAllTools(tools: McpToolDef[]): void {
+  const pi = getPi();
+  for (const tool of tools) {
+    const piName = piToolName(tool.name);
+    const override = TOOL_OVERRIDES[tool.name] ?? {};
+    const parameters =
+      tool.inputSchema && Object.keys(tool.inputSchema).length > 0
+        ? jsonSchemaToTypebox(tool.inputSchema)
+        : Type.Object({});
+
+    pi.registerTool({
+      name: piName,
+      label: override.label ?? `Memtrace: ${tool.name}`,
+      description: override.description ?? tool.description ?? `Memtrace ${tool.name} tool.`,
+      ...(override.promptSnippet ? { promptSnippet: override.promptSnippet } : {}),
+      ...(override.promptGuidelines ? { promptGuidelines: override.promptGuidelines } : {}),
+      parameters,
+      async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
+        return runMemtraceTool(tool.name, (params as Record<string, unknown> | undefined) ?? {}, signal);
+      },
+    });
+    state.registered.add(piName);
+  }
+  state.toolsEverRegistered = true;
+  debug(`registered ${tools.length} memtrace tools`);
+}
+
+async function runMemtraceTool(
+  name: string,
+  args: Record<string, unknown>,
+  signal: AbortSignal | undefined,
+): Promise<{ content: (TextContent | ImageContent)[]; details: Record<string, unknown> }> {
+  let client: MemtraceClient;
+  try {
+    client = await ensureConnected({ timeoutMs: state.cfg.requestTimeoutMs });
+  } catch (e) {
+    const msg = (e as Error).message;
+    return {
+      content: [{ type: "text", text: `memtrace is not available: ${msg}` }],
+      details: { error: msg, connected: false },
+    };
+  }
+
+  if (signal?.aborted) {
+    return { content: [{ type: "text", text: "cancelled" }], details: { aborted: true } };
+  }
+
+  let result;
+  try {
+    result = await client.callTool(name, args);
+  } catch (e) {
+    // Server died mid-call — try one reconnect before giving up.
+    debug(`callTool ${name} failed: ${(e as Error).message}; reconnecting`);
+    try {
+      state.client?.dispose();
+      state.client = null;
+      client = await ensureConnected({ timeoutMs: state.cfg.requestTimeoutMs, force: true });
+      result = await client.callTool(name, args);
+    } catch (e2) {
+      const msg = (e2 as Error).message;
+      return {
+        content: [{ type: "text", text: `memtrace call "${name}" failed: ${msg}` }],
+        details: { error: msg, connected: false },
+      };
+    }
+  }
+
+  const mapped = mapContent(result.content);
+  // Surface MCP soft errors to the model via content (do not throw if there is text).
+  if (result.isError && mapped.length === 0) {
+    throw new Error(`memtrace "${name}" returned an error with no content`);
+  }
+  return {
+    content: mapped,
+    details: { isError: !!result.isError, tool: name },
+  };
+}
+
+function mapContent(blocks: McpContentBlock[]): (TextContent | ImageContent)[] {
+  const out: (TextContent | ImageContent)[] = [];
+  let textBuf = "";
+  const flush = () => {
+    if (textBuf) {
+      const { text, fullFile } = truncateForContext(textBuf);
+      out.push({ type: "text", text });
+      if (fullFile) out.push({ type: "text", text: `Full output saved to: ${fullFile}` });
+      textBuf = "";
+    }
+  };
+  for (const b of blocks) {
+    if (b.type === "text" && typeof b.text === "string") {
+      textBuf += (textBuf ? "\n" : "") + b.text;
+    } else if (b.type === "image" && typeof b.data === "string") {
+      flush();
+      out.push({ type: "image", data: b.data, mimeType: b.mimeType ?? "image/png" });
+    } else if (b.type === "resource") {
+      flush();
+      const text = (b as { resource?: { text?: string; uri?: string } }).resource?.text;
+      if (text) out.push({ type: "text", text });
+    } else {
+      flush();
+      out.push({ type: "text", text: JSON.stringify(b) });
+    }
+  }
+  flush();
+  return out;
+}
+
+function truncateForContext(text: string): { text: string; fullFile?: string } {
+  const tr = truncateHead(text, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES });
+  if (!tr.truncated) return { text };
+  const fullFile = writeTemp(text);
+  return {
+    text:
+      tr.content +
+      `\n\n[Output truncated: ${tr.outputLines}/${tr.totalLines} lines ` +
+      `(${formatSize(tr.outputBytes)}/${formatSize(tr.totalBytes)}). Full output: ${fullFile}]`,
+    fullFile,
+  };
+}
+
+function writeTemp(text: string): string {
+  const path = join(tmpdir(), `memtrace-${Date.now()}-${Math.random().toString(36).slice(2, 8)}.txt`);
+  try {
+    writeFileSync(path, text, { mode: 0o600 });
+  } catch {
+    return "(failed to save full output)";
+  }
+  return path;
+}
+
+// ---- index coverage -------------------------------------------------------
+
+async function refreshIndexCoverage(): Promise<void> {
+  if (!state.client || !state.client.connected) return;
+  try {
+    const res = await state.client.callTool("list_indexed_repositories", {});
+    const text = (res.content?.[0] as { text?: string } | undefined)?.text ?? "";
+    const parsed = tryJson(text);
+    if (Array.isArray(parsed)) {
+      state.indexedRepos = parsed as State["indexedRepos"];
+    }
+    state.cwdRepoId = resolveCwdRepo();
+    updateStatus(getPi());
+  } catch (e) {
+    debug(`list_indexed_repositories failed: ${(e as Error).message}`);
+  }
+}
+
+function tryJson(text: string): unknown {
+  try {
+    return JSON.parse(text);
+  } catch {
+    // Some servers wrap JSON in markdown fences; try to extract.
+    const m = text.match(/```(?:json)?\s*([\s\S]*?)```/i);
+    if (m) {
+      try {
+        return JSON.parse(m[1]);
+      } catch { /* noop */ }
+    }
+    return null;
+  }
+}
+
+function resolveCwdRepo(): string | null {
+  const cwd = process.cwd();
+  let best: { repo_id: string; depth: number } | null = null;
+  for (const repo of state.indexedRepos) {
+    const rp = repo.path as string | undefined;
+    if (!rp) continue;
+    const abs = resolve(rp);
+    if (cwd === abs || cwd.startsWith(abs + sep)) {
+      const depth = abs.split(sep).length;
+      if (!best || depth > best.depth) best = { repo_id: repo.repo_id, depth };
+    }
+  }
+  return best?.repo_id ?? null;
+}
+
+// ---- status footer --------------------------------------------------------
+
+function statusText(): string {
+  if (state.installed === false) return `memtrace: not installed — npm i -g ${state.cfg.command}`;
+  if (!state.client || !state.client.connected) {
+    return state.installed === true ? "memtrace: idle — /memtrace connect" : "memtrace: connecting…";
+  }
+  const n = state.indexedRepos.length;
+  if (n === 0) return "memtrace: connected · 0 repos indexed — /memtrace index";
+  if (state.cwdRepoId) return `memtrace: connected · repo "${state.cwdRepoId}" indexed`;
+  return `memtrace: connected · ${n} repo${n === 1 ? "" : "s"} (cwd not indexed — /memtrace index)`;
+}
+
+function updateStatus(pi: ExtensionAPI): void {
+  try {
+    pi.getActiveTools; // ensure pi is bound
+    // setStatus works in TUI + RPC; no-op in print/json.
+    // Access ctx via a session-scoped capture instead.
+  } catch { /* noop */ }
+  statusCtx?.ui.setStatus("memtrace", statusText());
+}
+
+let statusCtx: { ui: { setStatus: (k: string, v: string | undefined) => void } } | null = null;
+
+// ---- the pi handle (captured for async work outside handlers) -------------
+
+let piHandle: ExtensionAPI | null = null;
+function getPi(): ExtensionAPI {
+  if (!piHandle) throw new Error("extension not initialized");
+  return piHandle;
+}
+
+// ---- commands -------------------------------------------------------------
+
+function registerCommands(pi: ExtensionAPI): void {
+  pi.registerCommand("memtrace", {
+    description: "Memtrace status (connection, indexed repos, cwd coverage)",
+    handler: async (_args, ctx) => {
+      statusCtx = ctx;
+      const lines = buildStatusLines();
+      if (ctx.hasUI) ctx.ui.setWidget("memtrace:status", lines);
+      ctx.ui.notify(lines[0], state.installed === false ? "warning" : "info");
+    },
+  });
+
+  pi.registerCommand("memtrace:connect", {
+    description: "Connect (or reconnect) to the Memtrace MCP server and load tools",
+    handler: async (_args, ctx) => {
+      statusCtx = ctx;
+      ctx.ui.notify("Connecting to memtrace…", "info");
+      try {
+        await ensureConnected({ timeoutMs: state.cfg.requestTimeoutMs, force: true });
+        ctx.ui.notify(
+          `memtrace connected · ${state.registered.size} tools · ${state.indexedRepos.length} repo(s)`,
+          "info",
+        );
+      } catch (e) {
+        ctx.ui.notify(`memtrace connect failed: ${(e as Error).message}`, "error");
+      }
+      updateStatus(pi);
+    },
+  });
+
+  pi.registerCommand("memtrace:disconnect", {
+    description: "Disconnect the Memtrace MCP server",
+    handler: async (_args, ctx) => {
+      state.client?.dispose();
+      state.client = null;
+      state.connecting = null;
+      ctx.ui.notify("memtrace disconnected", "info");
+      statusCtx = ctx;
+      updateStatus(pi);
+    },
+  });
+
+  pi.registerCommand("memtrace:tools", {
+    description: "List registered Memtrace tools",
+    handler: async (_args, ctx) => {
+      const tools = pi.getAllTools().filter((t) => t.name.startsWith("memtrace_"));
+      if (tools.length === 0) {
+        ctx.ui.notify("No memtrace tools registered. Try /memtrace:connect.", "warning");
+        return;
+      }
+      const lines = tools.map((t) => `• ${t.name} — ${oneLine(t.description)}`).slice(0, 60);
+      if (ctx.hasUI) ctx.ui.setWidget("memtrace:tools", [`Memtrace tools (${tools.length}):`, ...lines]);
+      ctx.ui.notify(`${tools.length} memtrace tools`, "info");
+    },
+  });
+
+  pi.registerCommand("memtrace:index", {
+    description: "Index the current directory (or a path) into Memtrace",
+    getArgumentCompletions: (prefix: string) => {
+      const item = { value: ".", label: ". (current directory)" };
+      return item.value.startsWith(prefix) ? [item] : null;
+    },
+    handler: async (args, ctx) => {
+      statusCtx = ctx;
+      const target = args.trim() || ctx.cwd || process.cwd();
+      const abs = resolve(target);
+      await indexFlow(pi, ctx, abs, { incremental: false, clearExisting: false });
+    },
+  });
+
+  pi.registerCommand("memtrace:reindex", {
+    description: "Incrementally re-index the current repo in Memtrace",
+    handler: async (_args, ctx) => {
+      statusCtx = ctx;
+      const repo = state.cwdRepoId;
+      if (!repo) {
+        ctx.ui.notify("cwd is not under an indexed repo — use /memtrace:index first", "warning");
+        return;
+      }
+      const abs = process.cwd();
+      await indexFlow(pi, ctx, abs, { incremental: true, clearExisting: false });
+    },
+  });
+
+  pi.registerCommand("memtrace:setup", {
+    description: "Check the memtrace install and show next steps",
+    handler: async (_args, ctx) => {
+      statusCtx = ctx;
+      const installed = await checkInstalled(pi);
+      if (!installed) {
+        ctx.ui.notify(`memtrace not found. Install: npm install -g memtrace`, "warning");
+        if (ctx.hasUI) {
+          ctx.ui.setWidget("memtrace:setup", [
+            "Memtrace setup",
+            "",
+            "1) Install the binary:  npm install -g memtrace",
+            "2) (optional) Start the workspace daemon:  memtrace start",
+            "   (indexer + file watcher + HTTP API on :3030 — `memtrace mcp`",
+            "    attaches to it automatically, or self-starts an embedded",
+            "    MemDB sidecar if it isn't running)",
+            "3) Reload pi:           /reload",
+            "",
+            "Optional env: MEMTRACE_EMBED_KEY (embeddings API key).",
+          ]);
+        }
+        return;
+      }
+      ctx.ui.notify(`memtrace ${state.installedVersion ?? "installed"} ✓`, "info");
+      if (ctx.hasUI) {
+        ctx.ui.setWidget("memtrace:setup", [
+          `Memtrace ${state.installedVersion ?? ""} is installed.`,
+          "",
+          "If tools aren't connecting:",
+          "  • run:  memtrace start     (workspace daemon: indexer + watcher + HTTP API)",
+          "  • then: /memtrace:connect",
+          "",
+          "Server command in use:",
+          `  ${state.cfg.command} ${state.cfg.args.join(" ")}`,
+        ]);
+      }
+    },
+  });
+}
+
+function oneLine(s: string | undefined): string {
+  if (!s) return "";
+  return s.replace(/\s+/g, " ").trim().slice(0, 160);
+}
+
+function buildStatusLines(): string[] {
+  const lines: string[] = [];
+  if (state.installed === false) {
+    lines.push(`memtrace: NOT installed — run: npm install -g memtrace`);
+    return lines;
+  }
+  lines.push(`memtrace ${state.installedVersion ?? ""}`.trim());
+  lines.push(
+    state.client?.connected
+      ? `connection: connected (${serverInfoString()})`
+      : "connection: idle — /memtrace:connect",
+  );
+  lines.push(`tools registered: ${state.registered.size}`);
+  lines.push(`indexed repos: ${state.indexedRepos.length}`);
+  for (const r of state.indexedRepos.slice(0, 8)) {
+    lines.push(`  • ${r.repo_id}${r.path ? ` — ${r.path}` : ""}`);
+  }
+  lines.push(
+    state.cwdRepoId
+      ? `cwd coverage: repo "${state.cwdRepoId}"`
+      : `cwd coverage: not indexed — /memtrace:index`,
+  );
+  return lines;
+}
+
+function serverInfoString(): string {
+  const s = state.client?.server ?? {};
+  return [s.name, s.version].filter(Boolean).join(" ") || "memtrace mcp";
+}
+
+// ---- indexing flow (used by /memtrace:index and :reindex) -----------------
+
+async function indexFlow(
+  pi: ExtensionAPI,
+  ctx: { ui: { notify: (m: string, l: "info" | "warning" | "error") => void }; cwd: string },
+  absPath: string,
+  opts: { incremental: boolean; clearExisting: boolean },
+): Promise<void> {
+  let client: MemtraceClient;
+  try {
+    client = await ensureConnected({ timeoutMs: state.cfg.requestTimeoutMs });
+  } catch (e) {
+    ctx.ui.notify(`Cannot connect to memtrace: ${(e as Error).message}`, "error");
+    return;
+  }
+
+  ctx.ui.notify(`Indexing ${absPath} (incremental=${opts.incremental})…`, "info");
+  let job: { job_id?: string };
+  try {
+    job = (await client.callTool("index_directory", {
+      path: absPath,
+      incremental: opts.incremental,
+      clear_existing: opts.clearExisting,
+    })) as { job_id?: string };
+  } catch (e) {
+    ctx.ui.notify(`index_directory failed: ${(e as Error).message}`, "error");
+    return;
+  }
+
+  const jobId = job.job_id ?? (tryJson(((job as unknown as { content?: { text?: string }[] }).content?.[0]?.text) ?? "") as { job_id?: string })?.job_id;
+  if (!jobId) {
+    // Some servers index synchronously; just refresh coverage.
+    await refreshIndexCoverage();
+    ctx.ui.notify("Indexing submitted (no job_id returned) — refreshed repo list.", "info");
+    return;
+  }
+
+  const start = Date.now();
+  for (let i = 0; i < 100; i++) {
+    await sleep(2500);
+    let st: { status?: string; error?: string };
+    try {
+      st = (await client.callTool("check_job_status", { job_id: jobId })) as { status?: string; error?: string };
+    } catch (e) {
+      ctx.ui.notify(`check_job_status failed: ${(e as Error).message}`, "error");
+      return;
+    }
+    const status = st.status ?? "running";
+    if (status === "completed") {
+      await refreshIndexCoverage();
+      const secs = ((Date.now() - start) / 1000).toFixed(1);
+      ctx.ui.notify(`Indexing complete in ${secs}s · repo "${state.cwdRepoId ?? "?"}"`, "info");
+      return;
+    }
+    if (status === "failed") {
+      ctx.ui.notify(`Indexing failed: ${st.error ?? "unknown error"}`, "error");
+      return;
+    }
+    if (i % 4 === 0) ctx.ui.notify(`Indexing… [${status}] ${((Date.now() - start) / 1000).toFixed(0)}s`, "info");
+  }
+  ctx.ui.notify(
+    `Indexing still running after ~5 min (job_id=${jobId}). It continues server-side — re-check with /memtrace later.`,
+    "warning",
+  );
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((r) => setTimeout(r, ms));
+}
+
+// ---- the always-on status tool -------------------------------------------
+
+function registerStatusTool(pi: ExtensionAPI): void {
+  pi.registerTool({
+    name: "memtrace_status",
+    label: "Memtrace: Status",
+    description:
+      "Report Memtrace connection state, whether the memtrace binary is installed, indexed repositories, and whether the current working directory is covered. Call this before any other memtrace_* tool if you are unsure whether Memtrace is available.",
+    promptSnippet: "Check if Memtrace is connected and the cwd is indexed",
+    promptGuidelines: [
+      "Call memtrace_status first if you are unsure whether Memtrace is connected or the repo is indexed. It returns repo_id values you need for other memtrace_* tools.",
+    ],
+    parameters: Type.Object({}),
+    async execute() {
+      const installed = state.installed ?? (await checkInstalled(pi).catch(() => false));
+      const connected = !!state.client?.connected;
+      const payload = {
+        installed,
+        installedVersion: state.installedVersion,
+        connected,
+        server: state.client?.server ?? null,
+        toolsRegistered: state.registered.size,
+        indexedRepos: state.indexedRepos.map((r) => ({ repo_id: r.repo_id, path: r.path })),
+        cwd: process.cwd(),
+        cwdCoveredRepoId: state.cwdRepoId,
+        hint: !installed
+          ? "Run: npm install -g memtrace"
+          : !connected
+            ? "Run /memtrace:connect (or ask the user to run `memtrace start` then /memtrace:connect)"
+            : !state.cwdRepoId
+              ? "Index the current repo with memtrace_index_directory (absolute path), then call memtrace_list_indexed_repositories."
+              : "Memtrace is ready — use memtrace_* tools for code discovery, impact, and change history.",
+      };
+      return {
+        content: [{ type: "text", text: JSON.stringify(payload, null, 2) }],
+        details: payload,
+      };
+    },
+  });
+  state.registered.add("memtrace_status");
+}
+
+// ---- system-prompt injection ---------------------------------------------
+
+const FIRST_DIRECTIVE = [
+  "",
+  "## Memtrace code-graph tools (memtrace_*)",
+  "Memtrace is the memory layer of this codebase: a typed knowledge graph of every symbol, call, community, process, and API, with a time dimension. For the indexed repo:",
+  "- Code discovery / 'where is X' / 'how does X work' → memtrace_find_symbol / memtrace_find_code → memtrace_get_symbol_context. Do not grep/glob to locate indexed source.",
+  "- 'What breaks if I change X?' → memtrace_get_impact (transitive blast radius), not a reference search.",
+  "- 'What changed since <date>?' → memtrace_get_evolution (requires repo_id + from; no `days` param), not git log.",
+  "- About to view source Memtrace pointed at → use memtrace_get_source_window or a bounded read (offset/limit), never a whole-file read.",
+  "- Call memtrace_status if unsure whether Memtrace is available, and memtrace_list_indexed_repositories to get repo_id values.",
+  "Prefer these over read/grep/find for anything inside the indexed repository. See the memtrace-first skill for the full routing rules.",
+].join("\n");
+
+function hasMemtraceFirstSkill(skills: unknown): boolean {
+  if (!Array.isArray(skills)) return false;
+  return skills.some((s) => {
+    const name = (s as { name?: string; path?: string })?.name ?? (s as { path?: string })?.path ?? "";
+    return /memtrace[-_]first/i.test(String(name));
+  });
+}
+
+// ---- entry point ----------------------------------------------------------
+
+export default async function memtraceExtension(pi: ExtensionAPI): Promise<void> {
+  piHandle = pi;
+  state = createState(process.cwd());
+
+  registerStatusTool(pi);
+  registerCommands(pi);
+
+  // Register tools from the on-disk cache so they appear in the very first
+  // prompt even when the server is slow to connect (warm start).
+  state.cachedTools = readToolCache();
+  if (state.cachedTools.length > 0) {
+    registerAllTools(state.cachedTools);
+  }
+
+  pi.on("session_start", async (_event, ctx) => {
+    statusCtx = ctx;
+    if (state.installed === null) void checkInstalled(pi).then(() => updateStatus(pi));
+    updateStatus(pi);
+    // Reconnect for this session in the background (non-blocking).
+    void ensureConnected({ timeoutMs: state.cfg.startupTimeoutMs })
+      .then(() => updateStatus(pi))
+      .catch((e) => debug(`startup connect skipped: ${(e as Error).message}`));
+  });
+
+  pi.on("session_shutdown", async () => {
+    state.client?.dispose();
+    state.client = null;
+    state.connecting = null;
+  });
+
+  pi.on("before_agent_start", async (event, _ctx) => {
+    if (!state.toolsEverRegistered && state.cachedTools.length === 0) return;
+    if (hasMemtraceFirstSkill((event.systemPromptOptions as { skills?: unknown })?.skills)) return;
+    return { systemPrompt: event.systemPrompt + "\n" + FIRST_DIRECTIVE };
+  });
+
+  // Eager connect with a short timeout so live tools (and their real schemas)
+  // replace the cached ones before the first prompt whenever possible.
+  try {
+    await raceTimeout(
+      ensureConnected({ timeoutMs: state.cfg.startupTimeoutMs }),
+      state.cfg.startupTimeoutMs,
+      "startup",
+    );
+  } catch (e) {
+    debug(`startup connect deferred: ${(e as Error).message}`);
+  }
+  updateStatus(pi);
+}