pi-soly 1.9.3 → 1.11.0

package/mcp/server-manager.ts

@@ -0,0 +1,545 @@
+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 { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
+import {
+  ElicitationCompleteNotificationSchema,
+  type CallToolRequest,
+  type CallToolResult,
+  type ReadResourceResult,
+  type UrlElicitationRequiredError,
+} from "@modelcontextprotocol/sdk/types.js";
+import type {
+  McpTool,
+  McpResource,
+  ServerDefinition,
+  ServerStreamResultPatchNotification,
+  Transport,
+} from "./types.ts";
+import { serverStreamResultPatchNotificationSchema } from "./types.ts";
+import { resolveNpxBinary } from "./npx-resolver.ts";
+import { logger } from "./logger.ts";
+import { McpOAuthProvider } from "./mcp-oauth-provider.ts";
+import { extractOAuthConfig, supportsOAuth } from "./mcp-auth-flow.ts";
+import { registerSamplingHandler, type ServerSamplingConfig } from "./sampling-handler.ts";
+import {
+  handleUrlElicitation,
+  registerElicitationHandler,
+  type ServerElicitationConfig,
+} from "./elicitation-handler.ts";
+import { interpolateEnvRecord, resolveBearerToken, resolveConfigPath } from "./utils.ts";
+
+interface ServerConnection {
+  client: Client;
+  transport: Transport;
+  definition: ServerDefinition;
+  tools: McpTool[];
+  resources: McpResource[];
+  lastUsedAt: number;
+  inFlight: number;
+  status: "connected" | "closed" | "needs-auth";
+}
+
+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>();
+  private samplingConfig: ServerSamplingConfig | undefined;
+  private elicitationConfig: ServerElicitationConfig | undefined;
+  private acceptedUrlElicitations = new Map<string, Set<string>>();
+  /** Called after an auto-recovery attempt: args are (serverName, ok, errorMsg).
+   *  Fired from `callTool()` when a session-expired error triggered close+reconnect.
+   *  Use to surface notifications ("session auto-recovered" / "recovery failed"). */
+  onSessionRecover?: (serverName: string, ok: boolean, errorMsg: string) => void;
+
+  setSamplingConfig(config: ServerSamplingConfig | undefined): void {
+    this.samplingConfig = config;
+  }
+
+  setElicitationConfig(config: ServerElicitationConfig | undefined): void {
+    this.elicitationConfig = config;
+  }
+
+  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 = this.createClient(name);
+
+    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: resolveConfigPath(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) {
+      // Check for UnauthorizedError - server requires OAuth
+      if (error instanceof UnauthorizedError && supportsOAuth(definition)) {
+        // Clean up both client and transport before reporting needs-auth.
+        await client.close().catch(() => {});
+        await transport.close().catch(() => {});
+
+        return {
+          client,
+          transport,
+          definition,
+          tools: [],
+          resources: [],
+          lastUsedAt: Date.now(),
+          inFlight: 0,
+          status: "needs-auth",
+        };
+      }
+
+      // Clean up both client and transport on any error
+      await client.close().catch(() => {});
+      await transport.close().catch(() => {});
+      throw error;
+    }
+  }
+
+  private buildClientCapabilities() {
+    return {
+      ...(this.samplingConfig ? { sampling: {} } : {}),
+      ...(this.elicitationConfig
+        ? {
+            elicitation: {
+              form: {},
+              ...(this.elicitationConfig.allowUrl ? { url: {} } : {}),
+            },
+          }
+        : {}),
+    };
+  }
+
+  private createClient(serverName: string): Client {
+    const capabilities = this.buildClientCapabilities();
+    const client = new Client(
+      { name: `pi-mcp-${serverName}`, version: "1.0.0" },
+      Object.keys(capabilities).length > 0 ? { capabilities } : undefined,
+    );
+    if (this.samplingConfig) {
+      registerSamplingHandler(client, { ...this.samplingConfig, serverName });
+    }
+    if (this.elicitationConfig) {
+      registerElicitationHandler(client, {
+        ...this.elicitationConfig,
+        serverName,
+        onUrlAccepted: elicitationId => this.rememberUrlElicitation(serverName, elicitationId),
+      });
+      if (this.elicitationConfig.allowUrl) {
+        client.setNotificationHandler(ElicitationCompleteNotificationSchema, notification => {
+          const accepted = this.acceptedUrlElicitations.get(serverName);
+          if (!accepted?.delete(notification.params.elicitationId)) return;
+          this.elicitationConfig?.ui.notify(
+            `MCP browser interaction for ${serverName} completed. You can retry the tool now.`,
+            "info",
+          );
+        });
+      }
+    }
+    return client;
+  }
+
+  async handleUrlElicitationRequired(
+    serverName: string,
+    error: UrlElicitationRequiredError,
+  ): Promise<"accept" | "decline" | "cancel"> {
+    if (!this.elicitationConfig?.allowUrl) return "cancel";
+    for (const params of error.elicitations) {
+      const result = await handleUrlElicitation({
+        ...this.elicitationConfig,
+        serverName,
+        onUrlAccepted: elicitationId => this.rememberUrlElicitation(serverName, elicitationId),
+      }, params);
+      if (result.action !== "accept") return result.action;
+    }
+    return "accept";
+  }
+
+  private rememberUrlElicitation(serverName: string, elicitationId: string): void {
+    let accepted = this.acceptedUrlElicitations.get(serverName);
+    if (!accepted) {
+      accepted = new Set();
+      this.acceptedUrlElicitations.set(serverName, accepted);
+    }
+    accepted.add(elicitationId);
+  }
+
+  private async createHttpTransport(
+    definition: ServerDefinition,
+    serverName: string
+  ): Promise<Transport> {
+    const url = new URL(definition.url!);
+
+    // Build headers first (including any bearer token)
+    const headers = resolveHeaders(definition.headers) ?? {};
+
+    // For bearer auth, add the token to headers BEFORE creating requestInit
+    if (definition.auth === "bearer") {
+      const token = resolveBearerToken(definition);
+      if (token) {
+        headers["Authorization"] = `Bearer ${token}`;
+      }
+    }
+
+    // Create request init with headers (Authorization now included for bearer auth)
+    const requestInit = Object.keys(headers).length > 0 ? { headers } : undefined;
+
+    // For OAuth servers, create an auth provider
+    let authProvider: McpOAuthProvider | undefined;
+    if (supportsOAuth(definition)) {
+      const oauthConfig = extractOAuthConfig(definition);
+      authProvider = new McpOAuthProvider(
+        serverName,
+        definition.url!,
+        oauthConfig,
+        {
+          onRedirect: async (_authUrl) => {
+            // URL is captured by startAuth, no need to log
+          },
+        }
+      );
+    }
+
+    // Try StreamableHTTP first (modern MCP servers)
+    const streamableTransport = new StreamableHTTPClientTransport(url, {
+      requestInit,
+      authProvider,
+    });
+
+    try {
+      // Create a test client to verify the transport works
+      const testClient = new Client({ name: "pi-mcp-probe", version: "2.1.2" });
+      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, authProvider });
+    } catch (error) {
+      // StreamableHTTP failed, close and try SSE fallback
+      await streamableTransport.close().catch(() => {});
+
+      // If this was an UnauthorizedError, don't try SSE - the server needs auth
+      if (error instanceof UnauthorizedError) {
+        throw error;
+      }
+
+      // SSE is the legacy transport
+      return new SSEClientTransport(url, { requestInit, authProvider });
+    }
+  }
+
+  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);
+    }
+  }
+
+  /**
+   * Call a tool with automatic session-recovery.
+   *
+   * Some MCP servers (notably UE5) return a session id in the `Mcp-Session-Id`
+   * response header on `initialize`, then immediately forget it. The next
+   * `tools/call` fails with "Unknown session id ... client should reinitialize".
+   *
+   * This wrapper catches that specific error pattern, closes + reconnects the
+   * server (which triggers a fresh `initialize` and new session id), and
+   * retries the call once. Other errors are propagated unchanged.
+   *
+   * Note: this is a per-call recovery. The MCP SDK's transport caches the
+   * session id internally; our fix works because `close()` + `connect()`
+   * creates a brand new transport + client, so the SDK starts fresh.
+   */
+  async callTool(
+    name: string,
+    params: CallToolRequest["params"],
+    options?: { signal?: AbortSignal; _bookkeep?: boolean },
+  ): Promise<CallToolResult> {
+    const doCall = async (): Promise<CallToolResult> => {
+      const connection = this.connections.get(name);
+      if (!connection || connection.status !== "connected") {
+        throw new Error(`Server "${name}" is not connected`);
+      }
+      // Skip bookkeeping if the caller already did it (legacy path).
+      const bookkeeping = options?._bookkeep !== false;
+      if (bookkeeping) {
+        this.touch(name);
+        this.incrementInFlight(name);
+      }
+      try {
+        // The SDK's callTool signature changed in 1.28 to accept an optional
+        // resultSchema; we still target 1.27.1. Cast through unknown to
+        // bridge the type skew without losing type safety elsewhere.
+        const result = await (connection.client.callTool as (
+          p: CallToolRequest["params"],
+          o?: unknown,
+        ) => Promise<CallToolResult>)(
+          params as CallToolRequest["params"],
+          options,
+        );
+        return result;
+      } finally {
+        if (bookkeeping) {
+          this.decrementInFlight(name);
+          this.touch(name);
+        }
+      }
+    };
+
+    try {
+      return await doCall();
+    } catch (firstError) {
+      if (!this.isSessionExpiredError(firstError)) throw firstError;
+      // Session expired — reconnect and retry once. Caller can listen via
+      // the optional onSessionRecover callback to show a notification.
+      const connection = this.connections.get(name);
+      if (!connection) throw firstError;
+      const errMsg = firstError instanceof Error ? firstError.message : String(firstError);
+      try {
+        await this.close(name);
+      } catch {
+        // ignore close errors during recovery
+      }
+      await this.connect(name, connection.definition);
+      try {
+        const result = await (connection.client.callTool as (
+          p: CallToolRequest["params"],
+          o?: unknown,
+        ) => Promise<CallToolResult>)(
+          params as CallToolRequest["params"],
+          options,
+        );
+        this.onSessionRecover?.(name, true, errMsg);
+        return result;
+      } catch (retryErr) {
+        this.onSessionRecover?.(name, false, errMsg);
+        throw retryErr;
+      }
+    }
+  }
+
+  /** Detect the UE5-style "session expired" error pattern. */
+  private isSessionExpiredError(error: unknown): boolean {
+    if (!error) return false;
+    const msg =
+      error instanceof Error
+        ? error.message
+        : typeof error === "object" && error !== null && "message" in error
+          ? String((error as { message: unknown }).message)
+          : String(error);
+    return /unknown session id|client should reinitialize|session.*expired|session.*invalid/i.test(
+      msg,
+    );
+  }
+
+  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);
+    this.acceptedUrlElicitations.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;
+
+  const overrides = interpolateEnvRecord(env);
+  return overrides ? { ...resolved, ...overrides } : resolved;
+}
+
+/**
+ * Resolve headers with environment variable interpolation.
+ */
+function resolveHeaders(headers?: Record<string, string>): Record<string, string> | undefined {
+  return interpolateEnvRecord(headers);
+}

package/mcp/state.ts

@@ -0,0 +1,41 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { ConsentManager } from "./consent-manager.ts";
+import type { McpLifecycleManager } from "./lifecycle.ts";
+import type { McpServerManager } from "./server-manager.ts";
+import type { ToolMetadata, McpConfig, UiSessionMessages, UiStreamSummary } from "./types.ts";
+import type { UiResourceHandler } from "./ui-resource-handler.ts";
+import type { UiServerHandle } from "./ui-server.ts";
+
+export interface CompletedUiSession {
+  serverName: string;
+  toolName: string;
+  completedAt: Date;
+  reason: string;
+  messages: UiSessionMessages;
+  stream?: UiStreamSummary;
+}
+
+export type SendMessageFn = (
+  message: {
+    customType: string;
+    content: Array<{ type: "text"; 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;
+}