@docyrus/docyrus 0.0.20 → 0.0.22

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

@@ -0,0 +1,319 @@
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { McpExtensionState } from "./state.js";
+import type { ToolMetadata } from "./types.js";
+import { existsSync } from "node:fs";
+import { loadMcpConfig } from "./config.js";
+import { ConsentManager } from "./consent-manager.js";
+import { McpLifecycleManager } from "./lifecycle.js";
+import {
+  computeServerHash,
+  getMetadataCachePath,
+  isServerCacheValid,
+  loadMetadataCache,
+  reconstructToolMetadata,
+  saveMetadataCache,
+  serializeResources,
+  serializeTools,
+  type ServerCacheEntry,
+} from "./metadata-cache.js";
+import { McpServerManager } from "./server-manager.js";
+import { buildToolMetadata, totalToolCount } from "./tool-metadata.js";
+import { UiResourceHandler } from "./ui-resource-handler.js";
+import { openUrl, parallelLimit } from "./utils.js";
+import { logger } from "./logger.js";
+
+const FAILURE_BACKOFF_MS = 60 * 1000;
+
+export async function initializeMcp(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext
+): Promise<McpExtensionState> {
+  const configPath = pi.getFlag("mcp-config") as string | undefined;
+  const config = loadMcpConfig(configPath);
+
+  const manager = new McpServerManager();
+  const lifecycle = new McpLifecycleManager(manager);
+  const toolMetadata = new Map<string, ToolMetadata[]>();
+  const failureTracker = new Map<string, number>();
+  const uiResourceHandler = new UiResourceHandler(manager);
+  const consentManager = new ConsentManager("once-per-server");
+  const ui = ctx.hasUI ? ctx.ui : undefined;
+  const state: McpExtensionState = {
+    manager,
+    lifecycle,
+    toolMetadata,
+    config,
+    failureTracker,
+    uiResourceHandler,
+    consentManager,
+    uiServer: null,
+    completedUiSessions: [],
+    openBrowser: (url: string) => openUrl(pi, url, process.env.BROWSER),
+    ui,
+    sendMessage: (message, options) => pi.sendMessage(message, options),
+  };
+
+  const serverEntries = Object.entries(config.mcpServers);
+  if (serverEntries.length === 0) {
+    return state;
+  }
+
+  const idleSetting = typeof config.settings?.idleTimeout === "number" ? config.settings.idleTimeout : 10;
+  lifecycle.setGlobalIdleTimeout(idleSetting);
+
+  const cachePath = getMetadataCachePath();
+  const cacheFileExists = existsSync(cachePath);
+  let cache = loadMetadataCache();
+  let bootstrapAll = false;
+
+  if (!cacheFileExists) {
+    bootstrapAll = true;
+    saveMetadataCache({ version: 1, servers: {} });
+  } else if (!cache) {
+    cache = { version: 1, servers: {} };
+    saveMetadataCache(cache);
+  }
+
+  const prefix = config.settings?.toolPrefix ?? "server";
+
+  for (const [name, definition] of serverEntries) {
+    const lifecycleMode = definition.lifecycle ?? "lazy";
+    const idleOverride = definition.idleTimeout ?? (lifecycleMode === "eager" ? 0 : undefined);
+    lifecycle.registerServer(
+      name,
+      definition,
+      idleOverride !== undefined ? { idleTimeout: idleOverride } : undefined
+    );
+    if (lifecycleMode === "keep-alive") {
+      lifecycle.markKeepAlive(name, definition);
+    }
+
+    if (cache?.servers?.[name] && isServerCacheValid(cache.servers[name], definition)) {
+      const metadata = reconstructToolMetadata(name, cache.servers[name], prefix, definition.exposeResources);
+      toolMetadata.set(name, metadata);
+    }
+  }
+
+  const startupServers = bootstrapAll
+    ? serverEntries
+    : serverEntries.filter(([, definition]) => {
+        const mode = definition.lifecycle ?? "lazy";
+        return mode === "keep-alive" || mode === "eager";
+      });
+
+  if (ctx.hasUI && startupServers.length > 0) {
+    ctx.ui.setStatus("mcp", `MCP: connecting to ${startupServers.length} servers...`);
+  }
+
+  const results = await parallelLimit(startupServers, 10, async ([name, definition]) => {
+    try {
+      const connection = await manager.connect(name, definition);
+      return { name, definition, connection, error: null };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      return { name, definition, connection: null, error: message };
+    }
+  });
+
+  for (const { name, definition, connection, error } of results) {
+    if (error || !connection) {
+      if (ctx.hasUI) {
+        ctx.ui.notify(`MCP: Failed to connect to ${name}: ${error}`, "error");
+      }
+      console.error(`MCP: Failed to connect to ${name}: ${error}`);
+      continue;
+    }
+
+    const { metadata, failedTools } = buildToolMetadata(connection.tools, connection.resources, definition, name, prefix);
+    toolMetadata.set(name, metadata);
+    updateMetadataCache(state, name);
+
+    if (failedTools.length > 0 && ctx.hasUI) {
+      ctx.ui.notify(
+        `MCP: ${name} - ${failedTools.length} tools skipped`,
+        "warning"
+      );
+    }
+  }
+
+  const connectedCount = results.filter(r => r.connection).length;
+  const failedCount = results.filter(r => r.error).length;
+  if (ctx.hasUI && connectedCount > 0) {
+    const totalTools = totalToolCount(state);
+    const msg = failedCount > 0
+      ? `MCP: ${connectedCount}/${startupServers.length} servers connected (${totalTools} tools)`
+      : `MCP: ${connectedCount} servers connected (${totalTools} tools)`;
+    ctx.ui.notify(msg, "info");
+  }
+
+  const envDirect = process.env.MCP_DIRECT_TOOLS;
+  if (envDirect !== "__none__") {
+    const missingCacheServers: string[] = [];
+    const currentCache = loadMetadataCache();
+    for (const [name, definition] of serverEntries) {
+      const hasDirect = definition.directTools !== undefined
+        ? !!definition.directTools
+        : !!config.settings?.directTools;
+      if (!hasDirect) continue;
+      const entry = currentCache?.servers?.[name];
+      if (!entry || !isServerCacheValid(entry, definition)) {
+        missingCacheServers.push(name);
+      }
+    }
+
+    if (missingCacheServers.length > 0) {
+      const bootstrapResults = await parallelLimit(
+        missingCacheServers.filter(name => !results.some(r => r.name === name && r.connection)),
+        10,
+        async (name) => {
+          const definition = config.mcpServers[name];
+          try {
+            const connection = await manager.connect(name, definition);
+            const { metadata } = buildToolMetadata(connection.tools, connection.resources, definition, name, prefix);
+            toolMetadata.set(name, metadata);
+            updateMetadataCache(state, name);
+            return { name, ok: true };
+          } catch {
+            return { name, ok: false };
+          }
+        },
+      );
+      const bootstrapped = bootstrapResults.filter(r => r.ok).map(r => r.name);
+      if (bootstrapped.length > 0 && ctx.hasUI) {
+        ctx.ui.notify(`MCP: direct tools for ${bootstrapped.join(", ")} will be available after restart`, "info");
+      }
+    }
+  }
+
+  lifecycle.setReconnectCallback((serverName) => {
+    updateServerMetadata(state, serverName);
+    updateMetadataCache(state, serverName);
+    state.failureTracker.delete(serverName);
+    updateStatusBar(state);
+  });
+
+  lifecycle.setIdleShutdownCallback((serverName) => {
+    const idleMinutes = getEffectiveIdleTimeoutMinutes(state, serverName);
+    logger.debug(`${serverName} shut down (idle ${idleMinutes}m)`);
+    updateStatusBar(state);
+  });
+
+  lifecycle.startHealthChecks();
+
+  return state;
+}
+
+export function updateServerMetadata(state: McpExtensionState, serverName: string): void {
+  const connection = state.manager.getConnection(serverName);
+  if (!connection || connection.status !== "connected") return;
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) return;
+
+  const prefix = state.config.settings?.toolPrefix ?? "server";
+
+  const { metadata } = buildToolMetadata(connection.tools, connection.resources, definition, serverName, prefix);
+  state.toolMetadata.set(serverName, metadata);
+}
+
+export function updateMetadataCache(state: McpExtensionState, serverName: string): void {
+  const connection = state.manager.getConnection(serverName);
+  if (!connection || connection.status !== "connected") return;
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) return;
+
+  const configHash = computeServerHash(definition);
+  const existing = loadMetadataCache();
+  const existingEntry = existing?.servers?.[serverName];
+
+  const tools = serializeTools(connection.tools);
+  let resources = definition.exposeResources === false ? [] : serializeResources(connection.resources);
+
+  if (
+    definition.exposeResources !== false &&
+    resources.length === 0 &&
+    existingEntry?.resources?.length &&
+    existingEntry.configHash === configHash
+  ) {
+    resources = existingEntry.resources;
+  }
+
+  const entry: ServerCacheEntry = {
+    configHash,
+    tools,
+    resources,
+    cachedAt: Date.now(),
+  };
+
+  saveMetadataCache({ version: 1, servers: { [serverName]: entry } });
+}
+
+export function flushMetadataCache(state: McpExtensionState): void {
+  for (const [name, connection] of state.manager.getAllConnections()) {
+    if (connection.status === "connected") {
+      updateMetadataCache(state, name);
+    }
+  }
+}
+
+export function updateStatusBar(state: McpExtensionState): void {
+  const ui = state.ui;
+  if (!ui) return;
+  const total = Object.keys(state.config.mcpServers).length;
+  if (total === 0) {
+    ui.setStatus("mcp", "");
+    return;
+  }
+  const connectedCount = state.manager.getAllConnections().size;
+  ui.setStatus("mcp", ui.theme.fg("accent", `MCP: ${connectedCount}/${total} servers`));
+}
+
+export function getFailureAgeSeconds(state: McpExtensionState, serverName: string): number | null {
+  const failedAt = state.failureTracker.get(serverName);
+  if (!failedAt) return null;
+  const ageMs = Date.now() - failedAt;
+  if (ageMs > FAILURE_BACKOFF_MS) return null;
+  return Math.round(ageMs / 1000);
+}
+
+export async function lazyConnect(state: McpExtensionState, serverName: string): Promise<boolean> {
+  const connection = state.manager.getConnection(serverName);
+  if (connection?.status === "connected") {
+    updateServerMetadata(state, serverName);
+    return true;
+  }
+
+  const failedAgo = getFailureAgeSeconds(state, serverName);
+  if (failedAgo !== null) return false;
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) return false;
+
+  try {
+    if (state.ui) {
+      state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
+    }
+    await state.manager.connect(serverName, definition);
+    state.failureTracker.delete(serverName);
+    updateServerMetadata(state, serverName);
+    updateMetadataCache(state, serverName);
+    updateStatusBar(state);
+    return true;
+  } catch {
+    state.failureTracker.set(serverName, Date.now());
+    updateStatusBar(state);
+    return false;
+  }
+}
+
+function getEffectiveIdleTimeoutMinutes(state: McpExtensionState, serverName: string): number {
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) {
+    return typeof state.config.settings?.idleTimeout === "number" ? state.config.settings.idleTimeout : 10;
+  }
+  if (typeof definition.idleTimeout === "number") return definition.idleTimeout;
+  const mode = definition.lifecycle ?? "lazy";
+  if (mode === "eager") return 0;
+  return typeof state.config.settings?.idleTimeout === "number" ? state.config.settings.idleTimeout : 10;
+}

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

@@ -0,0 +1,93 @@
+import type { ServerDefinition } from "./types.js";
+import type { McpServerManager } from "./server-manager.js";
+import { logger } from "./logger.js";
+
+export type ReconnectCallback = (serverName: string) => void;
+
+export class McpLifecycleManager {
+  private manager: McpServerManager;
+  private keepAliveServers = new Map<string, ServerDefinition>();
+  private allServers = new Map<string, ServerDefinition>();
+  private serverSettings = new Map<string, { idleTimeout?: number }>();
+  private globalIdleTimeout: number = 10 * 60 * 1000;
+  private healthCheckInterval?: NodeJS.Timeout;
+  private onReconnect?: ReconnectCallback;
+  private onIdleShutdown?: (serverName: string) => void;
+  
+  constructor(manager: McpServerManager) {
+    this.manager = manager;
+  }
+  
+  /**
+   * Set callback to be invoked after a successful auto-reconnect.
+   * Use this to update tool metadata when a server reconnects.
+   */
+  setReconnectCallback(callback: ReconnectCallback): void {
+    this.onReconnect = callback;
+  }
+  
+  markKeepAlive(name: string, definition: ServerDefinition): void {
+    this.keepAliveServers.set(name, definition);
+  }
+
+  registerServer(name: string, definition: ServerDefinition, settings?: { idleTimeout?: number }): void {
+    this.allServers.set(name, definition);
+    if (settings?.idleTimeout !== undefined) {
+      this.serverSettings.set(name, settings);
+    }
+  }
+
+  setGlobalIdleTimeout(minutes: number): void {
+    this.globalIdleTimeout = minutes * 60 * 1000;
+  }
+
+  setIdleShutdownCallback(callback: (serverName: string) => void): void {
+    this.onIdleShutdown = callback;
+  }
+  
+  startHealthChecks(intervalMs = 30000): void {
+    this.healthCheckInterval = setInterval(() => {
+      this.checkConnections();
+    }, intervalMs);
+    this.healthCheckInterval.unref();
+  }
+  
+  private async checkConnections(): Promise<void> {
+    for (const [name, definition] of this.keepAliveServers) {
+      const connection = this.manager.getConnection(name);
+      
+      if (!connection || connection.status !== "connected") {
+        try {
+          await this.manager.connect(name, definition);
+          logger.debug(`Reconnected to ${name}`);
+          // Notify extension to update metadata
+          this.onReconnect?.(name);
+        } catch (error) {
+          console.error(`MCP: Failed to reconnect to ${name}:`, error);
+        }
+      }
+    }
+
+    for (const [name] of this.allServers) {
+      if (this.keepAliveServers.has(name)) continue;
+      const timeout = this.getIdleTimeout(name);
+      if (timeout > 0 && this.manager.isIdle(name, timeout)) {
+        await this.manager.close(name);
+        this.onIdleShutdown?.(name);
+      }
+    }
+  }
+
+  private getIdleTimeout(name: string): number {
+    const perServer = this.serverSettings.get(name)?.idleTimeout;
+    if (perServer !== undefined) return perServer * 60 * 1000;
+    return this.globalIdleTimeout;
+  }
+  
+  async gracefulShutdown(): Promise<void> {
+    if (this.healthCheckInterval) {
+      clearInterval(this.healthCheckInterval);
+    }
+    await this.manager.closeAll();
+  }
+}

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

@@ -0,0 +1,169 @@
+/**
+ * Centralized logging for MCP UI operations.
+ * Provides structured, contextual logs with levels.
+ */
+
+export type LogLevel = "debug" | "info" | "warn" | "error";
+
+export interface LogContext {
+  server?: string;
+  session?: string;
+  tool?: string;
+  uri?: string;
+  [key: string]: unknown;
+}
+
+export interface LogEntry {
+  level: LogLevel;
+  message: string;
+  context?: LogContext;
+  error?: Error;
+  timestamp: Date;
+}
+
+type LogHandler = (entry: LogEntry) => void;
+
+const LEVEL_PRIORITY: Record<LogLevel, number> = {
+  debug: 0,
+  info: 1,
+  warn: 2,
+  error: 3,
+};
+
+const LEVEL_PREFIX: Record<LogLevel, string> = {
+  debug: "[MCP-UI:DEBUG]",
+  info: "[MCP-UI]",
+  warn: "[MCP-UI:WARN]",
+  error: "[MCP-UI:ERROR]",
+};
+
+class Logger {
+  private minLevel: LogLevel = "info";
+  private handlers: LogHandler[] = [];
+  private defaultContext: LogContext = {};
+
+  setLevel(level: LogLevel): void {
+    this.minLevel = level;
+  }
+
+  setDefaultContext(context: LogContext): void {
+    this.defaultContext = context;
+  }
+
+  addHandler(handler: LogHandler): void {
+    this.handlers.push(handler);
+  }
+
+  clearHandlers(): void {
+    this.handlers = [];
+  }
+
+  private shouldLog(level: LogLevel): boolean {
+    return LEVEL_PRIORITY[level] >= LEVEL_PRIORITY[this.minLevel];
+  }
+
+  private emit(level: LogLevel, message: string, context?: LogContext, error?: Error): void {
+    if (!this.shouldLog(level)) return;
+
+    const entry: LogEntry = {
+      level,
+      message,
+      context: { ...this.defaultContext, ...context },
+      error,
+      timestamp: new Date(),
+    };
+
+    // Default console output
+    const prefix = LEVEL_PREFIX[level];
+    const contextStr = formatContext(entry.context);
+    const fullMessage = contextStr ? `${prefix} ${message} ${contextStr}` : `${prefix} ${message}`;
+
+    if (level === "error") {
+      console.error(fullMessage, error ?? "");
+    } else if (level === "warn") {
+      console.warn(fullMessage);
+    } else if (level === "debug") {
+      console.debug(fullMessage);
+    } else {
+      console.log(fullMessage);
+    }
+
+    // Custom handlers
+    for (const handler of this.handlers) {
+      try {
+        handler(entry);
+      } catch {
+        // Ignore handler errors
+      }
+    }
+  }
+
+  debug(message: string, context?: LogContext): void {
+    this.emit("debug", message, context);
+  }
+
+  info(message: string, context?: LogContext): void {
+    this.emit("info", message, context);
+  }
+
+  warn(message: string, context?: LogContext): void {
+    this.emit("warn", message, context);
+  }
+
+  error(message: string, error?: Error, context?: LogContext): void {
+    this.emit("error", message, context, error);
+  }
+
+  /**
+   * Create a child logger with additional default context.
+   */
+  child(context: LogContext): ChildLogger {
+    return new ChildLogger(this, context);
+  }
+}
+
+class ChildLogger {
+  constructor(
+    private parent: Logger,
+    private context: LogContext
+  ) {}
+
+  debug(message: string, context?: LogContext): void {
+    this.parent.debug(message, { ...this.context, ...context });
+  }
+
+  info(message: string, context?: LogContext): void {
+    this.parent.info(message, { ...this.context, ...context });
+  }
+
+  warn(message: string, context?: LogContext): void {
+    this.parent.warn(message, { ...this.context, ...context });
+  }
+
+  error(message: string, error?: Error, context?: LogContext): void {
+    this.parent.error(message, error, { ...this.context, ...context });
+  }
+
+  child(context: LogContext): ChildLogger {
+    return new ChildLogger(this.parent, { ...this.context, ...context });
+  }
+}
+
+function formatContext(context?: LogContext): string {
+  if (!context || Object.keys(context).length === 0) return "";
+  const parts: string[] = [];
+  for (const [key, value] of Object.entries(context)) {
+    if (value !== undefined && value !== null) {
+      parts.push(`${key}=${typeof value === "string" ? value : JSON.stringify(value)}`);
+    }
+  }
+  return parts.length > 0 ? `(${parts.join(", ")})` : "";
+}
+
+// Singleton instance
+export const logger = new Logger();
+
+// Enable debug mode via environment variable
+if (process.env.MCP_UI_DEBUG === "1" || process.env.MCP_UI_DEBUG === "true") {
+  logger.setLevel("debug");
+}