decorated-pi 0.5.4 → 0.6.0

package/hooks/mcp.ts

@@ -0,0 +1,340 @@
+/**
+ * mcp — MCP connection lifecycle (session_start → connect, session_shutdown → disconnect).
+ *
+ * Tool registration is in tools/mcp/; this hook only manages the
+ * activeConnections array and cache persistence.
+ */
+
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { McpConnection } from "../tools/mcp/client.js";
+import { resolveMcpConfigs, type McpServerConfig } from "../tools/mcp/config.js";
+import { loadMcpCache, loadScopedMcpCache, updateServerCache, cleanupStaleCache, type McpToolCache } from "../tools/mcp/cache.js";
+import { buildMcpTool } from "../tools/mcp/tool-definition.js";
+import type { Module, Skeleton } from "./skeleton.js";
+
+let activeConnections: McpConnection[] = [];
+let allServers = new Map<string, any>();
+let cachedConfigs: McpServerConfig[] = [];
+let cachedCwd = "";
+
+interface ServerStatus {
+  name: string;
+  url: string;
+  source: string;
+  state: "connecting" | "connected" | "failed" | "disabled" | "waiting_reload";
+  toolCount: number;
+  tools: Array<{ name: string; description?: string; inputSchema?: Record<string, unknown> }>;
+  error?: string;
+}
+
+function cacheScopeForSource(source: string): "global" | "project" {
+  return source === "project" ? "project" : "global";
+}
+
+async function connectAll(
+  configs: McpServerConfig[],
+  registry: any,
+  ctx?: any,
+): Promise<{ schemaChanges: string[]; hasNewServer: boolean }> {
+  const cache = loadMcpCache(cachedCwd);
+  const schemaChanges: string[] = [];
+
+  // IMPORTANT: do NOT reassign `allServers`. After /reload, multiple
+  // connectAll calls can race — if the new one reassigns, the old
+  // in-flight connections' `allServers.set("connected")` updates a
+  // detached Map that the UI never reads. Mutate in place instead so
+  // every caller's updates are visible.
+  allServers.clear();
+  for (const s of configs) {
+    allServers.set(s.name, { name: s.name, url: s.url ?? s.command ?? "(unknown)", source: s.source, state: "connecting", toolCount: 0, tools: [] });
+  }
+
+  // Watchdog: the inner Promise.race in McpConnection.connect() has a
+  // 30s setTimeout, but in some environments (MCP SDK fetch hanging,
+  // event-loop pressure) that timer never fires and the connection
+  // promise never settles. After 35s we force-mark any still-connecting
+  // server as failed so the UI doesn't get stuck on "connecting" forever,
+  // and we make connectAll() return so a subsequent /reload can start a
+  // fresh connection attempt without racing a zombie Promise.all.
+  let watchdogFired = false;
+  const watchdog = new Promise<void>((resolve) => {
+    setTimeout(() => {
+      watchdogFired = true;
+      for (const config of configs) {
+        const server = allServers.get(config.name);
+        if (server && server.state === "connecting") {
+          server.state = "failed";
+          server.error = "Watchdog: connection did not settle within 35s (inner timeout missed)";
+        }
+      }
+      resolve();
+    }, 35_000);
+  });
+
+  // Connect SERIALLY, not via Promise.all. Three concurrent
+  // conn.connect() calls in a Promise.all fan-out reliably hang in the
+  // user's environment (likely the MCP SDK's HTTP transport is sensitive
+  // to concurrent DNS / connection-pool state). refreshServerCache
+  // (single, sequential) works fine — and pressing 'r' on a hung
+  // server in the UI also recovers it, which is what pointed us here.
+  let hasNewServer = false;
+  for (const server of configs) {
+    if (watchdogFired) break; // remaining servers are already failed
+
+    const conn = new McpConnection(server.name, server);
+    conn.source = server.source;
+    try {
+      await conn.connect(30_000);
+      if (watchdogFired) {
+        try { await conn.disconnect(); } catch { /* ignore */ }
+        return { schemaChanges, hasNewServer };
+      }
+      activeConnections.push(conn);
+      const actualTools = conn.tools.map((t) => ({ name: t.name, description: t.description, inputSchema: t.inputSchema }));
+      const cachedEntry = cache?.servers[server.name];
+      if (cachedEntry && cachedEntry.tools.length > 0) {
+        const cachedToolNames = new Set(cachedEntry.tools.map(t => t.name));
+        const added = actualTools.filter(t => !cachedToolNames.has(t.name));
+        const removed = cachedEntry.tools.filter(t => !cachedToolNames.has(t.name));
+        const changed = actualTools.filter(t => {
+          const cached = cachedEntry.tools.find(ct => ct.name === t.name);
+          return cached && JSON.stringify(cached.inputSchema) !== JSON.stringify(t.inputSchema);
+        });
+        if (added.length > 0 || removed.length > 0 || changed.length > 0) {
+          const parts: string[] = [];
+          if (added.length) parts.push(`${added.length} added`);
+          if (removed.length) parts.push(`${removed.length} removed`);
+          if (changed.length) parts.push(`${changed.length} changed`);
+          schemaChanges.push(`${server.name} (${parts.join(', ')})`);
+        }
+      }
+      allServers.set(server.name, {
+        name: server.name, url: server.url ?? server.command ?? "(unknown)", source: server.source,
+        state: "connected", toolCount: conn.tools.length, tools: actualTools,
+      });
+      const tools: McpToolCache[] = conn.tools.map(t => ({ name: t.name, description: t.description, inputSchema: t.inputSchema }));
+      updateServerCache(server.name, { tools, cachedAt: Date.now() }, cacheScopeForSource(server.source), cachedCwd || undefined);
+    } catch (err) {
+      if (watchdogFired) return { schemaChanges, hasNewServer };
+      const msg = err instanceof Error ? err.message : String(err);
+      allServers.set(server.name, {
+        name: server.name, url: server.url ?? server.command ?? "(unknown)", source: server.source,
+        state: "failed", toolCount: 0, tools: [], error: msg,
+      });
+    }
+  }
+
+  // Watchdog is no longer raced with the work — the for-loop above is
+  // already sequential and bounded by each conn.connect()'s own 30s
+  // timeout. We still keep a 35s ceiling for any single server just in
+  // case the inner timeout ever stops firing again.
+  void watchdog;
+
+  for (const server of configs) {
+    const cachedEntry = cache?.servers[server.name];
+    if (!cachedEntry || cachedEntry.tools.length === 0) { hasNewServer = true; break; }
+  }
+  return { schemaChanges, hasNewServer };
+}
+
+export function getMcpStatus(): ServerStatus[] {
+  const cache = loadMcpCache(cachedCwd);
+  const result: ServerStatus[] = [];
+  for (const config of cachedConfigs) {
+    const connected = allServers.get(config.name);
+    if (connected) { result.push(connected); continue; }
+    const cachedEntry = cache?.servers[config.name];
+    // Not in allServers: either disabled at startup (connectAll skipped it)
+    // or was toggled enabled after startup and needs /reload.
+    result.push({
+      name: config.name, url: config.url ?? config.command ?? "(unknown)", source: config.source,
+      state: config.enabled ? "waiting_reload" : "disabled",
+      toolCount: cachedEntry?.tools.length ?? 0, tools: cachedEntry?.tools ?? [],
+    });
+  }
+  return result;
+}
+
+export async function updateConfigEnabled(serverName: string, enabled: boolean): Promise<void> {
+  const config = cachedConfigs.find(c => c.name === serverName);
+  if (config) config.enabled = enabled;
+
+  if (!enabled) {
+    // Tear down the live connection immediately so teardownMcp() on
+    // /reload doesn't have to deal with a zombie conn.
+    const conn = activeConnections.find(c => c.serverName === serverName);
+    if (conn) {
+      try { await conn.disconnect(); } catch { /* ignore */ }
+      activeConnections = activeConnections.filter(c => c.serverName !== serverName);
+    }
+    // Mark the server as disabled in allServers (or add it if absent).
+    allServers.set(serverName, {
+      name: serverName,
+      url: config?.url ?? config?.command ?? "(unknown)",
+      source: config?.source ?? "unknown",
+      state: "disabled",
+      toolCount: 0,
+      tools: [],
+    });
+    return;
+  }
+
+  // Re-enabling at runtime: the config is now enabled but tools aren't
+  // registered until the next session_start. Mark it as waiting_reload
+  // so the user sees exactly why the server isn't usable yet.
+  allServers.set(serverName, {
+    name: serverName,
+    url: config?.url ?? config?.command ?? "(unknown)",
+    source: config?.source ?? "unknown",
+    state: "waiting reload",
+    toolCount: 0,
+    tools: [],
+  });
+}
+
+export async function refreshServerCache(serverName: string, registry: any): Promise<{ ok: boolean; error?: string }> {
+  const config = resolveMcpConfigs(cachedCwd).find(s => s.name === serverName);
+  if (!config) return { ok: false, error: `Server "${serverName}" not found in config.` };
+  const existing = activeConnections.find(c => c.serverName === serverName);
+  if (existing) {
+    try { await existing.disconnect(); } catch { /* ignore */ }
+    activeConnections = activeConnections.filter(c => c.serverName !== serverName);
+  }
+  const conn = new McpConnection(config.name, config);
+  conn.source = config.source;
+  try {
+    await conn.connect(30_000);
+    activeConnections.push(conn);
+    allServers.set(config.name, {
+      name: config.name, url: config.url ?? config.command ?? "(unknown)", source: config.source,
+      state: "connected", toolCount: conn.tools.length,
+      tools: conn.tools.map(t => ({ name: t.name, description: t.description, inputSchema: t.inputSchema })),
+    });
+    const tools: McpToolCache[] = conn.tools.map(t => ({ name: t.name, description: t.description, inputSchema: t.inputSchema }));
+    updateServerCache(config.name, { tools, cachedAt: Date.now() }, cacheScopeForSource(config.source), cachedCwd || undefined);
+    return { ok: true };
+  } catch (err) {
+    const msg = err instanceof Error ? err.message : String(err);
+    allServers.set(config.name, { name: config.name, url: config.url ?? config.command ?? "(unknown)", source: config.source, state: "failed", toolCount: 0, tools: [], error: msg });
+    return { ok: false, error: msg };
+  }
+}
+
+export function getActiveMcpConnections(): McpConnection[] {
+  return activeConnections;
+}
+
+export function getCachedMcpConfigs(): McpServerConfig[] {
+  return cachedConfigs;
+}
+
+/**
+ * Ensure one MCP server is ready and its tools are registered with pi.
+ *
+ * Behavior:
+ *   - Cache hit  → register tools from cache immediately (fast path).
+ *   - Cache miss → connect synchronously, write the cache, then
+ *     register the live tools.
+ *   - Connection failure + no cache → no tools registered (per design).
+ *
+ * Side effect: pushes a McpConnection into `activeConnections` so
+ * later `execute` calls (via buildMcpTool's findConnection callback)
+ * can find it.
+ */
+export async function ensureMcpServerReady(pi: ExtensionAPI, config: McpServerConfig, cwd?: string): Promise<void> {
+  const scope = cacheScopeForSource(config.source);
+  const cache = loadScopedMcpCache(scope, cwd);
+  const entry = cache?.servers[config.name];
+  const findConnection = (name: string) => activeConnections.find(c => c.serverName === name);
+
+  if (entry && entry.tools.length > 0) {
+    for (const t of entry.tools) {
+      try {
+        pi.registerTool(buildMcpTool(config, t, findConnection) as any);
+      } catch { /* duplicate name; previous registration still in effect */ }
+    }
+    return;
+  }
+
+  if (activeConnections.find(c => c.serverName === config.name)) return;
+
+  const conn = new McpConnection(config.name, config);
+  conn.source = config.source;
+  try {
+    await conn.connect(30_000);
+    activeConnections.push(conn);
+    const tools: McpToolCache[] = conn.tools.map(t => ({
+      name: t.name,
+      description: t.description,
+      inputSchema: t.inputSchema,
+    }));
+    updateServerCache(config.name, { tools, cachedAt: Date.now() }, scope, cwd);
+    for (const t of tools) {
+      try {
+        pi.registerTool(buildMcpTool(config, t, findConnection) as any);
+      } catch { /* duplicate name */ }
+    }
+  } catch (err) {
+    const msg = err instanceof Error ? err.message : String(err);
+    allServers.set(config.name, {
+      name: config.name,
+      url: config.url ?? config.command ?? "(unknown)",
+      source: config.source,
+      state: "failed",
+      toolCount: 0,
+      tools: [],
+      error: `No cache and initial connect failed: ${msg}`,
+    });
+  }
+}
+
+export const mcpModule: Module = {
+  name: "mcp",
+  hooks: {
+    session_start: [
+      async (_event, ctx) => {
+        await teardownMcp();
+        cachedCwd = ctx.cwd;
+        const configs = resolveMcpConfigs(ctx.cwd).sort((a, b) => a.name.localeCompare(b.name));
+        cachedConfigs = configs;
+        if (configs.length === 0) return;
+        cleanupStaleCache(configs, cachedCwd);
+        const enabledConfigs = configs.filter(s => s.enabled);
+
+        // connectAll only needs to handle servers that weren't already
+        // connected by index.ts's ensureMcpServerReady. We don't
+        // teardown first: the initial-sync connections are still
+        // healthy and we don't want to disconnect + reconnect on
+        // every session_start.
+        const toConnect = enabledConfigs.filter(s => !activeConnections.find(c => c.serverName === s.name));
+        if (toConnect.length === 0) return;
+
+        // connectAll runs in the background so the watchdog tests
+        // that mock hung connections don't block session_start.
+        void connectAll(toConnect, ctx.modelRegistry, ctx).then(({ schemaChanges }) => {
+          if (schemaChanges.length > 0 && ctx.hasUI) {
+            ctx.ui.notify(`mcp schema changed! please '/reload'`, "warning");
+          }
+        });
+      },
+    ],
+    session_shutdown: [
+      async () => { await teardownMcp(); },
+    ],
+  },
+};
+
+export async function teardownMcp(): Promise<void> {
+  await Promise.all(
+    activeConnections.map(async (conn) => {
+      try { await conn.disconnect(); } catch { /* ignore */ }
+    }),
+  );
+  activeConnections = [];
+  allServers.clear();
+  cachedConfigs = [];
+}
+
+export function setupMcp(sk: Skeleton, _pi: ExtensionAPI): void {
+  sk.register(mcpModule);
+}

package/hooks/normalize-codeblocks.ts

@@ -0,0 +1,88 @@
+/**
+ * normalize-codeblocks — fix markdown code block rendering in TUI.
+ *
+ * pi-tui's Markdown renderer doesn't expand tabs or handle CRLF in code
+ * blocks, causing garbled output (empty lines, misaligned content).
+ *
+ * This hook intercepts tool_result events and normalizes fenced code
+ * blocks (``` or ~~~):
+ *   - CRLF → LF (entire text)
+ *   - tabs → 4 spaces (inside code blocks only)
+ *   - strip trailing whitespace (inside code blocks only)
+ *
+ * Content outside code blocks is left untouched.
+ */
+
+import type { Module } from "./skeleton.js";
+
+/** Normalize CRLF line endings to LF. */
+function normalizeLineEndings(text: string): string {
+  return text.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
+}
+
+/** Normalize content inside markdown fenced code blocks.
+ *  - \t → 4 spaces
+ *  - strip trailing whitespace per line
+ *  Leaves everything outside code blocks untouched. */
+function normalizeCodeBlockContent(text: string): string {
+  const lines = text.split("\n");
+  let inCodeBlock = false;
+  let fenceChar = "";
+  let fenceLen = 0;
+
+  for (let i = 0; i < lines.length; i++) {
+    const line = lines[i];
+    const fenceMatch = line.match(/^(`{3,}|~{3,})/);
+
+    if (!inCodeBlock) {
+      if (fenceMatch) {
+        inCodeBlock = true;
+        fenceChar = fenceMatch[1][0];
+        fenceLen = fenceMatch[1].length;
+      }
+    } else {
+      // Check for closing fence (same char, length >= opening)
+      if (fenceMatch && fenceMatch[1][0] === fenceChar && fenceMatch[1].length >= fenceLen) {
+        inCodeBlock = false;
+      } else {
+        // Inside code block: normalize
+        lines[i] = line.replace(/\t/g, "    ").replace(/[ \t]+$/, "");
+      }
+    }
+  }
+
+  return lines.join("\n");
+}
+
+/** Normalize a tool result's text content.
+ *  Returns modified content array, or undefined if no changes. */
+export function normalizeToolResultContent(content: any[]): any[] | undefined {
+  let changed = false;
+  const newContent = content.map((block) => {
+    if (block.type !== "text" || typeof block.text !== "string") return block;
+    // First normalize line endings (entire text)
+    let text = normalizeLineEndings(block.text);
+    // Then normalize code block content
+    text = normalizeCodeBlockContent(text);
+    if (text !== block.text) {
+      changed = true;
+      return { ...block, text };
+    }
+    return block;
+  });
+  return changed ? newContent : undefined;
+}
+
+export const normalizeCodeblocksModule: Module = {
+  name: "normalize-codeblocks",
+  hooks: {
+    tool_result: [
+      (event) => {
+        if (!Array.isArray(event.content)) return undefined;
+        const newContent = normalizeToolResultContent(event.content);
+        if (!newContent) return undefined;
+        return { ...event, content: newContent };
+      },
+    ],
+  },
+};

package/hooks/pi-tool-filter.ts

@@ -0,0 +1,28 @@
+/**
+ * pi-tool-filter — unregister pi native tools that are replaced by our extensions.
+ *
+ *   edit   → replaced by patch
+ *   grep   → replaced by bash
+ *   find   → replaced by bash
+ *   ls     → replaced by bash
+ */
+
+import type { Module, Skeleton } from "./skeleton.js";
+
+const TOOLS_TO_DROP = new Set(["edit", "grep", "find", "ls"]);
+
+export const piToolFilterModule: Module = {
+  name: "pi-tool-filter",
+  hooks: {
+    session_start: [
+      (_event, ctx, pi) => {
+        const active = pi.getActiveTools();
+        pi.setActiveTools(active.filter((t) => !TOOLS_TO_DROP.has(t)));
+      },
+    ],
+  },
+};
+
+export function setupPiToolFilter(sk: Skeleton): void {
+  sk.register(piToolFilterModule);
+}

package/{extensions/safety → hooks/redact}/types.ts

@@ -1,9 +1,7 @@
 /**
- * Safety Detection — shared types and constants
+ * Redact — shared types and constants
  */
 
-// ─── Match types ──────────────────────────────────────────────────────────
-
 export type SecretMatchSource = "pattern" | "regex" | "entropy";
 
 export interface SecretMatch {
@@ -19,7 +17,6 @@ export interface SecretPattern {
   pattern: RegExp;
   minLength: number;
   allowsSpaces: boolean;
-  /** If true, skip safe-pattern exclusion (unambiguous prefix) */
   highConfidence: boolean;
 }
 
@@ -27,8 +24,6 @@ export interface DetectSecretsOptions {
   filePath?: string;
 }
 
-// ─── Internal types ──────────────────────────────────────────────────────
-
 export interface ConfigStringEntry {
   key: string;
   normalizedKey: string;
@@ -37,8 +32,6 @@ export interface ConfigStringEntry {
   end: number;
 }
 
-// ─── Constants ────────────────────────────────────────────────────────────
-
 export const MIN_SCAN_LENGTH = 10;
 export const CONFIG_VALUE_MIN_LENGTH = 32;
 export const CONFIG_FILE_EXTENSIONS = new Set([

package/hooks/redact.ts

@@ -0,0 +1,104 @@
+/**
+ * redact — secret redaction on read/bash tool_result.
+ *
+ * Detects API keys, tokens, passwords in tool output and replaces them with
+ * mask characters before the result enters the model context. Notifies the
+ * user via ctx.ui.notify with the count of redactions.
+ *
+ * Implementation is split into helper modules under hooks/redact/ to keep
+ * this file focused on the hook itself.
+ */
+
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type { Module, Skeleton } from "./skeleton.js";
+import { detectSecrets, maskSecret } from "./redact/detect.js";
+
+type TextContent = { type: "text"; text: string };
+
+function summarizeCommand(command: string, maxLength = 48): string {
+  const singleLine = command.replace(/\s+/g, " ").trim();
+  if (singleLine.length <= maxLength) return singleLine;
+  return `${singleLine.slice(0, maxLength - 1)}…`;
+}
+
+function formatRedactionContext(toolName: string, input: any): string {
+  if (toolName === "read") {
+    const filePath = input?.path ?? input?.file ?? input?.file_path;
+    return filePath ? `read ${filePath}` : "read";
+  }
+  if (toolName === "bash") {
+    const command = input?.command;
+    return typeof command === "string" && command.trim().length > 0
+      ? `bash ${summarizeCommand(command)}` : "bash";
+  }
+  return toolName;
+}
+
+export const redactModule: Module = {
+  name: "redact",
+  hooks: {
+    tool_result: [
+      (event, ctx) => {
+        // Patch failures emit error text that may include surrounding
+        // file content as a "did you mean" diagnostic. That path leaks
+        // secrets if the file is a config / .env / source with API keys.
+        // The success path returns the constant "Success" (no secrets)
+        // but we let it through anyway so the filter is uniform.
+        if (event.toolName !== "read" && event.toolName !== "bash" && event.toolName !== "patch") return;
+        if (!event.content || !Array.isArray(event.content)) return;
+
+        const textParts: Array<{ index: number; item: TextContent; text: string }> = [];
+        for (let i = 0; i < event.content.length; i++) {
+          const item = event.content[i];
+          if (item.type === "text" && typeof item.text === "string" && item.text.length > 0) {
+            textParts.push({ index: i, item: item as TextContent, text: item.text });
+          }
+        }
+        if (textParts.length === 0) return;
+
+        let totalCount = 0;
+        const counts: Record<"pattern" | "regex" | "entropy", number> = { pattern: 0, regex: 0, entropy: 0 };
+        const newContent = [...event.content];
+        const filePath = (event.input as any)?.path ?? (event.input as any)?.file ?? (event.input as any)?.file_path;
+
+        for (const { index, text, item } of textParts) {
+          const matches = detectSecrets(text, { filePath });
+          if (matches.length === 0) continue;
+          totalCount += matches.length;
+          let redacted = text;
+          for (const { start, end, source } of matches.sort((a, b) => b.start - a.start)) {
+            counts[source] += 1;
+            const original = redacted.slice(start, end);
+            redacted = redacted.slice(0, start) + maskSecret(original, source) + redacted.slice(end);
+          }
+          newContent[index] = { ...item, text: redacted };
+        }
+
+        if (totalCount === 0) return;
+        const label = totalCount === 1 ? "1 secret" : `${totalCount} secrets`;
+        const breakdown: string[] = [];
+        if (counts.pattern > 0) breakdown.push(`*:pattern=${counts.pattern}`);
+        if (counts.regex > 0) breakdown.push(`#:regex=${counts.regex}`);
+        if (counts.entropy > 0) breakdown.push(`?:entropy=${counts.entropy}`);
+        const suffix = breakdown.length > 0 ? ` · ${breakdown.join(" ")}` : "";
+        const contextLabel = formatRedactionContext(event.toolName, event.input);
+        if (ctx.hasUI) ctx.ui.notify(`🔒 [${contextLabel}] Redacted ${label}${suffix}`, "warning");
+        return { ...event, content: newContent };
+      },
+    ],
+  },
+};
+
+export function setupRedact(sk: Skeleton): void {
+  sk.register(redactModule);
+}
+
+/**
+ * System-prompt guidance for the redact hook — tells the LLM that
+ * masked values are real redactions (not "sk-xxx" literal strings),
+ * so it shouldn't try to read or reconstruct them.
+ */
+export const REDACT_GUIDANCE = [
+  "### Secret Masking, redacted values are real redactions",
+  "- When you see masked secret values (e.g. `sk-***...***` where `*`, `#`, or `?` are mask characters), the real value has been redacted by the system. Do not attempt to read or guess it. If you need the secret, use tools like `jq` or `grep` to extract it from the original source file.",
+].join("\n");