@herbertgao/pi-extensions 2026.8.6 → 2026.8.8

package/node_modules/pi-mcp-adapter/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);
@@ -45,6 +73,14 @@ export class McpLifecycleManager {
     if (settings?.idleTimeout !== undefined) this.serverSettings.set(name, settings);
   }
 
+  unregisterServer(name: string): void {
+    this.allServers.delete(name);
+    this.keepAliveServers.delete(name);
+    this.serverSettings.delete(name);
+    this.retryStates.delete(name);
+    this.pendingMetadataPublications.delete(name);
+  }
+
   setGlobalIdleTimeout(minutes: number): void {
     this.globalIdleTimeout = minutes * 60 * 1000;
   }
@@ -82,29 +118,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 +147,255 @@ 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;
+    // Fence against unregisterServer racing an in-flight convergence pass.
+    // Identity comparison also rejects stale passes after a same-name
+    // replacement registration, which a name check would wrongly accept.
+    if (this.keepAliveServers.get(name) !== definition) 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, definition, error, "reconnect", this.manager.getConnection(name));
+        return;
+      }
+      if (this.stopped || signal?.aborted) return;
+      if (freshConnection.status === "needs-auth") {
+        await this.notifyAuthRequired(name, definition, freshConnection);
+        return;
+      }
+      if (freshConnection.status !== "connected") {
+        this.reportConnectionFailure(
+          name,
+          definition,
+          new Error(`MCP server ${name} did not return a connected session`),
+          "reconnect",
+          freshConnection,
+        );
+        return;
+      }
+      logger.debug(`Reconnected to ${name}`);
+      await this.publishConnectedMetadata(name, definition, freshConnection);
+      return;
+    }
+
+    if (this.pendingMetadataPublications.has(name)) {
+      await this.publishConnectedMetadata(name, definition, 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, definition, 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, definition, reconnectError, "reconnect", this.manager.getConnection(name));
+        return;
+      }
+      if (this.stopped || signal?.aborted) return;
+      if (freshConnection.status === "needs-auth") {
+        await this.notifyAuthRequired(name, definition, freshConnection);
+        return;
+      }
+      if (freshConnection.status !== "connected") {
+        this.reportConnectionFailure(
+          name,
+          definition,
+          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, definition, freshConnection);
+      return;
+    }
+
+    if (refreshResult === "superseded") {
+      await this.handleSupersededConnection(name, definition, connection, signal, retrySuperseded);
+      return;
+    }
+    if (refreshResult === "refresh-timeout") {
+      this.deferRefreshTimeout(name, definition, connection);
+      return;
+    }
+    if (this.keepAliveServers.get(name) !== definition) 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 (this.keepAliveServers.get(name) !== definition) return;
+    if (current === staleConnection && current.status === "connected") {
+      if (this.retryStates.delete(name)) await this.onHealthRestored?.(name);
+      return;
+    }
+    if (current?.status === "connected") {
+      await this.publishConnectedMetadata(name, definition, current);
+      return;
+    }
+    if (current?.status === "needs-auth") {
+      await this.notifyAuthRequired(name, definition, current);
+      return;
+    }
+    if (retrySuperseded) {
+      await this.checkKeepAliveConnection(name, definition, signal, false);
+    }
+  }
+
+  private async publishConnectedMetadata(name: string, definition: ServerDefinition, connection: ServerConnection): Promise<void> {
+    // Fence stale convergence passes by definition identity so disposal, and
+    // disposal followed by a same-name replacement, never re-track this
+    // server. Close the pass's own connection only when it is still current,
+    // so a replacement's connection is never touched.
+    if (this.keepAliveServers.get(name) !== definition) {
+      if (this.manager.getConnection(name) === connection) await this.manager.close(name);
+      return;
+    }
+    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, definition, error, "publish", connection);
+    }
+  }
+
+  private async notifyAuthRequired(name: string, definition: ServerDefinition, connection: ServerConnection): Promise<void> {
+    // Fence stale convergence passes by definition identity. Close the
+    // leftover needs-auth connection only while it is still the manager's
+    // current entry, so a replacement's connection is never touched and the
+    // stale client/transport cannot leak when a replacement connects later.
+    if (this.keepAliveServers.get(name) !== definition) {
+      if (this.manager.getConnection(name) === connection) await this.manager.close(name);
+      return;
+    }
+    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,
+    definition: ServerDefinition,
+    error: unknown,
+    action: "refresh" | "reconnect" | "publish",
+    connection: ServerConnection | undefined,
+  ): void {
+    if (!this.recordRetry(name, definition, connection)) return;
+    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 deferRefreshTimeout(
+    name: string,
+    definition: ServerDefinition,
+    connection: ServerConnection | undefined,
+  ): void {
+    if (!this.recordRetry(name, definition, connection)) return;
+    logger.debug(`MCP: keep-alive tools/list refresh timed out for ${name}; retrying after backoff`);
+  }
+
+  private recordRetry(
+    name: string,
+    definition: ServerDefinition,
+    connection: ServerConnection | undefined,
+  ): boolean {
+    // Do not recreate retry/failure state from a stale convergence pass after
+    // disposal or a same-name replacement registration.
+    if (this.keepAliveServers.get(name) !== definition) return false;
+    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,
+    });
+    return true;
+  }
+
   private getIdleTimeout(name: string): number {
     const perServer = this.serverSettings.get(name)?.idleTimeout;
     if (perServer !== undefined) return perServer * 60 * 1000;
@@ -137,11 +416,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/node_modules/pi-mcp-adapter/mcp-auth-flow.ts

@@ -18,6 +18,7 @@ import {
   waitForCallback,
   cancelPendingCallback,
   stopCallbackServer,
+  stopCallbackServerIfIdle,
   releaseCallbackServer,
 } from "./mcp-callback-server.ts"
 import {
@@ -387,7 +388,7 @@ export async function startAuth(
   } catch (error) {
     releaseCallbackServer(oauthState)
     try {
-      await clearOAuthState(serverName, authStorageOptions)
+      await cleanupAndReleaseCallbackServerIfIdle(() => clearOAuthState(serverName, authStorageOptions))
     } catch (cleanupError) {
       throw new AggregateError([error, cleanupError], "OAuth startup cleanup failed")
     }
@@ -429,6 +430,7 @@ export async function startAuth(
       authProvider.deactivate()
       releaseCallbackServer(oauthState)
       await clearOAuthState(serverName, authStorageOptions)
+      await stopCallbackServerIfIdle()
       return { authorizationUrl: "" }
     }
     if (!capturedUrl) {
@@ -439,7 +441,7 @@ export async function startAuth(
   } catch (error) {
     authProvider.deactivate()
     try {
-      await clearPendingAuth(runtime, serverName, oauthState, authStorageOptions)
+      await clearPendingAuthAndReleaseIfIdle(runtime, serverName, oauthState, authStorageOptions)
     } catch (cleanupError) {
       throw new AggregateError([error, cleanupError], "OAuth startup cleanup failed")
     }
@@ -463,7 +465,7 @@ async function setPendingAuth(
   state.pendingAuths.set(key, pendingAuth)
   state.pendingAuthStates.set(key, oauthState)
   const cleanupTimer = setTimeout(() => {
-    void clearPendingAuth(runtime, serverName, oauthState, pendingAuth.authStorageOptions).catch(error => {
+    void clearPendingAuthAndReleaseIfIdle(runtime, serverName, oauthState, pendingAuth.authStorageOptions).catch(error => {
       console.error(`MCP Auth: Timed-out flow cleanup failed: ${formatTerminalError(error)}`)
     })
   }, MANUAL_AUTH_TIMEOUT_MS)
@@ -498,6 +500,37 @@ async function clearPendingAuth(runtime: McpOAuthRuntime, serverName: string, oa
   }
 }
 
+async function clearPendingAuthAndReleaseIfIdle(
+  runtime: McpOAuthRuntime,
+  serverName: string,
+  oauthState: string | undefined,
+  fallbackStorageOptions: AuthStorageOptions = {},
+): Promise<void> {
+  await cleanupAndReleaseCallbackServerIfIdle(
+    () => clearPendingAuth(runtime, serverName, oauthState, fallbackStorageOptions),
+  )
+}
+
+async function cleanupAndReleaseCallbackServerIfIdle(cleanup: () => void | Promise<void>): Promise<void> {
+  let cleanupFailure: { error: unknown } | undefined
+  try {
+    await cleanup()
+  } catch (error) {
+    cleanupFailure = { error }
+  }
+
+  try {
+    await stopCallbackServerIfIdle()
+  } catch (releaseError) {
+    if (cleanupFailure) {
+      throw new AggregateError([cleanupFailure.error, releaseError], "OAuth callback cleanup failed")
+    }
+    throw releaseError
+  }
+
+  if (cleanupFailure) throw cleanupFailure.error
+}
+
 function getSearchParamsFromInput(input: string): URLSearchParams | undefined {
   try {
     const url = new URL(input)
@@ -701,7 +734,7 @@ export async function completeAuth(
   } finally {
     if (!keepPendingForRetry) {
       try {
-        await clearPendingAuth(runtime, serverName, oauthState, authStorageOptions)
+        await clearPendingAuthAndReleaseIfIdle(runtime, serverName, oauthState, authStorageOptions)
       } catch (cleanupError) {
         if (caughtError !== undefined) {
           throw new AggregateError([caughtError, cleanupError], "OAuth completion cleanup failed")
@@ -801,7 +834,7 @@ export async function authenticate(
     } catch (error) {
       if (oauthState) cancelPendingCallback(oauthState)
       try {
-        await clearPendingAuth(runtime, serverName, oauthState, authStorageOptions)
+        await clearPendingAuthAndReleaseIfIdle(runtime, serverName, oauthState, authStorageOptions)
       } catch (cleanupError) {
         throw new AggregateError([error, cleanupError], "OAuth cancellation cleanup failed")
       }
@@ -922,7 +955,7 @@ export async function removeAuth(serverName: string, options: AuthenticateOption
   if (oauthState) {
     cancelPendingCallback(oauthState)
   }
-  await clearPendingAuth(runtime, serverName, oauthState, authStorageOptions)
+  await clearPendingAuthAndReleaseIfIdle(runtime, serverName, oauthState, authStorageOptions)
   throwIfAborted(signal)
   clearAllCredentials(serverName, authStorageOptions)
   await clearOAuthState(serverName, authStorageOptions)