@ikuma.cloud/pix-mcp 0.0.6 → 0.0.8

package/src/index.ts

@@ -2,13 +2,30 @@ import { resolve } from "node:path";
 import { Type } from "@earendil-works/pi-ai";
 import type {
   ExtensionAPI,
+  ExtensionCommandContext,
   ExtensionContext,
 } from "@earendil-works/pi-coding-agent";
-import type { Tool } from "@modelcontextprotocol/sdk/types.js";
+import type { Prompt, Tool } from "@modelcontextprotocol/sdk/types.js";
 import { compact, entry, search, summary, type Entry } from "./catalog.ts";
 import { Connection } from "./client.ts";
 import { readConfig, type ConfigIssue, type ServerConfig } from "./config.ts";
-import { formatResult } from "./output.ts";
+import {
+  formatEditableContent,
+  formatResult,
+  removeEditablePromptGuard,
+} from "./output.ts";
+import { pickPrompt } from "./prompt-picker.ts";
+import {
+  formatPromptList,
+  formatPromptResult,
+  parsePromptCommand,
+  promptCompletions,
+  promptKey,
+  resolvePromptArguments,
+  type PromptArgumentToken,
+  type PromptCommand,
+  type PromptEntry,
+} from "./prompts.ts";
 import { prepareSchema } from "./schema.ts";
 import { ToolRejectionError, type RejectionCode } from "./rejection.ts";
 
@@ -22,14 +39,22 @@ interface Rejections {
   }[];
 }
 
+interface PromptSelection {
+  command: PromptCommand;
+  item: PromptEntry;
+}
+
 interface State {
   alive: boolean;
   status: string;
   entries: Map<string, Entry>;
+  prompts: Map<string, PromptEntry>;
   loaded: Map<string, string>;
   connections: Map<string, Connection>;
   rejected: Map<string, Rejections>;
   configIssues: ConfigIssue[];
+  promptDraftCleanups: Set<() => Promise<void>>;
+  promptDraftGuards: Set<string>;
 }
 
 const discoveryParameters = Type.Object(
@@ -62,6 +87,175 @@ export default function mcp(pi: ExtensionAPI) {
     description:
       "Read this MCP config file for this session (default: .mcp.json).",
   });
+  pi.registerCommand("mcp-prompt", {
+    description: "Select, list, or run a user-controlled MCP prompt",
+    getArgumentCompletions: (prefix) =>
+      promptCompletions(prefix, [...(state?.prompts.values() ?? [])]),
+    async handler(input, ctx) {
+      try {
+        const owner = state;
+        if (!owner) throw new Error("MCP session has not started.");
+        current(owner);
+        const editBeforeSending = ctx.mode === "tui" && input.trim() === "";
+        const selection = editBeforeSending
+          ? await choosePrompt(owner, ctx)
+          : undefined;
+        const command =
+          selection?.command ??
+          (editBeforeSending ? undefined : parsePromptCommand(input));
+        if (!command) return;
+        if (command.action === "list") {
+          if (
+            command.server &&
+            !owner.connections.has(command.server) &&
+            !owner.configIssues.some((issue) => issue.name === command.server)
+          )
+            throw new Error("Unknown MCP server.");
+          const promptStatus = [
+            ...owner.configIssues
+              .filter(
+                (issue) => !command.server || issue.name === command.server,
+              )
+              .map((issue) => `${issue.name}: ${issue.message}`),
+            ...[...owner.connections.values()]
+              .filter(
+                (connection) =>
+                  !command.server || connection.config.name === command.server,
+              )
+              .map(
+                (connection) =>
+                  `${connection.config.name}: ${connection.promptStatus}`,
+              ),
+          ].join("\n");
+          const text = `${promptStatus}${promptStatus ? "\n\n" : ""}${formatPromptList(
+            [...owner.prompts.values()],
+            command.server,
+          )}`;
+          if (!ctx.hasUI) throw new Error(text);
+          ctx.ui.notify(text, "info");
+          return;
+        }
+        const key = promptKey(command.server, command.name);
+        const item = selection?.item ?? owner.prompts.get(key);
+        if (!item)
+          throw new Error(
+            "Unknown or unavailable MCP prompt. List prompts again.",
+          );
+        const ensureSelectedPromptCurrent = () => {
+          if (selection && owner.prompts.get(key) !== item)
+            throw new Error(
+              "MCP prompt catalog changed. Open the picker again.",
+            );
+        };
+        ensureSelectedPromptCurrent();
+        const connection = owner.connections.get(command.server);
+        if (!connection) throw new Error("MCP connection unavailable.");
+        const args = resolvePromptArguments(
+          item.prompt,
+          command.argumentTokens,
+        );
+        if (ctx.mode === "tui")
+          ctx.ui.setStatus(
+            "mcp-prompt",
+            `Loading ${item.server} / ${item.prompt.name}…`,
+          );
+        try {
+          const result = await connection.getPrompt(
+            item.prompt.name,
+            args,
+            ctx.signal,
+          );
+          current(owner);
+          ensureSelectedPromptCurrent();
+          const formatted = await formatPromptResult(result);
+          current(owner);
+          ensureSelectedPromptCurrent();
+          ctx.signal?.throwIfAborted();
+          if (editBeforeSending) {
+            const draft = await formatEditableContent(formatted.content);
+            try {
+              current(owner);
+              ensureSelectedPromptCurrent();
+              ctx.signal?.throwIfAborted();
+              ctx.ui.pasteToEditor(draft.text);
+              owner.promptDraftCleanups.add(draft.cleanup);
+              if (draft.guard) owner.promptDraftGuards.add(draft.guard);
+            } catch (error) {
+              await draft.cleanup();
+              throw error;
+            }
+          } else {
+            pi.sendUserMessage(
+              formatted.content,
+              ctx.isIdle() ? undefined : { deliverAs: "followUp" },
+            );
+          }
+        } finally {
+          if (ctx.mode === "tui") ctx.ui.setStatus("mcp-prompt", undefined);
+        }
+      } catch (error) {
+        const message =
+          error instanceof Error ? error.message : "MCP prompt failed.";
+        if (!ctx.hasUI) throw new Error(message);
+        ctx.ui.notify(message, "error");
+      }
+    },
+  });
+
+  async function choosePrompt(
+    owner: State,
+    ctx: ExtensionCommandContext,
+  ): Promise<PromptSelection | undefined> {
+    const entries = [...owner.prompts.values()].sort(
+      (a, b) =>
+        a.server.localeCompare(b.server, "en") ||
+        a.prompt.name.localeCompare(b.prompt.name, "en"),
+    );
+    if (entries.length === 0) {
+      ctx.ui.notify("No MCP prompts are available.", "warning");
+      return;
+    }
+
+    const item = await pickPrompt(entries, ctx);
+    if (!item) return;
+    ctx.signal?.throwIfAborted();
+
+    const key = promptKey(item.server, item.prompt.name);
+    const ensureCurrent = () => {
+      current(owner);
+      if (owner.prompts.get(key) !== item)
+        throw new Error("MCP prompt catalog changed. Open the picker again.");
+    };
+    ensureCurrent();
+
+    const dialogOptions = ctx.signal ? { signal: ctx.signal } : undefined;
+    const argumentTokens: PromptArgumentToken[] = [];
+    for (const argument of item.prompt.arguments ?? []) {
+      const value = await ctx.ui.input(
+        `${argument.name} (${argument.required ? "required" : "optional"})`,
+        argument.required ? "Enter a value" : "Leave empty to omit",
+        dialogOptions,
+      );
+      if (value === undefined) return;
+      ctx.signal?.throwIfAborted();
+      ensureCurrent();
+      if (argument.required || value !== "")
+        argumentTokens.push({
+          value: `${argument.name}=${value}`,
+          separator: argument.name.length,
+        });
+    }
+
+    return {
+      command: {
+        action: "run",
+        server: item.server,
+        name: item.prompt.name,
+        argumentTokens,
+      },
+      item,
+    };
+  }
 
   function deactivate(names: Iterable<string>) {
     const removed = new Set(names);
@@ -230,6 +424,23 @@ export default function mcp(pi: ExtensionAPI) {
     if (active) pi.setActiveTools(active);
   }
 
+  function synchronizePrompts(
+    owner: State,
+    config: ServerConfig,
+    prompts: Prompt[],
+  ) {
+    if (!owner.alive || state !== owner) return;
+    for (const [key, item] of owner.prompts) {
+      if (item.server === config.name) owner.prompts.delete(key);
+    }
+    for (const prompt of prompts) {
+      owner.prompts.set(promptKey(config.name, prompt.name), {
+        server: config.name,
+        prompt,
+      });
+    }
+  }
+
   function synchronize(owner: State, config: ServerConfig, tools: Tool[]) {
     if (!owner.alive || state !== owner) return;
     const previous = new Map(
@@ -340,9 +551,15 @@ export default function mcp(pi: ExtensionAPI) {
     owner.alive = false;
     state = undefined;
     deactivate(owned);
-    await Promise.all(
-      [...owner.connections.values()].map((connection) => connection.close()),
-    );
+    const promptDraftCleanups = [...owner.promptDraftCleanups];
+    owner.promptDraftCleanups.clear();
+    owner.promptDraftGuards.clear();
+    await Promise.all([
+      ...[...owner.connections.values()].map((connection) =>
+        connection.close(),
+      ),
+      ...promptDraftCleanups.map((cleanup) => cleanup()),
+    ]);
   }
 
   async function start(ctx: ExtensionContext) {
@@ -351,10 +568,13 @@ export default function mcp(pi: ExtensionAPI) {
       alive: true,
       status: "No MCP servers configured.",
       entries: new Map(),
+      prompts: new Map(),
       loaded: new Map(),
       connections: new Map(),
       rejected: new Map(),
       configIssues: [],
+      promptDraftCleanups: new Set(),
+      promptDraftGuards: new Set(),
     };
     state = owner;
     try {
@@ -382,7 +602,11 @@ export default function mcp(pi: ExtensionAPI) {
       for (const server of config.servers) {
         owner.connections.set(
           server.name,
-          new Connection(server, (tools) => synchronize(owner, server, tools)),
+          new Connection(
+            server,
+            (tools) => synchronize(owner, server, tools),
+            (prompts) => synchronizePrompts(owner, server, prompts),
+          ),
         );
       }
       // Bound startup concurrency without letting one broken server hide healthy ones.
@@ -405,6 +629,35 @@ export default function mcp(pi: ExtensionAPI) {
     }
   }
 
+  // Keep each invisible editor guard through Pi's command/template dispatch,
+  // then remove only guards created by this session before model context.
+  pi.on("message_end", (event) => {
+    const owner = state;
+    if (!owner || event.message.role !== "user") return;
+    const consumed = new Set<string>();
+    const removeKnownGuards = (text: string) => {
+      let next = text;
+      for (const guard of owner.promptDraftGuards) {
+        if (!next.includes(guard)) continue;
+        next = removeEditablePromptGuard(next, guard);
+        consumed.add(guard);
+      }
+      return next;
+    };
+    let content: typeof event.message.content;
+    if (typeof event.message.content === "string") {
+      content = removeKnownGuards(event.message.content);
+    } else {
+      content = event.message.content.map((block) =>
+        block.type === "text"
+          ? { ...block, text: removeKnownGuards(block.text) }
+          : block,
+      );
+    }
+    for (const guard of consumed) owner.promptDraftGuards.delete(guard);
+    if (consumed.size > 0) return { message: { ...event.message, content } };
+  });
+
   // Any extension's registerTool can reactivate allowlisted tools in Pi 0.87.
   // Independently track loaded fingerprints, enforce them in execute, and scrub
   // request-local declarations at the last provider-neutral context boundary.

package/src/output.ts

@@ -1,8 +1,12 @@
-import { mkdtemp, writeFile } from "node:fs/promises";
+import { randomBytes } from "node:crypto";
+import { mkdtemp, rm, writeFile } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
-import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type {
+  CallToolResult,
+  ContentBlock,
+} from "@modelcontextprotocol/sdk/types.js";
 
 export const MAX_TEXT_BYTES = 24 * 1024;
 const MAX_IMAGE_BYTES = 4 * 1024 * 1024;
@@ -20,12 +24,27 @@ function preview(text: string): string {
     .join("\n");
 }
 
-export async function formatResult(result: CallToolResult) {
+interface ContentOptions {
+  artifact: unknown;
+  structuredContent?: unknown;
+  preserveErrorImages?: boolean;
+  preserveOrder?: boolean;
+}
+
+export async function formatContent(
+  blocks: ContentBlock[],
+  options: ContentOptions,
+) {
   const text: string[] = [];
   const images: ImageContent[] = [];
+  const ordered: (TextContent | ImageContent)[] = [];
+  const addText = (value: string) => {
+    text.push(value);
+    ordered.push({ type: "text", text: value });
+  };
   let omitted = false;
-  for (const block of result.content) {
-    if (block.type === "text") text.push(block.text);
+  for (const block of blocks) {
+    if (block.type === "text") addText(block.text);
     else if (
       block.type === "image" &&
       images.length < MAX_IMAGES &&
@@ -35,53 +54,80 @@ export async function formatResult(result: CallToolResult) {
       // Padding can give three decoded sizes the same encoded length.
       Buffer.byteLength(block.data, "base64") <= MAX_IMAGE_BYTES
     ) {
-      images.push({
+      const image: ImageContent = {
         type: "image",
         data: block.data,
         mimeType: block.mimeType,
-      });
+      };
+      images.push(image);
+      ordered.push(image);
     } else {
-      text.push(`[MCP ${block.type} content omitted; see full result.]`);
+      addText(`[MCP ${block.type} content omitted; see full result.]`);
       omitted = true;
     }
   }
-  if (result.structuredContent !== undefined) {
-    text.push(
-      `Structured content:\n${JSON.stringify(result.structuredContent)}`,
+  if (options.structuredContent !== undefined) {
+    addText(
+      `Structured content:\n${JSON.stringify(options.structuredContent)}`,
     );
   }
-  const fullText = text.join("\n\n") || "(No text output)";
+  const fullText =
+    text.join("\n\n") || (options.preserveOrder ? "" : "(No text output)");
   const bounded = preview(fullText);
-  const structured = result.structuredContent;
+  const structured = options.structuredContent;
   const largeDetails =
     structured !== undefined &&
     Buffer.byteLength(JSON.stringify(structured)) > MAX_DETAILS_BYTES;
-  // Pi 0.87 turns thrown tool errors into text only. Preserve error images in
-  // the full-result artifact before the native execution wrapper throws.
+  const textTruncated = bounded !== fullText;
   const truncated =
-    bounded !== fullText ||
+    textTruncated ||
     omitted ||
     largeDetails ||
-    (result.isError === true && images.length > 0);
+    (options.preserveErrorImages === true && images.length > 0);
   let fullOutputPath: string | undefined;
   if (truncated) {
     const directory = await mkdtemp(join(tmpdir(), "pix-mcp-"));
     fullOutputPath = join(directory, "result.json");
-    await writeFile(fullOutputPath, JSON.stringify(result, null, 2), {
+    await writeFile(fullOutputPath, JSON.stringify(options.artifact, null, 2), {
       mode: 0o600,
     });
   }
-  const content: (TextContent | ImageContent)[] = [
-    {
-      type: "text",
-      text:
-        bounded +
-        (fullOutputPath
-          ? `\n\nFull MCP result: ${fullOutputPath}\nUse read with offset/limit to inspect it.`
-          : ""),
-    },
-    ...images,
-  ];
+  const artifactNotice = fullOutputPath
+    ? `Full MCP result: ${fullOutputPath}\nUse read with offset/limit to inspect it.`
+    : undefined;
+  let content: (TextContent | ImageContent)[];
+  if (options.preserveOrder) {
+    content = [];
+    let offset = 0;
+    let textIndex = 0;
+    for (const block of ordered) {
+      if (block.type === "image") {
+        if (!textTruncated || offset < bounded.length) content.push(block);
+        continue;
+      }
+      const value = `${textIndex++ > 0 ? "\n\n" : ""}${block.text}`;
+      // Slice the already byte/line-bounded preview itself; value supplies only
+      // the original block boundary needed to preserve image ordering.
+      const visible = bounded.slice(offset, offset + value.length);
+      offset += visible.length;
+      if (!visible) continue;
+      const previous = content.at(-1);
+      if (previous?.type === "text") previous.text += visible;
+      else content.push({ type: "text", text: visible });
+    }
+    if (textIndex === 0 && bounded)
+      content.unshift({ type: "text", text: bounded });
+    if (artifactNotice)
+      content.push({ type: "text", text: `\n\n${artifactNotice}` });
+  } else {
+    content = [
+      {
+        type: "text",
+        text: bounded + (artifactNotice ? `\n\n${artifactNotice}` : ""),
+      },
+      ...images,
+    ];
+  }
   return {
     content,
     details: {
@@ -94,3 +140,90 @@ export async function formatResult(result: CallToolResult) {
     },
   };
 }
+
+function createEditablePromptGuard(): string {
+  const bytes = randomBytes(8);
+  let guard = "\u2060\u2063";
+  for (const byte of bytes)
+    for (let bit = 7; bit >= 0; bit--)
+      guard += byte & (1 << bit) ? "\u2063" : "\u2060";
+  return guard;
+}
+
+export function removeEditablePromptGuard(text: string, guard: string): string {
+  return text.replaceAll(guard, "");
+}
+
+function prepareEditablePrompt(text: string): {
+  text: string;
+  guard?: string;
+} {
+  // The TUI interprets leading / and ! as commands. An invisible guard keeps
+  // untrusted prompt text inert in the editor and is removed at message_end.
+  const sanitized = text
+    .replace(/\r\n?/gu, "\n")
+    .replace(/\p{Cc}/gu, (character) =>
+      character === "\n" || character === "\t" ? character : "",
+    )
+    .replace(/[\u061C\u200E\u200F\u202A-\u202E\u2066-\u2069]/gu, "");
+  if (!/^[!/]/u.test(sanitized.trimStart())) return { text: sanitized };
+  const guard = createEditablePromptGuard();
+  return { text: `${guard}${sanitized}`, guard };
+}
+
+const imageExtensions: Record<string, string> = {
+  "image/png": "png",
+  "image/jpeg": "jpg",
+  "image/gif": "gif",
+  "image/webp": "webp",
+};
+
+export async function formatEditableContent(
+  content: (TextContent | ImageContent)[],
+): Promise<{ text: string; guard?: string; cleanup: () => Promise<void> }> {
+  let text = "";
+  let imageDirectory: string | undefined;
+  let imageIndex = 0;
+  try {
+    for (const block of content) {
+      if (block.type === "text") {
+        text += block.text;
+        continue;
+      }
+      imageDirectory ??= await mkdtemp(join(tmpdir(), "pix-mcp-prompt-"));
+      const extension = imageExtensions[block.mimeType] ?? "img";
+      const path = join(
+        imageDirectory,
+        `prompt-image-${++imageIndex}.${extension}`,
+      );
+      await writeFile(path, Buffer.from(block.data, "base64"), { mode: 0o600 });
+      const separator = text.endsWith("\n\n")
+        ? ""
+        : text.endsWith("\n")
+          ? "\n"
+          : "\n\n";
+      text += `${separator}@${JSON.stringify(path.replaceAll("\\", "/"))}`;
+    }
+  } catch (error) {
+    if (imageDirectory)
+      await rm(imageDirectory, { recursive: true, force: true });
+    throw error;
+  }
+  return {
+    ...prepareEditablePrompt(text),
+    cleanup: async () => {
+      if (imageDirectory)
+        await rm(imageDirectory, { recursive: true, force: true });
+    },
+  };
+}
+
+export function formatResult(result: CallToolResult) {
+  // Pi 0.87 turns thrown tool errors into text only. Preserve error images in
+  // the full-result artifact before the native execution wrapper throws.
+  return formatContent(result.content, {
+    artifact: result,
+    structuredContent: result.structuredContent,
+    preserveErrorImages: result.isError === true,
+  });
+}

package/src/prompt-picker.ts

@@ -0,0 +1,120 @@
+import {
+  DynamicBorder,
+  type ExtensionCommandContext,
+  type Theme,
+} from "@earendil-works/pi-coding-agent";
+import {
+  Container,
+  SelectList,
+  Text,
+  type Component,
+  type SelectItem,
+  type TUI,
+  type TuiMouseEvent,
+} from "@earendil-works/pi-tui";
+import { compact } from "./catalog.ts";
+import { promptKey, type PromptEntry } from "./prompts.ts";
+
+function details(entry: PromptEntry, theme: Theme): string {
+  const title = compact(entry.prompt.title ?? "", 120);
+  const description = compact(entry.prompt.description ?? "", 320);
+  return [
+    ...(title ? [theme.fg("accent", theme.bold(title))] : []),
+    theme.fg(
+      description ? "muted" : "dim",
+      description || "No description provided.",
+    ),
+  ].join("\n");
+}
+
+export function createPromptPicker(
+  entries: PromptEntry[],
+  tui: TUI,
+  theme: Theme,
+  done: (result: string | undefined) => void,
+  signal?: AbortSignal,
+): Component & { dispose(): void } {
+  const items: SelectItem[] = entries.map((entry) => ({
+    value: promptKey(entry.server, entry.prompt.name),
+    label: `${entry.server} / ${entry.prompt.name}`,
+  }));
+  const entriesByKey = new Map(
+    entries.map((entry) => [promptKey(entry.server, entry.prompt.name), entry]),
+  );
+  const container = new Container();
+  container.addChild(new DynamicBorder((text) => theme.fg("accent", text)));
+  container.addChild(
+    new Text(theme.fg("accent", theme.bold("Select MCP prompt"))),
+  );
+
+  const list = new SelectList(items, Math.min(items.length, 10), {
+    selectedPrefix: (text) => theme.fg("accent", text),
+    selectedText: (text) => theme.fg("accent", text),
+    description: (text) => theme.fg("muted", text),
+    scrollInfo: (text) => theme.fg("dim", text),
+    noMatch: (text) => theme.fg("warning", text),
+  });
+  container.addChild(list);
+
+  const focusedDescription = new Text("", 1, 1);
+  const updateDescription = (item: SelectItem) => {
+    const entry = entriesByKey.get(item.value);
+    if (entry) focusedDescription.setText(details(entry, theme));
+  };
+  const initialItem = items[0];
+  if (initialItem) updateDescription(initialItem);
+  list.onSelectionChange = updateDescription;
+  container.addChild(focusedDescription);
+  container.addChild(
+    new Text(theme.fg("dim", "↑↓ navigate • enter select • esc cancel")),
+  );
+  container.addChild(new DynamicBorder((text) => theme.fg("accent", text)));
+
+  let finished = false;
+  const finish = (result: string | undefined) => {
+    if (finished) return;
+    finished = true;
+    done(result);
+  };
+  list.onSelect = (item) => finish(item.value);
+  list.onCancel = () => finish(undefined);
+  const onAbort = () => finish(undefined);
+  signal?.addEventListener("abort", onAbort, { once: true });
+  if (signal?.aborted) onAbort();
+
+  return {
+    render(width: number) {
+      return container.render(width);
+    },
+    invalidate() {
+      container.invalidate();
+    },
+    handleInput(data: string) {
+      list.handleInput(data);
+      tui.requestRender();
+    },
+    handleMouse(event: TuiMouseEvent) {
+      const result = container.handleMouse(event);
+      if (result?.render) tui.requestRender();
+      return result;
+    },
+    dispose() {
+      signal?.removeEventListener("abort", onAbort);
+    },
+  };
+}
+
+export async function pickPrompt(
+  entries: PromptEntry[],
+  ctx: ExtensionCommandContext,
+): Promise<PromptEntry | undefined> {
+  if (ctx.signal?.aborted) return;
+  const selected = await ctx.ui.custom<string | undefined>(
+    (tui, theme, _keybindings, done) =>
+      createPromptPicker(entries, tui, theme, done, ctx.signal),
+  );
+  if (selected === undefined) return;
+  return entries.find(
+    (entry) => promptKey(entry.server, entry.prompt.name) === selected,
+  );
+}