decorated-pi 0.4.0 → 0.5.0

package/extensions/mcp/index.ts

@@ -1,102 +1,332 @@
-import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { Type } from "typebox";
+import { keyHint, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { Text } from "@earendil-works/pi-tui";
 import { McpConnection } from "./client.js";
-import { resolveMcpConfigs } from "./builtin.js";
+import {
+  resolveMcpConfigs, saveProjectMcpDescription,
+  loadMcpCache, updateServerCache, cleanupStaleCache,
+  type McpServerConfig, type McpToolCache,
+} from "./builtin.js";
+import {
+  getMcpBrokerModelKey, getCompactModelKey,
+  getMcpDescription, setMcpDescription,
+  parseModelKey,
+} from "../settings.js";
 
 export interface McpServerStatus {
   name: string;
   url: string;
   source: string;
-  state: "connecting" | "connected" | "failed";
+  description?: string;
+  state: "connecting" | "connected" | "failed" | "disabled";
   toolCount: number;
-  tools: Array<{ name: string; description: string }>;
+  tools: Array<{ name: string; description?: string; inputSchema?: Record<string, unknown> }>;
   error?: string;
 }
 
 let activeConnections: McpConnection[] = [];
 let allServers = new Map<string, McpServerStatus>();
+let cachedConfigs: McpServerConfig[] = [];
 let connectPromise: Promise<void> | null = null;
+let cachedCwd = "";
 
-export function setupMcp(pi: ExtensionAPI) {
-  pi.on("session_start", (_event, ctx: ExtensionContext) => {
-    void (async () => {
-      await teardownMcp();
-
-      const configs = resolveMcpConfigs(ctx.cwd);
-      if (configs.length === 0) return;
-
-      // Initialise every target server as "connecting"
-      allServers = new Map(
-        configs.map((s) => [
-          s.name,
-          {
-            name: s.name,
-            url: s.url,
-            source: s.source,
-            state: "connecting" as const,
-            toolCount: 0,
-            tools: [],
-          },
-        ]),
-      );
-
-      connectPromise = Promise.all(
-        configs.map(async (server) => {
-          const conn = new McpConnection(server.name, server.url);
-          conn.source = server.source;
+const MCP_RESULT_FOLD_LINES = 45;
+
+function trimTrailingEmptyLines(lines: string[]): string[] {
+  let end = lines.length;
+  while (end > 0 && lines[end - 1] === "") end -= 1;
+  return lines.slice(0, end);
+}
+
+function collapseMcpText(text: string, maxLines = MCP_RESULT_FOLD_LINES) {
+  const lines = trimTrailingEmptyLines(text.split("\n"));
+  const totalLines = lines.length;
+  const displayLines = lines.slice(0, maxLines);
+  const remainingLines = Math.max(0, totalLines - maxLines);
+  return { totalLines, displayLines, remainingLines };
+}
+
+function getTextContent(result: { content?: Array<{ type: string; text?: string }> }): string {
+  return (result.content ?? [])
+    .filter((c): c is { type: "text"; text?: string } => c.type === "text")
+    .map((c) => c.text ?? "")
+    .join("\n");
+}
+
+function formatMcpResultText(text: string, expanded: boolean, theme: any): string {
+  const { totalLines, displayLines, remainingLines } = collapseMcpText(
+    text,
+    expanded ? Number.MAX_SAFE_INTEGER : MCP_RESULT_FOLD_LINES,
+  );
+  let rendered = displayLines.join("\n") ? theme.fg("toolOutput", displayLines.join("\n")) : "";
+  if (!expanded && remainingLines > 0) {
+    rendered += `${theme.fg("muted", `\n... (${remainingLines} more lines, ${totalLines} total,`)} ${keyHint("app.tools.expand", "to expand")})`;
+  }
+  return rendered;
+}
+
+function renderMcpResult(result: any, options: { expanded: boolean }, theme: any, context: any) {
+  const component = context.lastComponent ?? new Text("", 0, 0);
+  component.setText(formatMcpResultText(getTextContent(result), options.expanded, theme));
+  return component;
+}
+
+// ── config helpers ────────────────────────────────────────────────────────
+
+export function updateConfigEnabled(serverName: string, enabled: boolean): void {
+  const config = cachedConfigs.find(c => c.name === serverName);
+  if (config) config.enabled = enabled;
+  const server = allServers.get(serverName);
+  if (server) {
+    if (!enabled) {
+      // Stash the real connection state, set to disabled
+      server.state = "disabled";
+    } else {
+      // Re-enable: if there's still an active connection, restore it
+      const conn = activeConnections.find(c => c.serverName === serverName);
+      server.state = conn ? "connected" : "connecting";
+    }
+  }
+}
+
+// ── helpers ───────────────────────────────────────────────────────────────
+
+function serverDescription(s: McpServerConfig): string | undefined {
+  return s.description || getMcpDescription(s.name, cachedCwd);
+}
+
+function makeToolName(serverName: string, toolName: string): string {
+  return `${serverName}_${toolName}`;
+}
+
+function makeToolLabel(serverName: string, toolName: string, desc?: string): string {
+  return `MCP ${serverName}: ${toolName}${desc ? ` (${desc.slice(0, 20)})` : ""}`;
+}
+
+// ── cache helpers ─────────────────────────────────────────────────────────
+
+function cacheScopeForSource(source: string): "global" | "project" {
+  return source === "project" ? "project" : "global";
+}
+
+// ── auto-summary ──────────────────────────────────────────────────────────
+
+async function autoDescribeServer(
+  conn: McpConnection,
+  serverName: string,
+  registry: any,
+): Promise<string> {
+  const descs = conn.tools.map(t => `- ${t.name}: ${t.description || "(no description)"}`).join("\n");
+
+  const prompt = `Describe what this MCP server is and what it does, based on the tools it exposes. Start with action verbs directly, like a capability summary.
+
+Server: "${serverName}"
+Tools:
+${descs}
+
+Respond with ONLY one short sentence. No quotes.`;
+
+  return await summarizeWithBroker(registry, prompt) || `${serverName} MCP server (${conn.tools.length} tools)`;
+}
+
+async function summarizeWithBroker(registry: any, prompt: string): Promise<string | undefined> {
+  if (!registry) return undefined;
+
+  const brokerKey = getMcpBrokerModelKey() || getCompactModelKey();
+  const model = brokerKey
+    ? (() => {
+        const parsed = parseModelKey(brokerKey);
+        return parsed ? registry.find(parsed.provider, parsed.modelId) : undefined;
+      })()
+    : undefined;
+
+  if (!model) return undefined;
+
+  try {
+    const auth = await registry.getApiKeyAndHeaders(model);
+    if (!auth.ok) return undefined;
+
+    const { complete } = await import("@earendil-works/pi-ai");
+    const resp = await complete(model, {
+      systemPrompt: "You are a concise MCP server description generator.",
+      messages: [{
+        role: "user" as const,
+        content: [{ type: "text" as const, text: prompt }],
+        timestamp: Date.now(),
+      }],
+    }, {
+      maxTokens: 128,
+      apiKey: auth.apiKey ?? "",
+      headers: auth.headers,
+      signal: AbortSignal.timeout(15_000),
+    });
+
+    if (resp.stopReason === "error") return undefined;
+    return resp.content
+      .filter((c: any): c is { type: "text"; text: string } => c.type === "text")
+      .map((c: any) => c.text).join(" ").trim() || undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+// ── register cached tools ─────────────────────────────────────────────────
 
+function registerCachedTools(pi: ExtensionAPI, configs: McpServerConfig[]): void {
+  const cache = loadMcpCache(cachedCwd);
+  if (!cache) return;
+
+  for (const config of configs) {
+    if (!config.enabled) continue;
+    const entry = cache.servers[config.name];
+    if (!entry || entry.tools.length === 0) continue;
+
+    for (const t of entry.tools) {
+      const toolName = makeToolName(config.name, t.name);
+      const desc = t.description || `${t.name} (MCP tool)`;
+      pi.registerTool({
+        name: toolName,
+        label: makeToolLabel(config.name, t.name, t.description),
+        description: desc,
+        promptSnippet: desc || `MCP tool ${config.name}/${t.name}`,
+        renderResult: renderMcpResult,
+        parameters: t.inputSchema,
+        execute: async (_id, params, _signal, _update, _ctx) => {
+          const conn = activeConnections.find(c => c.serverName === config.name);
+          if (!conn) {
+            return {
+              content: [{ type: "text", text: `MCP server "${config.name}" is not connected.` }],
+              isError: true,
+              details: {},
+            };
+          }
           try {
-            await conn.connect();
-            activeConnections.push(conn);
-
-            for (const tool of conn.tools) {
-              const prefixedName = `${server.name}_${tool.name}`;
-              pi.registerTool({
-                name: prefixedName,
-                label: `MCP: ${tool.name}`,
-                description: tool.description,
-                promptSnippet: tool.description.slice(0, 120),
-                parameters: Type.Unsafe(tool.inputSchema as never),
-                execute: async (_toolCallId, params, _signal, _onUpdate, _ctx2) => {
-                  const text = await conn.callTool(
-                    tool.name,
-                    params as Record<string, unknown>,
-                  );
-                  return {
-                    content: [{ type: "text" as const, text }],
-                    isError: false,
-                    details: { server: server.name, tool: tool.name },
-                  };
-                },
-              });
-            }
-
-            allServers.set(server.name, {
-              name: server.name,
-              url: server.url,
-              source: server.source,
-              state: "connected",
-              toolCount: conn.tools.length,
-              tools: conn.tools.map((t) => ({ name: t.name, description: t.description })),
-            });
+            const text = await conn.callTool(t.name, params as Record<string, unknown>);
+            return { content: [{ type: "text", text }], isError: false, details: {} };
           } catch (err) {
             const msg = err instanceof Error ? err.message : String(err);
-            allServers.set(server.name, {
-              name: server.name,
-              url: server.url,
-              source: server.source,
-              state: "failed",
-              toolCount: 0,
-              tools: [],
-              error: msg,
-            });
+            return {
+              content: [{ type: "text", text: `MCP call failed on "${config.name}/${t.name}": ${msg}` }],
+              isError: true,
+              details: {},
+            };
+          }
+        },
+      });
+    }
+  }
+}
+
+// ── connect ───────────────────────────────────────────────────────────────
+
+async function connectAll(configs: McpServerConfig[], registry: any): Promise<void> {
+  allServers = new Map(
+    configs.map((s) => [
+      s.name,
+      {
+        name: s.name,
+        url: s.url ?? s.command ?? "(unknown)",
+        source: s.source,
+        description: serverDescription(s),
+        state: "connecting" as const,
+        toolCount: 0,
+        tools: [],
+      },
+    ]),
+  );
+
+  connectPromise = Promise.all(
+    configs.map(async (server) => {
+      const conn = new McpConnection(server.name, server);
+      conn.source = server.source;
+
+      try {
+        await conn.connect(30_000);
+        activeConnections.push(conn);
+
+        let desc = serverDescription(server);
+        if (!desc) {
+          desc = await autoDescribeServer(conn, server.name, registry);
+          if (desc) {
+            if (server.source === "project") saveProjectMcpDescription(cachedCwd, server.name, desc);
+            else setMcpDescription(server.name, desc);
           }
-        }),
-      ).then(() => undefined);
+        }
+
+        allServers.set(server.name, {
+          name: server.name,
+          url: server.url ?? server.command ?? "(unknown)",
+          source: server.source,
+          description: desc,
+          state: "connected",
+          toolCount: conn.tools.length,
+          tools: conn.tools.map((t) => ({
+            name: t.name,
+            description: t.description,
+            inputSchema: t.inputSchema,
+          })),
+        });
+
+        // Update cache with this server's tools
+        const tools: McpToolCache[] = conn.tools.map(t => ({
+          name: t.name,
+          description: t.description,
+          inputSchema: t.inputSchema,
+        }));
+        updateServerCache(
+          server.name,
+          { description: desc, tools, cachedAt: Date.now() },
+          cacheScopeForSource(server.source),
+          cachedCwd || undefined,
+        );
+      } catch (err) {
+        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,
+          description: serverDescription(server),
+          state: "failed",
+          toolCount: 0,
+          tools: [],
+          error: msg,
+        });
+      }
+    }),
+  ).then(() => undefined);
+
+  await connectPromise;
+  connectPromise = null;
+}
 
-      await connectPromise;
-      connectPromise = null;
-    })();
+// ── setup ─────────────────────────────────────────────────────────────────
+
+export function setupMcp(pi: ExtensionAPI) {
+  pi.on("session_start", async (_event, ctx: ExtensionContext) => {
+    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;
+
+    // Clean stale cache entries for removed servers
+    cleanupStaleCache(configs, cachedCwd);
+
+    const enabledConfigs = configs.filter(s => s.enabled);
+
+    // Register tools from cache — prompt-stable, works even if MCP is down
+    registerCachedTools(pi, enabledConfigs);
+
+    const needSummary = enabledConfigs.filter(s => !serverDescription(s));
+
+    if (needSummary.length === 0) {
+      // All servers have descriptions — connect in background, update cache
+      void connectAll(enabledConfigs, ctx.modelRegistry);
+      return;
+    }
+
+    // Some servers lack description — connect and auto-summarize synchronously
+    await connectAll(enabledConfigs, ctx.modelRegistry);
   });
 
   pi.on("session_shutdown", () => {
@@ -105,7 +335,103 @@ export function setupMcp(pi: ExtensionAPI) {
 }
 
 export function getMcpStatus(): McpServerStatus[] {
-  return [...allServers.values()];
+  const cache = loadMcpCache(cachedCwd);
+  const result: McpServerStatus[] = [];
+  for (const config of cachedConfigs) {
+    const connected = allServers.get(config.name);
+    if (connected) {
+      result.push(connected);
+    } else {
+      const cachedEntry = cache?.servers[config.name];
+      result.push({
+        name: config.name,
+        url: config.url ?? config.command ?? "(unknown)",
+        source: config.source,
+        description: serverDescription(config),
+        state: config.enabled ? "connecting" : "disabled",
+        toolCount: cachedEntry?.tools.length ?? 0,
+        tools: cachedEntry?.tools ?? [],
+      });
+    }
+  }
+  return result;
+}
+
+// ── refresh single server cache ───────────────────────────────────────────
+
+export const __mcpIndexTest = { collapseMcpText };
+
+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.` };
+
+  // Disconnect existing connection for this server
+  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);
+
+    let desc = serverDescription(config);
+    if (!desc) {
+      desc = await autoDescribeServer(conn, config.name, registry);
+      if (desc) {
+        if (config.source === "project") saveProjectMcpDescription(cachedCwd, config.name, desc);
+        else setMcpDescription(config.name, desc);
+      }
+    }
+
+    allServers.set(config.name, {
+      name: config.name,
+      url: config.url ?? config.command ?? "(unknown)",
+      source: config.source,
+      description: desc,
+      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,
+      { description: desc, 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,
+      description: serverDescription(config),
+      state: "failed",
+      toolCount: 0,
+      tools: [],
+      error: msg,
+    });
+    return { ok: false, error: msg };
+  }
 }
 
 async function teardownMcp(): Promise<void> {
@@ -120,4 +446,5 @@ async function teardownMcp(): Promise<void> {
   );
   activeConnections = [];
   allServers = new Map();
+  cachedConfigs = [];
 }

package/extensions/model-integration.ts

@@ -25,8 +25,8 @@ import { fileTypeFromFile } from "file-type";
 import { isContextOverflow, type Model } from "@earendil-works/pi-ai";
 import {
   loadConfig, saveConfig, parseModelKey, formatModelKey,
-  getImageModelKey, getCompactModelKey,
-  setImageModelKey, setCompactModelKey,
+  getImageModelKey, getCompactModelKey, getMcpBrokerModelKey,
+  setImageModelKey, setCompactModelKey, setMcpBrokerModelKey,
 } from "./settings.js";
 import * as fs from "node:fs";
 import * as os from "node:os";
@@ -172,6 +172,7 @@ function setupImageReadFallback(pi: ExtensionAPI) {
 
 const TAB_IMAGE = 0;
 const TAB_COMPACT = 1;
+const TAB_BROKER = 2;
 
 export class ModelPickerComponent extends Container {
   private searchInput: Input;
@@ -182,6 +183,7 @@ export class ModelPickerComponent extends Container {
   private activeTab = TAB_IMAGE;
   private imageKey: string | null;
   private compactKey: string | null;
+  private brokerKey: string | null;
   private allItems: { label: string; desc: string; model: Model<any> | null; modelName?: string }[] = [];
   private filtered: typeof this.allItems = [];
   private selectedIndex = 0;
@@ -197,6 +199,7 @@ export class ModelPickerComponent extends Container {
     this.onDone = onDone;
     this.imageKey = getImageModelKey();
     this.compactKey = getCompactModelKey();
+    this.brokerKey = getMcpBrokerModelKey();
 
     this.addChild(new DynamicBorder());
     this.addChild(new Spacer(1));
@@ -232,8 +235,8 @@ export class ModelPickerComponent extends Container {
     }
   }
 
-  private currentKey() { return this.activeTab === TAB_IMAGE ? this.imageKey : this.compactKey; }
-  private currentKind() { return this.activeTab === TAB_IMAGE ? "image" : "compact"; }
+  private currentKey() { return this.activeTab === TAB_IMAGE ? this.imageKey : this.activeTab === TAB_BROKER ? this.brokerKey : this.compactKey; }
+  private currentKind() { return this.activeTab === TAB_IMAGE ? "image" : this.activeTab === TAB_BROKER ? "broker" : "compact"; }
 
   private switchTab(tab: number) {
     this.activeTab = tab;
@@ -260,16 +263,21 @@ export class ModelPickerComponent extends Container {
     const t = this.theme;
     const im = this.activeTab === TAB_IMAGE ? t.fg("accent", "●") : "○";
     const cm = this.activeTab === TAB_COMPACT ? t.fg("accent", "●") : "○";
-    const il = this.activeTab === TAB_IMAGE ? t.bold("Image Model") : t.fg("dim", "Image Model");
-    const cl = this.activeTab === TAB_COMPACT ? t.bold("Compact Model") : t.fg("dim", "Compact Model");
-    this.tabTitle.setText(`${im} ${il}  |  ${cm} ${cl}`);
+    const bm = this.activeTab === TAB_BROKER ? t.fg("accent", "●") : "○";
+    const il = this.activeTab === TAB_IMAGE ? t.bold("Image") : t.fg("dim", "Image");
+    const cl = this.activeTab === TAB_COMPACT ? t.bold("Compact") : t.fg("dim", "Compact");
+    const bl = this.activeTab === TAB_BROKER ? t.bold("Broker") : t.fg("dim", "Broker");
+    this.tabTitle.setText(`${im} ${il}  |  ${cm} ${cl}  |  ${bm} ${bl}`);
     const key = this.currentKey();
     this.subtitleText.setText(key ? t.fg("warning", `Current ${this.currentKind()} model: ${key}`) : t.fg("warning", `No ${this.currentKind()} model set`));
   }
 
   handleInput(keyData: string) {
     const kb = getKeybindings();
-    if (kb.matches(keyData, "tui.input.tab")) { this.switchTab(this.activeTab === TAB_IMAGE ? TAB_COMPACT : TAB_IMAGE); this.tui.requestRender(); return; }
+    if (kb.matches(keyData, "tui.input.tab")) {
+      const next = this.activeTab === TAB_IMAGE ? TAB_COMPACT : this.activeTab === TAB_COMPACT ? TAB_BROKER : TAB_IMAGE;
+      this.switchTab(next); this.tui.requestRender(); return;
+    }
     if (kb.matches(keyData, "tui.select.up")) { this.selectedIndex = this.selectedIndex === 0 ? this.filtered.length - 1 : this.selectedIndex - 1; this.updateList(); return; }
     if (kb.matches(keyData, "tui.select.down")) { this.selectedIndex = this.selectedIndex === this.filtered.length - 1 ? 0 : this.selectedIndex + 1; this.updateList(); return; }
     if (kb.matches(keyData, "tui.select.confirm")) { const s = this.filtered[this.selectedIndex]; if (s) this.selectModel(s.model); return; }
@@ -289,12 +297,15 @@ export class ModelPickerComponent extends Container {
     const kind = this.currentKind();
     if (model) {
       if (kind === "image") setImageModelKey(formatModelKey(model));
+      else if (kind === "broker") setMcpBrokerModelKey(formatModelKey(model));
       else setCompactModelKey(formatModelKey(model));
     } else {
       if (kind === "image") setImageModelKey(null);
+      else if (kind === "broker") setMcpBrokerModelKey(null);
       else setCompactModelKey(null);
     }
     if (kind === "image") this.imageKey = model ? formatModelKey(model) : null;
+    else if (kind === "broker") this.brokerKey = model ? formatModelKey(model) : null;
     else this.compactKey = model ? formatModelKey(model) : null;
     this.switchTab(this.activeTab); this.tui.requestRender();
   }