@docyrus/docyrus 0.0.20 → 0.0.21

package/resources/pi-agent/extensions/pi-mcp-adapter/server-manager.ts

@@ -0,0 +1,330 @@
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
+import type { ReadResourceResult } from "@modelcontextprotocol/sdk/types.js";
+import type {
+  McpTool,
+  McpResource,
+  ServerDefinition,
+  ServerStreamResultPatchNotification,
+  Transport,
+} from "./types.js";
+import { serverStreamResultPatchNotificationSchema } from "./types.js";
+import { getStoredTokens } from "./oauth-handler.js";
+import { resolveNpxBinary } from "./npx-resolver.js";
+import { logger } from "./logger.js";
+
+interface ServerConnection {
+  client: Client;
+  transport: Transport;
+  definition: ServerDefinition;
+  tools: McpTool[];
+  resources: McpResource[];
+  lastUsedAt: number;
+  inFlight: number;
+  status: "connected" | "closed";
+}
+
+type UiStreamListener = (serverName: string, notification: ServerStreamResultPatchNotification["params"]) => void;
+
+export class McpServerManager {
+  private connections = new Map<string, ServerConnection>();
+  private connectPromises = new Map<string, Promise<ServerConnection>>();
+  private uiStreamListeners = new Map<string, UiStreamListener>();
+  
+  async connect(name: string, definition: ServerDefinition): Promise<ServerConnection> {
+    // Dedupe concurrent connection attempts
+    if (this.connectPromises.has(name)) {
+      return this.connectPromises.get(name)!;
+    }
+    
+    // Reuse existing connection if healthy
+    const existing = this.connections.get(name);
+    if (existing?.status === "connected") {
+      existing.lastUsedAt = Date.now();
+      return existing;
+    }
+    
+    const promise = this.createConnection(name, definition);
+    this.connectPromises.set(name, promise);
+    
+    try {
+      const connection = await promise;
+      this.connections.set(name, connection);
+      return connection;
+    } finally {
+      this.connectPromises.delete(name);
+    }
+  }
+  
+  private async createConnection(
+    name: string,
+    definition: ServerDefinition
+  ): Promise<ServerConnection> {
+    const client = new Client({ name: `pi-mcp-${name}`, version: "1.0.0" });
+    
+    let transport: Transport;
+    
+    if (definition.command) {
+      let command = definition.command;
+      let args = definition.args ?? [];
+
+      if (command === "npx" || command === "npm") {
+        const resolved = await resolveNpxBinary(command, args);
+        if (resolved) {
+          command = resolved.isJs ? "node" : resolved.binPath;
+          args = resolved.isJs ? [resolved.binPath, ...resolved.extraArgs] : resolved.extraArgs;
+          logger.debug(`${name} resolved to ${resolved.binPath} (skipping npm parent)`);
+        }
+      }
+
+      transport = new StdioClientTransport({
+        command,
+        args,
+        env: resolveEnv(definition.env),
+        cwd: definition.cwd,
+        stderr: definition.debug ? "inherit" : "ignore",
+      });
+    } else if (definition.url) {
+      // HTTP transport with fallback
+      transport = await this.createHttpTransport(definition, name);
+    } else {
+      throw new Error(`Server ${name} has no command or url`);
+    }
+    
+    try {
+      await client.connect(transport);
+      this.attachAdapterNotificationHandlers(name, client);
+      
+      // Discover tools and resources
+      const [tools, resources] = await Promise.all([
+        this.fetchAllTools(client),
+        this.fetchAllResources(client),
+      ]);
+      
+      return {
+        client,
+        transport,
+        definition,
+        tools,
+        resources,
+        lastUsedAt: Date.now(),
+        inFlight: 0,
+        status: "connected",
+      };
+    } catch (error) {
+      // Clean up both client and transport on any error
+      await client.close().catch(() => {});
+      await transport.close().catch(() => {});
+      throw error;
+    }
+  }
+  
+  private async createHttpTransport(definition: ServerDefinition, serverName?: string): Promise<Transport> {
+    const url = new URL(definition.url!);
+    const headers = resolveHeaders(definition.headers) ?? {};
+    
+    // Add bearer token if configured
+    if (definition.auth === "bearer") {
+      const token = definition.bearerToken 
+        ?? (definition.bearerTokenEnv ? process.env[definition.bearerTokenEnv] : undefined);
+      if (token) {
+        headers["Authorization"] = `Bearer ${token}`;
+      }
+    }
+    
+    // Handle OAuth auth - use stored tokens
+    if (definition.auth === "oauth") {
+      if (!serverName) {
+        throw new Error("Server name required for OAuth authentication");
+      }
+      const tokens = getStoredTokens(serverName);
+      if (!tokens) {
+        throw new Error(
+          `No OAuth tokens found for "${serverName}". Run /mcp-auth ${serverName} to authenticate.`
+        );
+      }
+      headers["Authorization"] = `Bearer ${tokens.access_token}`;
+    }
+    
+    const requestInit = Object.keys(headers).length > 0 ? { headers } : undefined;
+    
+    // Try StreamableHTTP first (modern MCP servers)
+    const streamableTransport = new StreamableHTTPClientTransport(url, { requestInit });
+    
+    try {
+      // Create a test client to verify the transport works
+      const testClient = new Client({ name: "pi-mcp-probe", version: "1.0.0" });
+      await testClient.connect(streamableTransport);
+      await testClient.close().catch(() => {});
+      // Close probe transport before creating fresh one
+      await streamableTransport.close().catch(() => {});
+      
+      // StreamableHTTP works - create fresh transport for actual use
+      return new StreamableHTTPClientTransport(url, { requestInit });
+    } catch {
+      // StreamableHTTP failed, close and try SSE fallback
+      await streamableTransport.close().catch(() => {});
+      
+      // SSE is the legacy transport
+      return new SSEClientTransport(url, { requestInit });
+    }
+  }
+  
+  private async fetchAllTools(client: Client): Promise<McpTool[]> {
+    const allTools: McpTool[] = [];
+    let cursor: string | undefined;
+    
+    do {
+      const result = await client.listTools(cursor ? { cursor } : undefined);
+      allTools.push(...(result.tools ?? []));
+      cursor = result.nextCursor;
+    } while (cursor);
+    
+    return allTools;
+  }
+  
+  private async fetchAllResources(client: Client): Promise<McpResource[]> {
+    try {
+      const allResources: McpResource[] = [];
+      let cursor: string | undefined;
+      
+      do {
+        const result = await client.listResources(cursor ? { cursor } : undefined);
+        allResources.push(...(result.resources ?? []));
+        cursor = result.nextCursor;
+      } while (cursor);
+      
+      return allResources;
+    } catch {
+      // Server may not support resources
+      return [];
+    }
+  }
+
+  private attachAdapterNotificationHandlers(serverName: string, client: Client): void {
+    client.setNotificationHandler(serverStreamResultPatchNotificationSchema, (notification) => {
+      const listener = this.uiStreamListeners.get(notification.params.streamToken);
+      if (!listener) return;
+      listener(serverName, notification.params);
+    });
+  }
+
+  registerUiStreamListener(streamToken: string, listener: UiStreamListener): void {
+    this.uiStreamListeners.set(streamToken, listener);
+  }
+
+  removeUiStreamListener(streamToken: string): void {
+    this.uiStreamListeners.delete(streamToken);
+  }
+
+  async readResource(name: string, uri: string): Promise<ReadResourceResult> {
+    const connection = this.connections.get(name);
+    if (!connection || connection.status !== "connected") {
+      throw new Error(`Server "${name}" is not connected`);
+    }
+
+    try {
+      this.touch(name);
+      this.incrementInFlight(name);
+      return await connection.client.readResource({ uri });
+    } finally {
+      this.decrementInFlight(name);
+      this.touch(name);
+    }
+  }
+  
+  async close(name: string): Promise<void> {
+    const connection = this.connections.get(name);
+    if (!connection) return;
+    
+    // Delete from map BEFORE async cleanup to prevent a race where a
+    // concurrent connect() creates a new connection that our deferred
+    // delete() would then remove, orphaning the new server process.
+    connection.status = "closed";
+    this.connections.delete(name);
+    await connection.client.close().catch(() => {});
+    await connection.transport.close().catch(() => {});
+  }
+  
+  async closeAll(): Promise<void> {
+    const names = [...this.connections.keys()];
+    await Promise.all(names.map(name => this.close(name)));
+  }
+  
+  getConnection(name: string): ServerConnection | undefined {
+    return this.connections.get(name);
+  }
+  
+  getAllConnections(): Map<string, ServerConnection> {
+    return new Map(this.connections);
+  }
+
+  touch(name: string): void {
+    const connection = this.connections.get(name);
+    if (connection) {
+      connection.lastUsedAt = Date.now();
+    }
+  }
+
+  incrementInFlight(name: string): void {
+    const connection = this.connections.get(name);
+    if (connection) {
+      connection.inFlight = (connection.inFlight ?? 0) + 1;
+    }
+  }
+
+  decrementInFlight(name: string): void {
+    const connection = this.connections.get(name);
+    if (connection && connection.inFlight) {
+      connection.inFlight--;
+    }
+  }
+
+  isIdle(name: string, timeoutMs: number): boolean {
+    const connection = this.connections.get(name);
+    if (!connection || connection.status !== "connected") return false;
+    if (connection.inFlight > 0) return false;
+    return (Date.now() - connection.lastUsedAt) > timeoutMs;
+  }
+}
+
+/**
+ * Resolve environment variables with interpolation.
+ */
+function resolveEnv(env?: Record<string, string>): Record<string, string> {
+  // Copy process.env, filtering out undefined values
+  const resolved: Record<string, string> = {};
+  for (const [key, value] of Object.entries(process.env)) {
+    if (value !== undefined) {
+      resolved[key] = value;
+    }
+  }
+  
+  if (!env) return resolved;
+  
+  for (const [key, value] of Object.entries(env)) {
+    // Support ${VAR} and $env:VAR interpolation
+    resolved[key] = value
+      .replace(/\$\{(\w+)\}/g, (_, name) => process.env[name] ?? "")
+      .replace(/\$env:(\w+)/g, (_, name) => process.env[name] ?? "");
+  }
+  
+  return resolved;
+}
+
+/**
+ * Resolve headers with environment variable interpolation.
+ */
+function resolveHeaders(headers?: Record<string, string>): Record<string, string> | undefined {
+  if (!headers) return undefined;
+  
+  const resolved: Record<string, string> = {};
+  for (const [key, value] of Object.entries(headers)) {
+    resolved[key] = value
+      .replace(/\$\{(\w+)\}/g, (_, name) => process.env[name] ?? "")
+      .replace(/\$env:(\w+)/g, (_, name) => process.env[name] ?? "");
+  }
+  return resolved;
+}

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

@@ -0,0 +1,41 @@
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { ConsentManager } from "./consent-manager.js";
+import type { McpLifecycleManager } from "./lifecycle.js";
+import type { McpServerManager } from "./server-manager.js";
+import type { ToolMetadata, McpConfig, UiSessionMessages, UiStreamSummary } from "./types.js";
+import type { UiResourceHandler } from "./ui-resource-handler.js";
+import type { UiServerHandle } from "./ui-server.js";
+
+export interface CompletedUiSession {
+  serverName: string;
+  toolName: string;
+  completedAt: Date;
+  reason: string;
+  messages: UiSessionMessages;
+  stream?: UiStreamSummary;
+}
+
+export type SendMessageFn = (
+  message: {
+    customType: string;
+    content: Array<{ type: string; text: string }>;
+    display?: string;
+    details?: unknown;
+  },
+  options?: { triggerTurn?: boolean }
+) => void;
+
+export interface McpExtensionState {
+  manager: McpServerManager;
+  lifecycle: McpLifecycleManager;
+  toolMetadata: Map<string, ToolMetadata[]>;
+  config: McpConfig;
+  failureTracker: Map<string, number>;
+  uiResourceHandler: UiResourceHandler;
+  consentManager: ConsentManager;
+  uiServer: UiServerHandle | null;
+  completedUiSessions: CompletedUiSession[];
+  openBrowser: (url: string) => Promise<void>;
+  ui?: ExtensionContext["ui"];
+  sendMessage?: SendMessageFn;
+}

package/resources/pi-agent/extensions/pi-mcp-adapter/tool-metadata.ts

@@ -0,0 +1,144 @@
+import { getToolUiResourceUri } from "@modelcontextprotocol/ext-apps/app-bridge";
+import type { McpExtensionState } from "./state.js";
+import type { ToolMetadata, McpTool, McpResource, ServerEntry } from "./types.js";
+import { formatToolName } from "./types.js";
+import { resourceNameToToolName } from "./resource-tools.js";
+import { extractToolUiStreamMode } from "./utils.js";
+
+export function buildToolMetadata(
+  tools: McpTool[],
+  resources: McpResource[],
+  definition: ServerEntry,
+  serverName: string,
+  prefix: "server" | "none" | "short"
+): { metadata: ToolMetadata[]; failedTools: string[] } {
+  const metadata: ToolMetadata[] = [];
+  const failedTools: string[] = [];
+
+  for (const tool of tools) {
+    if (!tool?.name) {
+      failedTools.push("(unnamed)");
+      continue;
+    }
+    let uiResourceUri: string | undefined;
+    try {
+      uiResourceUri = getToolUiResourceUri({ _meta: tool._meta });
+    } catch {
+      failedTools.push(tool.name);
+    }
+    metadata.push({
+      name: formatToolName(tool.name, serverName, prefix),
+      originalName: tool.name,
+      description: tool.description ?? "",
+      inputSchema: tool.inputSchema,
+      uiResourceUri,
+      uiStreamMode: extractToolUiStreamMode(tool._meta),
+    });
+  }
+
+  if (definition.exposeResources !== false) {
+    for (const resource of resources) {
+      const baseName = `get_${resourceNameToToolName(resource.name)}`;
+      metadata.push({
+        name: formatToolName(baseName, serverName, prefix),
+        originalName: baseName,
+        description: resource.description ?? `Read resource: ${resource.uri}`,
+        resourceUri: resource.uri,
+      });
+    }
+  }
+
+  return { metadata, failedTools };
+}
+
+export function getToolNames(state: McpExtensionState, serverName: string): string[] {
+  return state.toolMetadata.get(serverName)?.map(m => m.name) ?? [];
+}
+
+export function totalToolCount(state: McpExtensionState): number {
+  let count = 0;
+  for (const metadata of state.toolMetadata.values()) {
+    count += metadata.length;
+  }
+  return count;
+}
+
+export function findToolByName(metadata: ToolMetadata[] | undefined, toolName: string): ToolMetadata | undefined {
+  if (!metadata) return undefined;
+  const exact = metadata.find(m => m.name === toolName);
+  if (exact) return exact;
+  const normalized = toolName.replace(/-/g, "_");
+  return metadata.find(m => m.name.replace(/-/g, "_") === normalized);
+}
+
+export function formatSchema(schema: unknown, indent = "  "): string {
+  if (!schema || typeof schema !== "object") {
+    return `${indent}(no schema)`;
+  }
+
+  const s = schema as Record<string, unknown>;
+
+  if (s.type === "object" && s.properties && typeof s.properties === "object") {
+    const props = s.properties as Record<string, unknown>;
+    const required = Array.isArray(s.required) ? s.required as string[] : [];
+
+    if (Object.keys(props).length === 0) {
+      return `${indent}(no parameters)`;
+    }
+
+    const lines: string[] = [];
+    for (const [name, propSchema] of Object.entries(props)) {
+      const isRequired = required.includes(name);
+      const propLine = formatProperty(name, propSchema, isRequired, indent);
+      lines.push(propLine);
+    }
+    return lines.join("\n");
+  }
+
+  if (s.type) {
+    return `${indent}(${s.type})`;
+  }
+
+  return `${indent}(complex schema)`;
+}
+
+function formatProperty(name: string, schema: unknown, required: boolean, indent: string): string {
+  if (!schema || typeof schema !== "object") {
+    return `${indent}${name}${required ? " *required*" : ""}`;
+  }
+
+  const s = schema as Record<string, unknown>;
+  const parts: string[] = [];
+
+  let typeStr = "";
+  if (s.type) {
+    if (Array.isArray(s.type)) {
+      typeStr = s.type.join(" | ");
+    } else {
+      typeStr = String(s.type);
+    }
+  } else if (s.enum) {
+    typeStr = "enum";
+  } else if (s.anyOf || s.oneOf) {
+    typeStr = "union";
+  }
+
+  if (Array.isArray(s.enum)) {
+    const enumVals = s.enum.map(v => JSON.stringify(v)).join(", ");
+    typeStr = `enum: ${enumVals}`;
+  }
+
+  parts.push(`${indent}${name}`);
+  if (typeStr) parts.push(`(${typeStr})`);
+  if (required) parts.push("*required*");
+
+  if (s.description && typeof s.description === "string") {
+    parts.push(`- ${s.description}`);
+  }
+
+  if (s.default !== undefined) {
+    parts.push(`[default: ${JSON.stringify(s.default)}]`);
+  }
+
+  return parts.join(" ");
+}

package/resources/pi-agent/extensions/pi-mcp-adapter/tool-registrar.ts

@@ -0,0 +1,46 @@
+// tool-registrar.ts - MCP content transformation
+// NOTE: Tools are NOT registered with Pi - only the unified `mcp` proxy tool is registered.
+// This keeps the LLM context small (1 tool instead of 100s).
+
+import type { McpContent, ContentBlock } from "./types.js";
+
+/**
+ * Transform MCP content types to Pi content blocks.
+ */
+export function transformMcpContent(content: McpContent[]): ContentBlock[] {
+  return content.map(c => {
+    if (c.type === "text") {
+      return { type: "text" as const, text: c.text ?? "" };
+    }
+    if (c.type === "image") {
+      return {
+        type: "image" as const,
+        data: c.data ?? "",
+        mimeType: c.mimeType ?? "image/png",
+      };
+    }
+    if (c.type === "resource") {
+      const resourceUri = c.resource?.uri ?? "(no URI)";
+      const resourceContent = c.resource?.text ?? (c.resource ? JSON.stringify(c.resource) : "(no content)");
+      return {
+        type: "text" as const,
+        text: `[Resource: ${resourceUri}]\n${resourceContent}`,
+      };
+    }
+    if (c.type === "resource_link") {
+      const linkName = c.name ?? c.uri ?? "unknown";
+      const linkUri = c.uri ?? "(no URI)";
+      return {
+        type: "text" as const,
+        text: `[Resource Link: ${linkName}]\nURI: ${linkUri}`,
+      };
+    }
+    if (c.type === "audio") {
+      return {
+        type: "text" as const,
+        text: `[Audio content: ${c.mimeType ?? "audio/*"}]`,
+      };
+    }
+    return { type: "text" as const, text: JSON.stringify(c) };
+  });
+}