@ikuma.cloud/pix-mcp 0.0.7 → 0.0.9

package/src/index.ts

@@ -2,13 +2,39 @@ 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 { Prompt, Tool } from "@modelcontextprotocol/sdk/types.js";
+import type {
+  Prompt,
+  Resource,
+  ResourceTemplate,
+  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 { pickResource } from "./resource-picker.ts";
+import {
+  formatResourceList,
+  formatResourceResult,
+  parseResourceCommand,
+  resourceCompletions,
+  resourceEntryKey,
+  resourceKey,
+  resourceTemplateVariables,
+  resolveResourceUri,
+  type ResourceArgumentToken,
+  type ResourceCommand,
+  type ResourceEntry,
+  type TemplateResourceEntry,
+} from "./resources.ts";
 import {
   formatPromptList,
   formatPromptResult,
@@ -16,6 +42,8 @@ import {
   promptCompletions,
   promptKey,
   resolvePromptArguments,
+  type PromptArgumentToken,
+  type PromptCommand,
   type PromptEntry,
 } from "./prompts.ts";
 import { prepareSchema } from "./schema.ts";
@@ -31,15 +59,28 @@ interface Rejections {
   }[];
 }
 
+interface PromptSelection {
+  command: PromptCommand;
+  item: PromptEntry;
+}
+
+interface ResourceSelection {
+  command: ResourceCommand;
+  item: ResourceEntry;
+}
+
 interface State {
   alive: boolean;
   status: string;
   entries: Map<string, Entry>;
   prompts: Map<string, PromptEntry>;
+  resources: Map<string, ResourceEntry>;
   loaded: Map<string, string>;
   connections: Map<string, Connection>;
   rejected: Map<string, Rejections>;
   configIssues: ConfigIssue[];
+  draftCleanups: Set<() => Promise<void>>;
+  draftGuards: Set<string>;
 }
 
 const discoveryParameters = Type.Object(
@@ -73,7 +114,7 @@ export default function mcp(pi: ExtensionAPI) {
       "Read this MCP config file for this session (default: .mcp.json).",
   });
   pi.registerCommand("mcp-prompt", {
-    description: "List or run a user-selected MCP prompt",
+    description: "Select, list, or run a user-controlled MCP prompt",
     getArgumentCompletions: (prefix) =>
       promptCompletions(prefix, [...(state?.prompts.values() ?? [])]),
     async handler(input, ctx) {
@@ -81,7 +122,14 @@ export default function mcp(pi: ExtensionAPI) {
         const owner = state;
         if (!owner) throw new Error("MCP session has not started.");
         current(owner);
-        const command = parsePromptCommand(input);
+        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 &&
@@ -113,43 +161,319 @@ export default function mcp(pi: ExtensionAPI) {
           ctx.ui.notify(text, "info");
           return;
         }
-        const item = owner.prompts.get(promptKey(command.server, command.name));
+        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,
         );
-        const result = await connection.getPrompt(
-          item.prompt.name,
-          args,
-          ctx.signal,
-        );
+        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.draftCleanups.add(draft.cleanup);
+              if (draft.guard) owner.draftGuards.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");
+      }
+    },
+  });
+
+  pi.registerCommand("mcp-resource", {
+    description: "Select, list, or read a user-controlled MCP resource",
+    getArgumentCompletions: (prefix) =>
+      resourceCompletions(
+        prefix,
+        [...(state?.resources.values() ?? [])],
+        state?.connections.keys() ?? [],
+      ),
+    async handler(input, ctx) {
+      try {
+        const owner = state;
+        if (!owner) throw new Error("MCP session has not started.");
         current(owner);
-        const formatted = await formatPromptResult(
-          result,
-          item.server,
-          item.prompt.name,
+        const editBeforeSending = ctx.mode === "tui" && input.trim() === "";
+        const selection = editBeforeSending
+          ? await chooseResource(owner, ctx)
+          : undefined;
+        const command =
+          selection?.command ??
+          (editBeforeSending ? undefined : parseResourceCommand(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 resourceStatus = [
+            ...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.resourceStatus}`,
+              ),
+          ].join("\n");
+          const text = `${resourceStatus}${resourceStatus ? "\n\n" : ""}${formatResourceList(
+            [...owner.resources.values()],
+            command.server,
+          )}`;
+          if (!ctx.hasUI) throw new Error(text);
+          ctx.ui.notify(text, "info");
+          return;
+        }
+
+        const direct = owner.resources.get(
+          resourceKey(command.server, "resource", command.target),
         );
-        current(owner);
-        ctx.signal?.throwIfAborted();
-        pi.sendUserMessage(
-          formatted.content,
-          ctx.isIdle() ? undefined : { deliverAs: "followUp" },
+        const template = owner.resources.get(
+          resourceKey(command.server, "template", command.target),
+        );
+        const item =
+          selection?.item ??
+          (command.argumentTokens.length > 0 ? template : (direct ?? template));
+        const selectedKey = item ? resourceEntryKey(item) : undefined;
+        const ensureSelectedResourceCurrent = () => {
+          if (
+            selection &&
+            selectedKey &&
+            owner.resources.get(selectedKey) !== item
+          )
+            throw new Error(
+              "MCP resource catalog changed. Open the picker again.",
+            );
+        };
+        ensureSelectedResourceCurrent();
+        const templateEntry =
+          item?.kind === "template"
+            ? item
+            : template?.kind === "template"
+              ? template
+              : undefined;
+        const uri = resolveResourceUri(
+          command.target,
+          templateEntry,
+          command.argumentTokens,
         );
+        const connection = owner.connections.get(command.server);
+        if (!connection) throw new Error("MCP connection unavailable.");
+        if (ctx.mode === "tui")
+          ctx.ui.setStatus(
+            "mcp-resource",
+            `Loading ${command.server} / ${compact(uri, 120)}…`,
+          );
+        try {
+          const result = await connection.readResource(uri, ctx.signal);
+          current(owner);
+          ensureSelectedResourceCurrent();
+          const formatted = await formatResourceResult(command.server, result);
+          current(owner);
+          ensureSelectedResourceCurrent();
+          ctx.signal?.throwIfAborted();
+          if (editBeforeSending) {
+            const draft = await formatEditableContent(formatted.content);
+            try {
+              current(owner);
+              ensureSelectedResourceCurrent();
+              ctx.signal?.throwIfAborted();
+              ctx.ui.pasteToEditor(draft.text);
+              owner.draftCleanups.add(draft.cleanup);
+              if (draft.guard) owner.draftGuards.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-resource", undefined);
+        }
       } catch (error) {
         const message =
-          error instanceof Error ? error.message : "MCP prompt failed.";
+          error instanceof Error ? error.message : "MCP resource failed.";
         if (!ctx.hasUI) throw new Error(message);
         ctx.ui.notify(message, "error");
       }
     },
   });
 
+  async function chooseResource(
+    owner: State,
+    ctx: ExtensionCommandContext,
+  ): Promise<ResourceSelection | undefined> {
+    const entries = [...owner.resources.values()].sort(
+      (a, b) =>
+        a.server.localeCompare(b.server, "en") ||
+        resourceEntryKey(a).localeCompare(resourceEntryKey(b), "en"),
+    );
+    if (entries.length === 0) {
+      ctx.ui.notify("No MCP resources are available.", "warning");
+      return;
+    }
+
+    const item = await pickResource(entries, ctx);
+    if (!item) return;
+    ctx.signal?.throwIfAborted();
+
+    const key = resourceEntryKey(item);
+    const ensureCurrent = () => {
+      current(owner);
+      if (owner.resources.get(key) !== item)
+        throw new Error("MCP resource catalog changed. Open the picker again.");
+    };
+    ensureCurrent();
+
+    const dialogOptions = ctx.signal ? { signal: ctx.signal } : undefined;
+    const argumentTokens: ResourceArgumentToken[] = [];
+    if (item.kind === "template") {
+      for (const variable of item.variables) {
+        const value = await ctx.ui.input(
+          `${variable} (optional)`,
+          "Leave empty to omit",
+          dialogOptions,
+        );
+        if (value === undefined) return;
+        ctx.signal?.throwIfAborted();
+        ensureCurrent();
+        if (value !== "")
+          argumentTokens.push({
+            value: `${variable}=${value}`,
+            separator: variable.length,
+          });
+      }
+    }
+
+    return {
+      command: {
+        action: "read",
+        server: item.server,
+        target:
+          item.kind === "resource"
+            ? item.resource.uri
+            : item.template.uriTemplate,
+        argumentTokens,
+      },
+      item,
+    };
+  }
+
+  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);
     pi.setActiveTools(pi.getActiveTools().filter((name) => !removed.has(name)));
@@ -334,6 +658,35 @@ export default function mcp(pi: ExtensionAPI) {
     }
   }
 
+  function synchronizeResources(
+    owner: State,
+    config: ServerConfig,
+    resources: Resource[],
+    templates: ResourceTemplate[],
+  ) {
+    if (!owner.alive || state !== owner) return;
+    for (const [key, item] of owner.resources) {
+      if (item.server === config.name) owner.resources.delete(key);
+    }
+    for (const resource of resources) {
+      const item: ResourceEntry = {
+        kind: "resource",
+        server: config.name,
+        resource,
+      };
+      owner.resources.set(resourceEntryKey(item), item);
+    }
+    for (const template of templates) {
+      const item: TemplateResourceEntry = {
+        kind: "template",
+        server: config.name,
+        template,
+        variables: resourceTemplateVariables(template.uriTemplate),
+      };
+      owner.resources.set(resourceEntryKey(item), item);
+    }
+  }
+
   function synchronize(owner: State, config: ServerConfig, tools: Tool[]) {
     if (!owner.alive || state !== owner) return;
     const previous = new Map(
@@ -444,9 +797,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 draftCleanups = [...owner.draftCleanups];
+    owner.draftCleanups.clear();
+    owner.draftGuards.clear();
+    await Promise.all([
+      ...[...owner.connections.values()].map((connection) =>
+        connection.close(),
+      ),
+      ...draftCleanups.map((cleanup) => cleanup()),
+    ]);
   }
 
   async function start(ctx: ExtensionContext) {
@@ -456,10 +815,13 @@ export default function mcp(pi: ExtensionAPI) {
       status: "No MCP servers configured.",
       entries: new Map(),
       prompts: new Map(),
+      resources: new Map(),
       loaded: new Map(),
       connections: new Map(),
       rejected: new Map(),
       configIssues: [],
+      draftCleanups: new Set(),
+      draftGuards: new Set(),
     };
     state = owner;
     try {
@@ -491,6 +853,8 @@ export default function mcp(pi: ExtensionAPI) {
             server,
             (tools) => synchronize(owner, server, tools),
             (prompts) => synchronizePrompts(owner, server, prompts),
+            (resources, templates) =>
+              synchronizeResources(owner, server, resources, templates),
           ),
         );
       }
@@ -514,6 +878,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.draftGuards) {
+        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.draftGuards.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,4 +1,5 @@
-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";
@@ -28,6 +29,7 @@ interface ContentOptions {
   structuredContent?: unknown;
   preserveErrorImages?: boolean;
   preserveOrder?: boolean;
+  artifactRequired?: boolean;
 }
 
 export async function formatContent(
@@ -70,16 +72,19 @@ export async function formatContent(
       `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 = options.structuredContent;
   const largeDetails =
     structured !== undefined &&
     Buffer.byteLength(JSON.stringify(structured)) > MAX_DETAILS_BYTES;
+  const textTruncated = bounded !== fullText;
   const truncated =
-    bounded !== fullText ||
+    textTruncated ||
     omitted ||
     largeDetails ||
+    options.artifactRequired === true ||
     (options.preserveErrorImages === true && images.length > 0);
   let fullOutputPath: string | undefined;
   if (truncated) {
@@ -99,7 +104,7 @@ export async function formatContent(
     let textIndex = 0;
     for (const block of ordered) {
       if (block.type === "image") {
-        content.push(block);
+        if (!textTruncated || offset < bounded.length) content.push(block);
         continue;
       }
       const value = `${textIndex++ > 0 ? "\n\n" : ""}${block.text}`;
@@ -112,7 +117,8 @@ export async function formatContent(
       if (previous?.type === "text") previous.text += visible;
       else content.push({ type: "text", text: visible });
     }
-    if (textIndex === 0) content.unshift({ type: "text", text: bounded });
+    if (textIndex === 0 && bounded)
+      content.unshift({ type: "text", text: bounded });
     if (artifactNotice)
       content.push({ type: "text", text: `\n\n${artifactNotice}` });
   } else {
@@ -137,6 +143,80 @@ export async function formatContent(
   };
 }
 
+function createEditableContentGuard(): 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 prepareEditableContent(text: string): {
+  text: string;
+  guard?: string;
+} {
+  // The TUI interprets leading / and ! as commands. An invisible guard keeps
+  // untrusted staged 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 = createEditableContentGuard();
+  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-draft-"));
+      const extension = imageExtensions[block.mimeType] ?? "img";
+      const path = join(imageDirectory, `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 {
+    ...prepareEditableContent(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.