pi-mcp-adapter 2.25.0 → 2.26.1

package/lifecycle.ts

@@ -1,11 +1,26 @@
+import { SdkError, SdkErrorCode } from "@modelcontextprotocol/client";
 import { isServerDisabled, type ServerDefinition } from "./types.ts";
-import type { McpServerManager } from "./server-manager.ts";
+import type { McpServerManager, ServerConnection } from "./server-manager.ts";
 import { hasPendingAuth } from "./mcp-auth-flow.ts";
 import { logger } from "./logger.ts";
-import { formatTerminalError, sanitizeTerminalText } from "./utils.ts";
+import { formatTerminalError, parallelLimit, sanitizeTerminalText } from "./utils.ts";
+import { isTerminatedSession } from "./session-recovery.ts";
 
-export type ReconnectCallback = (serverName: string) => void;
+export type ReconnectCallback = (serverName: string) => void | Promise<void>;
 export type ReconnectFailureCallback = (serverName: string, error: unknown) => void;
+export type HealthRestoredCallback = (serverName: string) => void | Promise<void>;
+export type AuthRequiredCallback = (serverName: string) => void | Promise<void>;
+
+const KEEP_ALIVE_RETRY_BASE_MS = 30_000;
+const KEEP_ALIVE_RETRY_MAX_MS = 5 * 60_000;
+const KEEP_ALIVE_CHECK_CONCURRENCY = 10;
+
+interface RetryState {
+  attempts: number;
+  nextAttemptAt: number;
+  connection: ServerConnection | undefined;
+  status: ServerConnection["status"] | undefined;
+}
 
 export class McpLifecycleManager {
   private keepAliveServers = new Map<string, ServerDefinition>();
@@ -15,9 +30,14 @@ export class McpLifecycleManager {
   private healthCheckInterval: NodeJS.Timeout | undefined;
   private onReconnect: ReconnectCallback | undefined;
   private onReconnectFailure: ReconnectFailureCallback | undefined;
+  private onHealthRestored: HealthRestoredCallback | undefined;
+  private onAuthRequired: AuthRequiredCallback | undefined;
   private onIdleShutdown: ((serverName: string) => void) | undefined;
   private activeHealthCheck: Promise<void> | undefined;
+  private activeConvergence: Promise<void> | undefined;
   private shutdownPromise: Promise<void> | undefined;
+  private retryStates = new Map<string, RetryState>();
+  private pendingMetadataPublications = new Set<string>();
   private stopped = false;
   private removeHealthAbortListener: (() => void) | undefined;
 
@@ -34,6 +54,14 @@ export class McpLifecycleManager {
     this.onReconnectFailure = callback;
   }
 
+  setHealthRestoredCallback(callback: HealthRestoredCallback): void {
+    this.onHealthRestored = callback;
+  }
+
+  setAuthRequiredCallback(callback: AuthRequiredCallback): void {
+    this.onAuthRequired = callback;
+  }
+
   markKeepAlive(name: string, definition: ServerDefinition): void {
     if (isServerDisabled(definition)) return;
     this.keepAliveServers.set(name, definition);
@@ -82,29 +110,23 @@ export class McpLifecycleManager {
     this.healthCheckInterval.unref();
   }
 
-  private async checkConnections(signal?: AbortSignal): Promise<void> {
+  async ensureConverged(signal?: AbortSignal): Promise<void> {
     if (this.stopped || signal?.aborted) return;
-    for (const [name, definition] of this.keepAliveServers) {
-      if (isServerDisabled(definition)) continue;
-      const connection = this.manager.getConnection(name);
-      if (!connection || connection.status !== "connected") {
-        if (this.hasPendingAuthForServer(name)) {
-          logger.debug(`Skipping reconnect for ${name} while OAuth authorization is pending`);
-          continue;
-        }
-        try {
-          await this.manager.connect(name, definition, signal);
-          if (this.stopped || signal?.aborted) return;
-          logger.debug(`Reconnected to ${name}`);
-          this.onReconnect?.(name);
-        } catch (error) {
-          if (this.stopped || signal?.aborted) return;
-          this.onReconnectFailure?.(name, error);
-          const message = error instanceof Error ? error.message : String(error);
-          console.error(`MCP: Failed to reconnect to ${name}: ${sanitizeTerminalText(message)}`);
-        }
-      }
+    if (this.activeConvergence) return this.activeConvergence;
+
+    const check = this.checkKeepAliveConnections(signal);
+    this.activeConvergence = check;
+    try {
+      await check;
+    } finally {
+      if (this.activeConvergence === check) this.activeConvergence = undefined;
     }
+  }
+
+  private async checkConnections(signal?: AbortSignal): Promise<void> {
+    if (this.stopped || signal?.aborted) return;
+    await this.ensureConverged(signal);
+    if (this.stopped || signal?.aborted) return;
 
     for (const [name] of this.allServers) {
       if (this.keepAliveServers.has(name)) continue;
@@ -117,6 +139,205 @@ export class McpLifecycleManager {
     }
   }
 
+  private async checkKeepAliveConnections(signal?: AbortSignal): Promise<void> {
+    if (this.stopped || signal?.aborted) return;
+    await parallelLimit(
+      [...this.keepAliveServers],
+      KEEP_ALIVE_CHECK_CONCURRENCY,
+      ([name, definition]) => this.checkKeepAliveConnection(name, definition, signal),
+    );
+  }
+
+  private async checkKeepAliveConnection(
+    name: string,
+    definition: ServerDefinition,
+    signal?: AbortSignal,
+    retrySuperseded = true,
+  ): Promise<void> {
+    if (isServerDisabled(definition) || this.stopped || signal?.aborted) return;
+    const connection = this.manager.getConnection(name);
+    if (connection?.status === "needs-auth") {
+      this.pendingMetadataPublications.delete(name);
+      return;
+    }
+    if (!this.shouldAttemptConnection(name, connection)) return;
+    if (!connection || connection.status !== "connected") {
+      if (this.hasPendingAuthForServer(name)) {
+        logger.debug(`Skipping reconnect for ${name} while OAuth authorization is pending`);
+        return;
+      }
+      let freshConnection: ServerConnection;
+      try {
+        freshConnection = await this.manager.connect(name, definition, signal);
+      } catch (error) {
+        if (this.stopped || signal?.aborted) return;
+        this.reportConnectionFailure(name, error, "reconnect", this.manager.getConnection(name));
+        return;
+      }
+      if (this.stopped || signal?.aborted) return;
+      if (freshConnection.status === "needs-auth") {
+        await this.notifyAuthRequired(name);
+        return;
+      }
+      if (freshConnection.status !== "connected") {
+        this.reportConnectionFailure(
+          name,
+          new Error(`MCP server ${name} did not return a connected session`),
+          "reconnect",
+          freshConnection,
+        );
+        return;
+      }
+      logger.debug(`Reconnected to ${name}`);
+      await this.publishConnectedMetadata(name, freshConnection);
+      return;
+    }
+
+    if (this.pendingMetadataPublications.has(name)) {
+      await this.publishConnectedMetadata(name, connection);
+      return;
+    }
+    if (!definition.url) return;
+    const hadSessionId = (connection.transport as { sessionId?: string } | undefined)?.sessionId != null;
+    let refreshResult: Awaited<ReturnType<McpServerManager["refreshTools"]>>;
+    try {
+      refreshResult = await this.manager.refreshTools(name, connection, signal);
+    } catch (error) {
+      if (this.stopped || signal?.aborted) return;
+      const current = this.manager.getConnection(name);
+      if (current !== connection || connection.status !== "connected") {
+        await this.handleSupersededConnection(name, definition, connection, signal, retrySuperseded);
+        return;
+      }
+      if (!shouldReconnectAfterRefresh(error, hadSessionId)) {
+        this.reportConnectionFailure(name, error, "refresh", connection);
+        return;
+      }
+      if (this.hasPendingAuthForServer(name)) {
+        logger.debug(`Skipping reconnect for ${name} while OAuth authorization is pending`);
+        return;
+      }
+      let freshConnection: ServerConnection;
+      try {
+        freshConnection = await this.manager.reconnect(name, definition, connection, signal);
+      } catch (reconnectError) {
+        if (this.stopped || signal?.aborted) return;
+        this.reportConnectionFailure(name, reconnectError, "reconnect", this.manager.getConnection(name));
+        return;
+      }
+      if (this.stopped || signal?.aborted) return;
+      if (freshConnection.status === "needs-auth") {
+        await this.notifyAuthRequired(name);
+        return;
+      }
+      if (freshConnection.status !== "connected") {
+        this.reportConnectionFailure(
+          name,
+          new Error(`MCP server ${name} did not return a connected session`),
+          "reconnect",
+          freshConnection,
+        );
+        return;
+      }
+      logger.debug(`Reconnected stale MCP session for ${name}`);
+      await this.publishConnectedMetadata(name, freshConnection);
+      return;
+    }
+
+    if (refreshResult === "superseded") {
+      await this.handleSupersededConnection(name, definition, connection, signal, retrySuperseded);
+      return;
+    }
+    if (this.retryStates.delete(name)) {
+      await this.onHealthRestored?.(name);
+    }
+  }
+
+  private async handleSupersededConnection(
+    name: string,
+    definition: ServerDefinition,
+    staleConnection: ServerConnection,
+    signal: AbortSignal | undefined,
+    retrySuperseded: boolean,
+  ): Promise<void> {
+    const current = this.manager.getConnection(name);
+    if (current === staleConnection && current.status === "connected") {
+      if (this.retryStates.delete(name)) await this.onHealthRestored?.(name);
+      return;
+    }
+    if (current?.status === "connected") {
+      await this.publishConnectedMetadata(name, current);
+      return;
+    }
+    if (current?.status === "needs-auth") {
+      await this.notifyAuthRequired(name);
+      return;
+    }
+    if (retrySuperseded) {
+      await this.checkKeepAliveConnection(name, definition, signal, false);
+    }
+  }
+
+  private async publishConnectedMetadata(name: string, connection: ServerConnection): Promise<void> {
+    this.pendingMetadataPublications.add(name);
+    try {
+      await this.onReconnect?.(name);
+      this.pendingMetadataPublications.delete(name);
+      this.retryStates.delete(name);
+    } catch (error) {
+      if (this.stopped) return;
+      this.reportConnectionFailure(name, error, "publish", connection);
+    }
+  }
+
+  private async notifyAuthRequired(name: string): Promise<void> {
+    this.pendingMetadataPublications.delete(name);
+    this.retryStates.delete(name);
+    try {
+      await this.onAuthRequired?.(name);
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      logger.debug(`MCP: auth-required callback failed for ${name}: ${sanitizeTerminalText(message)}`);
+    }
+  }
+
+  private shouldAttemptConnection(name: string, connection: ServerConnection | undefined): boolean {
+    const retry = this.retryStates.get(name);
+    if (!retry) return true;
+    if (retry.connection !== connection || retry.status !== connection?.status) {
+      this.retryStates.delete(name);
+      return true;
+    }
+    return Date.now() >= retry.nextAttemptAt;
+  }
+
+  private reportConnectionFailure(
+    name: string,
+    error: unknown,
+    action: "refresh" | "reconnect" | "publish",
+    connection: ServerConnection | undefined,
+  ): void {
+    const attempts = (this.retryStates.get(name)?.attempts ?? 0) + 1;
+    const delay = Math.min(
+      KEEP_ALIVE_RETRY_BASE_MS * 2 ** Math.min(attempts - 1, 10),
+      KEEP_ALIVE_RETRY_MAX_MS,
+    );
+    this.retryStates.set(name, {
+      attempts,
+      nextAttemptAt: Date.now() + delay,
+      connection,
+      status: connection?.status,
+    });
+    this.onReconnectFailure?.(name, error);
+    const message = error instanceof Error ? error.message : String(error);
+    const target = action === "reconnect"
+      ? `reconnect to ${name}`
+      : action === "publish"
+        ? `publish metadata for ${name}`
+        : `refresh ${name}`;
+    console.error(`MCP: Failed to ${target}: ${sanitizeTerminalText(message)}`);
+  }
+
   private getIdleTimeout(name: string): number {
     const perServer = this.serverSettings.get(name)?.idleTimeout;
     if (perServer !== undefined) return perServer * 60 * 1000;
@@ -137,11 +358,23 @@ export class McpLifecycleManager {
     this.removeHealthAbortListener = undefined;
     await this.activeHealthCheck;
     this.activeHealthCheck = undefined;
+    await this.activeConvergence;
+    this.activeConvergence = undefined;
     this.onReconnect = undefined;
     this.onReconnectFailure = undefined;
+    this.onHealthRestored = undefined;
+    this.onAuthRequired = undefined;
     this.onIdleShutdown = undefined;
+    this.retryStates.clear();
+    this.pendingMetadataPublications.clear();
     if (typeof this.manager.closeAll === "function") {
       await this.manager.closeAll();
     }
   }
 }
+
+function shouldReconnectAfterRefresh(error: unknown, hadSessionId: boolean): boolean {
+  if (isTerminatedSession(error, hadSessionId)) return true;
+  return error instanceof SdkError
+    && (error.code === SdkErrorCode.NotConnected || error.code === SdkErrorCode.ConnectionClosed);
+}

package/mcp-setup-panel.ts

@@ -1,6 +1,7 @@
 import { matchesKey, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
 import { createPanelKeys, type PanelKeybindings, type PanelKeys } from "./panel-keys.ts";
 import type { ImportKind } from "./types.ts";
+import { getConfigDirName } from "./agent-dir.ts";
 import { KNOWN_SERVER_PRESETS, type ConfigWritePreview, type KnownServerPreset, type McpDiscoverySummary } from "./config.ts";
 import type { McpOnboardingState } from "./onboarding-state.ts";
 
@@ -556,7 +557,7 @@ export class McpSetupPanel {
           "3. ~/.agents/mcp/mcp.json",
           "4. <Pi agent dir>/mcp.json",
           "5. .mcp.json",
-          "6. .pi/mcp.json",
+          `6. ${getConfigDirName()}/mcp.json`,
           `Host discovery: ${this.discovery.hostConfigDiscovery}. Conflicts reported: ${this.discovery.conflicts.length}.`,
           ...this.discovery.conflicts.slice(0, 8).map((conflict) =>
             `${conflict.serverName}: ${conflict.sources.map((source) => source.path).join(" -> ")} (winner: ${conflict.winner.path})`,

package/metadata-cache.ts

@@ -19,11 +19,12 @@ import type {
   ToolMetadata,
   PromptMetadata,
 } from "./types.ts";
-import { formatPromptCommandName, formatToolName, getToolNameCandidates, isServerDisabled, isToolAllowed, resolveToolPrefix, type ToolPrefix } from "./types.ts";
+import { createToolSelectorCandidateIndex, formatPromptCommandName, formatToolName, getToolNameCandidates, isServerDisabled, isToolAllowed, resolveToolPrefix, type ToolPrefix } from "./types.ts";
 import { resourceNameToToolName } from "./resource-tools.ts";
 import {
   extractToolUiStreamMode,
   interpolateEnvRecord,
+  interpolateEnvVars,
   resolveBearerToken,
   resolveConfigPath,
   resolveServerUrl,
@@ -90,6 +91,14 @@ export function computeServerHash(definition: ServerEntry): string {
     cwd: resolveConfigPath(definition.cwd),
     url: resolveServerUrl(definition),
     headers: interpolateEnvRecord(definition.headers),
+    requestHeadersCommand: definition.requestHeadersCommand
+      ? {
+          command: interpolateEnvVars(definition.requestHeadersCommand.command),
+          args: definition.requestHeadersCommand.args?.map(interpolateEnvVars),
+          env: interpolateEnvRecord(definition.requestHeadersCommand.env),
+          timeoutMs: definition.requestHeadersCommand.timeoutMs,
+        }
+      : undefined,
     auth: definition.auth,
     protocolVersion: definition.protocolVersion,
     bearerToken: resolveBearerToken(definition),
@@ -184,8 +193,10 @@ export function reconstructToolMetadata(
   const metadata: ToolMetadata[] = [];
   const seenNames = new Set<string>();
   const effectivePrefix = resolveToolPrefix(definition, prefix);
-  const getOtherCurrentCandidates = (toolName: string): Set<string> | undefined => {
-    if (!configuredServers || !cache) return undefined;
+  const hasToolFilters =
+    (Array.isArray(definition.includeTools) && definition.includeTools.length > 0) ||
+    (Array.isArray(definition.excludeTools) && definition.excludeTools.length > 0);
+  const selectorCandidateIndex = hasToolFilters && configuredServers && cache ? (() => {
     const candidates = new Set<string>();
     for (const [otherServerName, otherDefinition] of Object.entries(configuredServers)) {
       const otherEntry = cache.servers[otherServerName];
@@ -202,16 +213,15 @@ export function reconstructToolMetadata(
         }
       }
     }
-    for (const candidate of getToolNameCandidates(toolName, serverName, effectivePrefix, false)) candidates.delete(candidate);
-    return candidates;
-  };
+    return createToolSelectorCandidateIndex(candidates);
+  })() : undefined;
 
   for (const tool of entry.tools ?? []) {
     if (!tool?.name) continue;
     if (!isUiToolVisibleToModel(tool.uiVisibility)) {
       continue;
     }
-    if (!isToolAllowed(tool.name, serverName, effectivePrefix, definition.includeTools, definition.excludeTools, getOtherCurrentCandidates(tool.name))) {
+    if (!isToolAllowed(tool.name, serverName, effectivePrefix, definition.includeTools, definition.excludeTools, selectorCandidateIndex)) {
       continue;
     }
 
@@ -236,7 +246,7 @@ export function reconstructToolMetadata(
     for (const resource of entry.resources ?? []) {
       if (!resource?.name || !resource?.uri) continue;
       const baseName = `read_${resourceNameToToolName(resource.name)}`;
-      if (!isToolAllowed(baseName, serverName, effectivePrefix, definition.includeTools, definition.excludeTools, getOtherCurrentCandidates(baseName))) {
+      if (!isToolAllowed(baseName, serverName, effectivePrefix, definition.includeTools, definition.excludeTools, selectorCandidateIndex)) {
         continue;
       }
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-mcp-adapter",
-  "version": "2.25.0",
+  "version": "2.26.1",
   "description": "MCP (Model Context Protocol) adapter extension for Pi coding agent",
   "type": "module",
   "types": "./index.ts",
@@ -89,6 +89,7 @@
     "ui-app-bridge-helpers.ts",
     "config.ts",
     "server-manager.ts",
+    "request-headers-command.ts",
     "unix-socket-transport.ts",
     "json-schema-validator.ts",
     "session-recovery.ts",