@docyrus/docyrus 0.0.20 → 0.0.21

package/resources/pi-agent/extensions/pi-mcp-adapter/ui-session.ts

@@ -0,0 +1,384 @@
+import { randomUUID } from "node:crypto";
+import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import type { McpExtensionState } from "./state.js";
+import {
+  extractUiPromptText,
+  UI_STREAM_HOST_CONTEXT_KEY,
+  UI_STREAM_REQUEST_META_KEY,
+  UI_STREAM_STRUCTURED_CONTENT_KEY,
+  type UiHostContext,
+  type UiMessageParams,
+  type UiModelContextParams,
+  type UiStreamMode,
+} from "./types.js";
+import { logger } from "./logger.js";
+import { startUiServer, type UiServerHandle } from "./ui-server.js";
+import { isGlimpseAvailable, openGlimpseWindow } from "./glimpse-ui.js";
+
+let activeGlimpseWindow: { close(): void } | null = null;
+
+export interface UiSessionRequest {
+  serverName: string;
+  toolName: string;
+  toolArgs: Record<string, unknown>;
+  uiResourceUri: string;
+  streamMode?: UiStreamMode;
+}
+
+export interface UiSessionRuntime {
+  serverName: string;
+  toolName: string;
+  reused: boolean;
+  streamId?: string;
+  streamToken?: string;
+  streamMode?: UiStreamMode;
+  requestMeta?: Record<string, unknown>;
+  url: string;
+  isActive: () => boolean;
+  sendToolResult: (result: CallToolResult) => void;
+  sendResultPatch: (result: CallToolResult) => void;
+  sendToolCancelled: (reason: string) => void;
+  close: (reason?: string) => void;
+}
+
+const MAX_COMPLETED_SESSIONS = 10;
+
+function withStreamEnvelope(
+  result: CallToolResult,
+  streamId: string | undefined,
+  sequence: number,
+): CallToolResult {
+  if (!streamId) {
+    return result;
+  }
+
+  const structuredContent = result.structuredContent && typeof result.structuredContent === "object" && !Array.isArray(result.structuredContent)
+    ? { ...result.structuredContent }
+    : {};
+
+  const rawEnvelope = structuredContent[UI_STREAM_STRUCTURED_CONTENT_KEY];
+  const envelope = rawEnvelope && typeof rawEnvelope === "object" && !Array.isArray(rawEnvelope)
+    ? { ...rawEnvelope as Record<string, unknown> }
+    : {
+        frameType: "final",
+        phase: "settled",
+        status: result.isError ? "error" : "ok",
+      };
+
+  structuredContent[UI_STREAM_STRUCTURED_CONTENT_KEY] = {
+    ...envelope,
+    streamId,
+    sequence,
+  };
+
+  return {
+    ...result,
+    structuredContent,
+  };
+}
+
+async function openInBrowser(state: McpExtensionState, url: string): Promise<void> {
+  try {
+    await state.openBrowser(url);
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    state.ui?.notify(`MCP UI browser open failed: ${message}`, "warning");
+    state.ui?.notify(`Open manually: ${url}`, "info");
+  }
+}
+
+export async function maybeStartUiSession(
+  state: McpExtensionState,
+  request: UiSessionRequest,
+): Promise<UiSessionRuntime | null> {
+  const log = logger.child({
+    component: "UiSession",
+    server: request.serverName,
+    tool: request.toolName,
+  });
+
+  try {
+    if (
+      state.uiServer &&
+      state.uiServer.serverName === request.serverName &&
+      state.uiServer.toolName === request.toolName
+    ) {
+      const existingHandle = state.uiServer;
+      const streamMode = request.streamMode;
+      const streamId = streamMode ? randomUUID() : undefined;
+      const streamToken = streamMode ? randomUUID() : undefined;
+      let active = true;
+      let nextStreamSequence = 0;
+
+      const cleanupStreamListener = () => {
+        if (streamToken) {
+          state.manager.removeUiStreamListener(streamToken);
+        }
+      };
+
+      existingHandle.sendToolInput(request.toolArgs);
+
+      if (streamToken) {
+        state.manager.registerUiStreamListener(streamToken, (serverName, notification) => {
+          if (!active || state.uiServer !== existingHandle) return;
+          if (serverName !== request.serverName) return;
+          nextStreamSequence += 1;
+          existingHandle.sendResultPatch(
+            withStreamEnvelope(notification.result as CallToolResult, streamId, nextStreamSequence),
+          );
+        });
+      }
+
+      return {
+        serverName: request.serverName,
+        toolName: request.toolName,
+        reused: true,
+        streamId,
+        streamToken,
+        streamMode,
+        requestMeta: streamToken ? { [UI_STREAM_REQUEST_META_KEY]: streamToken } : undefined,
+        url: existingHandle.url,
+        isActive: () => active && state.uiServer === existingHandle,
+        sendToolResult: (result: CallToolResult) => {
+          if (!active || state.uiServer !== existingHandle) return;
+          nextStreamSequence += 1;
+          existingHandle.sendToolResult(withStreamEnvelope(result, streamId, nextStreamSequence));
+        },
+        sendResultPatch: (result: CallToolResult) => {
+          if (!active || state.uiServer !== existingHandle) return;
+          nextStreamSequence += 1;
+          existingHandle.sendResultPatch(withStreamEnvelope(result, streamId, nextStreamSequence));
+        },
+        sendToolCancelled: (reason: string) => {
+          if (!active || state.uiServer !== existingHandle) return;
+          nextStreamSequence += 1;
+          existingHandle.sendToolResult(
+            withStreamEnvelope(
+              {
+                isError: true,
+                content: [{ type: "text", text: reason }],
+              },
+              streamId,
+              nextStreamSequence,
+            ),
+          );
+        },
+        close: () => {
+          active = false;
+          cleanupStreamListener();
+        },
+      };
+    }
+
+    const resource = await state.uiResourceHandler.readUiResource(request.serverName, request.uiResourceUri);
+
+    if (state.uiServer) {
+      state.uiServer.close("replaced");
+      state.uiServer = null;
+    }
+    if (activeGlimpseWindow) {
+      activeGlimpseWindow.close();
+      activeGlimpseWindow = null;
+    }
+
+    const streamMode = request.streamMode;
+    const streamId = streamMode ? randomUUID() : undefined;
+    const streamToken = streamMode ? randomUUID() : undefined;
+    const hostContext: UiHostContext | undefined = streamMode && streamId
+      ? {
+          [UI_STREAM_HOST_CONTEXT_KEY]: {
+            mode: streamMode,
+            streamId,
+            intermediateResultPatches: streamMode === "stream-first",
+            partialInput: false,
+          },
+        }
+      : undefined;
+
+    let active = true;
+    let nextStreamSequence = 0;
+    let handle: UiServerHandle | null = null;
+
+    const cleanupStreamListener = () => {
+      if (streamToken) {
+        state.manager.removeUiStreamListener(streamToken);
+      }
+    };
+
+    handle = await startUiServer({
+      serverName: request.serverName,
+      toolName: request.toolName,
+      toolArgs: streamMode === "stream-first" ? {} : request.toolArgs,
+      resource,
+      manager: state.manager,
+      consentManager: state.consentManager,
+      hostContext,
+
+      onMessage: (params: UiMessageParams) => {
+        const prompt = extractUiPromptText(params);
+        if (prompt) {
+          if (state.sendMessage) {
+            state.sendMessage(
+              {
+                customType: "mcp-ui-prompt",
+                content: [{ type: "text", text: `User sent prompt from ${request.serverName} UI: "${prompt}"` }],
+                display: `💬 UI Prompt: ${prompt}`,
+                details: { server: request.serverName, tool: request.toolName, prompt },
+              },
+              { triggerTurn: true },
+            );
+            log.debug("Triggered agent turn for UI prompt", { prompt: prompt.slice(0, 50) });
+          }
+        } else if (params.type === "intent" || params.intent) {
+          const intent = params.intent ?? "";
+          const intentParams = params.params;
+          if (intent && state.sendMessage) {
+            const paramsStr = intentParams ? ` ${JSON.stringify(intentParams)}` : "";
+            state.sendMessage(
+              {
+                customType: "mcp-ui-intent",
+                content: [{ type: "text", text: `User triggered intent from ${request.serverName} UI: ${intent}${paramsStr}` }],
+                display: `🎯 UI Intent: ${intent}`,
+                details: { server: request.serverName, tool: request.toolName, intent, params: intentParams },
+              },
+              { triggerTurn: true },
+            );
+            log.debug("Triggered agent turn for UI intent", { intent });
+          }
+        } else if (params.type === "notify" || params.message) {
+          const text = params.message ?? "";
+          if (text && state.ui) {
+            state.ui.notify(`[${request.serverName}] ${text}`, "info");
+          }
+        }
+      },
+
+      onContextUpdate: (params: UiModelContextParams) => {
+        log.debug("Model context update from UI", {
+          hasContent: !!params.content,
+          hasStructured: !!params.structuredContent,
+        });
+      },
+
+      onComplete: (reason: string) => {
+        active = false;
+        cleanupStreamListener();
+
+        if (state.uiServer === handle) {
+          const messages = handle.getSessionMessages();
+          const stream = handle.getStreamSummary();
+          const hasContent =
+            messages.prompts.length > 0 ||
+            messages.intents.length > 0 ||
+            messages.notifications.length > 0 ||
+            !!stream;
+
+          if (hasContent) {
+            state.completedUiSessions.push({
+              serverName: handle.serverName,
+              toolName: handle.toolName,
+              completedAt: new Date(),
+              reason,
+              messages,
+              stream,
+            });
+
+            while (state.completedUiSessions.length > MAX_COMPLETED_SESSIONS) {
+              state.completedUiSessions.shift();
+            }
+
+            log.debug("Session completed", {
+              reason,
+              prompts: messages.prompts.length,
+              intents: messages.intents.length,
+              notifications: messages.notifications.length,
+              streamFrames: stream?.frames ?? 0,
+            });
+          }
+
+          state.uiServer = null;
+          if (activeGlimpseWindow) {
+            activeGlimpseWindow.close();
+            activeGlimpseWindow = null;
+          }
+        }
+      },
+    });
+
+    if (streamToken) {
+      state.manager.registerUiStreamListener(streamToken, (serverName, notification) => {
+        if (!active || state.uiServer !== handle) return;
+        if (serverName !== request.serverName) return;
+        nextStreamSequence += 1;
+        handle.sendResultPatch(withStreamEnvelope(notification.result as CallToolResult, streamId, nextStreamSequence));
+      });
+    }
+
+    state.uiServer = handle;
+
+    const glimpseDetected = isGlimpseAvailable();
+    const viewerPref = process.env.MCP_UI_VIEWER?.toLowerCase();
+    const useGlimpse = viewerPref === "glimpse" ||
+      (viewerPref !== "browser" && glimpseDetected);
+
+    if (useGlimpse) {
+      try {
+        const glimpseHtml = `<!DOCTYPE html><html><head><meta charset="utf-8"><style>body{margin:0;padding:0;width:100vw;height:100vh;overflow:hidden}iframe{width:100%;height:100%;border:none}</style></head><body><iframe src="${handle.url}"></iframe></body></html>`;
+        activeGlimpseWindow = await openGlimpseWindow(glimpseHtml, {
+          title: `MCP · ${request.serverName} · ${request.toolName}`,
+          width: 1000,
+          height: 800,
+          onClosed: () => {
+            if (active) handle.close("glimpse-closed");
+          },
+        });
+      } catch (error) {
+        log.debug("Glimpse unavailable, using browser", {
+          error: error instanceof Error ? error.message : String(error),
+        });
+        await openInBrowser(state, handle.url);
+      }
+    } else {
+      await openInBrowser(state, handle.url);
+    }
+
+    return {
+      serverName: request.serverName,
+      toolName: request.toolName,
+      reused: false,
+      streamId,
+      streamToken,
+      streamMode,
+      requestMeta: streamToken ? { [UI_STREAM_REQUEST_META_KEY]: streamToken } : undefined,
+      url: handle.url,
+      isActive: () => active && state.uiServer === handle,
+      sendToolResult: (result: CallToolResult) => {
+        if (!active || state.uiServer !== handle) return;
+        nextStreamSequence += 1;
+        handle.sendToolResult(withStreamEnvelope(result, streamId, nextStreamSequence));
+      },
+      sendResultPatch: (result: CallToolResult) => {
+        if (!active || state.uiServer !== handle) return;
+        nextStreamSequence += 1;
+        handle.sendResultPatch(withStreamEnvelope(result, streamId, nextStreamSequence));
+      },
+      sendToolCancelled: (reason: string) => {
+        if (!active || state.uiServer !== handle) return;
+        handle.sendToolCancelled(reason);
+      },
+      close: (reason?: string) => {
+        active = false;
+        cleanupStreamListener();
+        handle.close(reason);
+      },
+    };
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    log.error("Failed to start UI session", error instanceof Error ? error : undefined);
+    state.ui?.notify(
+      `MCP UI unavailable for ${request.toolName} (${request.serverName}): ${message}`,
+      "warning",
+    );
+    return null;
+  }
+}

package/resources/pi-agent/extensions/pi-mcp-adapter/ui-stream-types.ts

@@ -0,0 +1,89 @@
+import { z } from "zod";
+
+export const UI_STREAM_HOST_CONTEXT_KEY = "pi-mcp-adapter/stream";
+export const UI_STREAM_REQUEST_META_KEY = "pi-mcp-adapter/stream-token";
+export const UI_STREAM_RESULT_PATCH_METHOD = "notifications/pi-mcp-adapter/ui-result-patch";
+export const SERVER_STREAM_RESULT_PATCH_METHOD = "notifications/pi-mcp-adapter/result-patch";
+export const UI_STREAM_STRUCTURED_CONTENT_KEY = "pi-mcp-adapter/stream";
+
+export const uiStreamModeSchema = z.enum(["eager", "stream-first"]);
+export type UiStreamMode = z.infer<typeof uiStreamModeSchema>;
+
+export const visualizationStreamPhaseSchema = z.enum(["shell", "narrative", "structure", "detail", "settled"]);
+export type VisualizationStreamPhase = z.infer<typeof visualizationStreamPhaseSchema>;
+
+export const visualizationStreamFrameTypeSchema = z.enum(["patch", "checkpoint", "final"]);
+export type VisualizationStreamFrameType = z.infer<typeof visualizationStreamFrameTypeSchema>;
+
+export const visualizationStreamStatusSchema = z.enum(["ok", "error"]);
+export type VisualizationStreamStatus = z.infer<typeof visualizationStreamStatusSchema>;
+
+const looseRecordSchema = z.record(z.string(), z.unknown());
+const looseArraySchema = z.array(z.unknown());
+
+export const uiStreamHostContextSchema = z.object({
+  mode: uiStreamModeSchema,
+  streamId: z.string().min(1),
+  intermediateResultPatches: z.boolean(),
+  partialInput: z.boolean(),
+});
+export type UiStreamHostContext = z.infer<typeof uiStreamHostContextSchema>;
+
+export const visualizationStreamEnvelopeSchema = z.object({
+  streamId: z.string().min(1),
+  sequence: z.number().int().nonnegative(),
+  frameType: visualizationStreamFrameTypeSchema,
+  phase: visualizationStreamPhaseSchema,
+  status: visualizationStreamStatusSchema,
+  message: z.string().optional(),
+  spec: looseRecordSchema.optional(),
+  checkpoint: looseRecordSchema.optional(),
+});
+export type VisualizationStreamEnvelope = z.infer<typeof visualizationStreamEnvelopeSchema>;
+
+export const uiStreamCallToolResultSchema = z.object({
+  content: looseArraySchema.optional(),
+  structuredContent: looseRecordSchema.optional(),
+  isError: z.boolean().optional(),
+  _meta: looseRecordSchema.optional(),
+}).passthrough();
+export type UiStreamCallToolResult = z.infer<typeof uiStreamCallToolResultSchema>;
+
+export const uiStreamResultPatchNotificationSchema = z.object({
+  method: z.literal(UI_STREAM_RESULT_PATCH_METHOD),
+  params: uiStreamCallToolResultSchema,
+});
+export type UiStreamResultPatchNotification = z.infer<typeof uiStreamResultPatchNotificationSchema>;
+
+export const serverStreamResultPatchNotificationSchema = z.object({
+  method: z.literal(SERVER_STREAM_RESULT_PATCH_METHOD),
+  params: z.object({
+    streamToken: z.string().min(1),
+    result: uiStreamCallToolResultSchema,
+  }),
+});
+export type ServerStreamResultPatchNotification = z.infer<typeof serverStreamResultPatchNotificationSchema>;
+
+export interface UiStreamSummary {
+  streamId: string;
+  mode: UiStreamMode;
+  frames: number;
+  phases: VisualizationStreamPhase[];
+  finalStatus?: VisualizationStreamStatus;
+  lastMessage?: string;
+}
+
+export function getUiStreamHostContext(hostContext: Record<string, unknown> | undefined): UiStreamHostContext | undefined {
+  const candidate = hostContext?.[UI_STREAM_HOST_CONTEXT_KEY];
+  const parsed = uiStreamHostContextSchema.safeParse(candidate);
+  return parsed.success ? parsed.data : undefined;
+}
+
+export function getVisualizationStreamEnvelope(structuredContent: unknown): VisualizationStreamEnvelope | undefined {
+  if (!structuredContent || typeof structuredContent !== "object" || Array.isArray(structuredContent)) {
+    return undefined;
+  }
+  const candidate = (structuredContent as Record<string, unknown>)[UI_STREAM_STRUCTURED_CONTENT_KEY];
+  const parsed = visualizationStreamEnvelopeSchema.safeParse(candidate);
+  return parsed.success ? parsed.data : undefined;
+}

package/resources/pi-agent/extensions/pi-mcp-adapter/utils.ts

@@ -0,0 +1,75 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { platform } from "node:os";
+
+export async function openUrl(pi: ExtensionAPI, url: string, browser?: string): Promise<void> {
+  const os = platform();
+  let result;
+
+  if (os === "darwin") {
+    result = browser ? await pi.exec("open", ["-a", browser, url]) : await pi.exec("open", [url]);
+  } else if (os === "win32") {
+    result = browser
+      ? await pi.exec("cmd", ["/c", "start", "", browser, url])
+      : await pi.exec("cmd", ["/c", "start", "", url]);
+  } else {
+    result = browser ? await pi.exec(browser, [url]) : await pi.exec("xdg-open", [url]);
+  }
+
+  if (result.code !== 0) {
+    throw new Error(result.stderr || `Failed to open browser (exit code ${result.code})`);
+  }
+}
+
+export async function parallelLimit<T, R>(
+  items: T[],
+  limit: number,
+  fn: (item: T) => Promise<R>
+): Promise<R[]> {
+  const results: R[] = [];
+  let index = 0;
+
+  async function worker() {
+    while (index < items.length) {
+      const i = index++;
+      results[i] = await fn(items[i]);
+    }
+  }
+
+  const workers = Array(Math.min(limit, items.length)).fill(null).map(() => worker());
+  await Promise.all(workers);
+  return results;
+}
+
+export function getConfigPathFromArgv(): string | undefined {
+  const idx = process.argv.indexOf("--mcp-config");
+  if (idx >= 0 && idx + 1 < process.argv.length) {
+    return process.argv[idx + 1];
+  }
+  return undefined;
+}
+
+export function truncateAtWord(text: string, target: number): string {
+  if (!text || text.length <= target) return text;
+
+  const truncated = text.slice(0, target);
+  const lastSpace = truncated.lastIndexOf(" ");
+
+  if (lastSpace > target * 0.6) {
+    return truncated.slice(0, lastSpace) + "...";
+  }
+
+  return truncated + "...";
+}
+
+/**
+ * Extract the adapter-owned UI stream mode from tool metadata.
+ */
+export function extractToolUiStreamMode(toolMeta: Record<string, unknown> | undefined): "eager" | "stream-first" | undefined {
+  const uiMeta = toolMeta?.ui;
+  if (!uiMeta || typeof uiMeta !== "object") return undefined;
+  const streamMode = (uiMeta as Record<string, unknown>)["pi-mcp-adapter.streamMode"];
+  if (streamMode === "eager" || streamMode === "stream-first") {
+    return streamMode;
+  }
+  return undefined;
+}