pi-mcp-extension/src/server-manager.ts

Full-featured, lightweight MCP extension for the Pi coding agent: stdio, SSE & Streamable HTTP transports, OAuth 2.1 (PKCE + DCR), lazy/eager lifecycles — minimal dependencies, zero bloat.

/**
 * Server lifecycle manager for pi-mcp.
 *
 * Manages MCP server connections using the official SDK.
 * Deliberately thin: the SDK handles protocol state, transport, and process lifecycle,
 * `src/mcp-client.ts` handles the stateless client/ping adapter, and this module owns only
 * what needs a state machine:
 *   - 3-state lifecycle per server (stopped / starting / ready)
 *   - Retry with a fixed delay schedule
 *   - Background recovery on transport close, stream error, or a failed health-check ping
 *   - Stderr capture (circular buffer)
 *   - PID tracking for safety-net SIGKILL on shutdown failure
 */

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import { AuthCancelledError, AuthRequiredError, McpError } from "./errors.js";
import type { McpConfig, ServerConfig, Settings } from "./config.js";
import { getAuthStatus, invalidateStoredTokens, resetAuth } from "./oauth-status.js";
import type { ToolRefreshOptions } from "./tool-bridge.js";
import { redactSecrets } from "./secret-redaction.js";
import { sanitizeTerminalText } from "./terminal-text.js";
import { describeError, sanitizeLogText } from "./debug-log-format.js";
import { reportDiagnostic } from "./debug-log.js";
import {
  ensurePidTerminated,
  instrumentStdioTransport,
  resolveTransportPid,
  terminateStdioTree,
} from "./stdio-process.js";
import { createTransport } from "./transport-factory.js";
import {
  classifyClientError,
  classifyConnectFailure,
  type ClientErrorAction,
  type ConnectFailureAction,
} from "./connect-policy.js";
import {
  connectClientWithTimeout,
  createMcpClient,
  startHealthCheck,
  type FormattedServerLog,
} from "./mcp-client.js";

/**
 * Strip carriage returns, ANSI CSI/OSC/DCS/PM/APC escape sequences, and C0/C1
 * control characters from a single server stderr line before storing or logging.
 */
export function sanitizeStderrLine(line: string): string {
  return sanitizeTerminalText(line);
}

/** Sanitized message text of a client/transport error, for diagnostics. */
function clientErrorMessage(err: unknown): string {
  return sanitizeTerminalText(err instanceof Error ? err.message : String(err));
}

// ─── Types ────────────────────────────────────────────────────────────────────

export type ServerState = "stopped" | "starting" | "ready";



/** Maximum stderr lines stored per server (circular). */
const STDERR_BUFFER_SIZE = 100;

export interface ManagedServer {
  name: string;
  config: ServerConfig;
  state: ServerState;
  client: Client | null;
  transport?: Transport | undefined;
  cwd?: string | undefined;
  /** PID of the child process (stdio transport only). Used for safety-net cleanup. */
  childPid: number | null;
  retryCount: number;
  lastError: Error | null;
  /** Recent stderr lines from the server subprocess. */
  stderrLog: string[];
  healthCheckTimer: ReturnType<typeof setInterval> | null;
  /** Pending retry timeout — cleared on shutdown to prevent ghost reconnects. */
  retryTimer: ReturnType<typeof setTimeout> | null;
  attempt?: Promise<void> | undefined;
  controller?: AbortController | undefined;
  generation?: number;
  /** Timestamp of the last state transition — used for the debug log's `ms` field. */
  stateSince?: number | undefined;
}

/** Called after tool list is refreshed for a server (e.g. on list_changed notification). */
export type ToolRefreshCallback = (
  serverName: string,
  client: Client,
  options?: ToolRefreshOptions,
) => Promise<void>;

/** Called when background auto-recovery after a disconnect or health check fails. */
export type RecoveryFailedCallback = (serverName: string, error: Error) => void;

/**
 * Called every time a server reaches `ready`, whether the transition came from an
 * explicit connect, an eager startup, or background `_recover()`.
 *
 * WHY: Readiness is the only state change that makes a previously declined
 * `run /mcp:auth <name>` footer hint obsolete, and eager starts plus background
 * recovery bypass `ConnectionCoordinator.connect()`, which is where the hint is
 * otherwise repainted.
 */
export type ServerReadyCallback = (serverName: string) => void;

interface HandshakeParams {
  server: ManagedServer;
  client: Client;
  transport: Transport;
  timeoutMs: number;
  signal: AbortSignal;
  currentCwd: string;
}

interface ConnectFailureContext {
  server: ManagedServer;
  client: Client;
  transport: Transport;
  cwd: string;
  signal: AbortSignal;
  corruptStateRecovered: boolean;
  err: unknown;
}

// ─── ServerManager ────────────────────────────────────────────────────────────

export class ServerManager {
  private readonly servers = new Map<string, ManagedServer>();
  private settings: Settings;
  private onToolRefresh: ToolRefreshCallback | null = null;
  private onResourceListChanged: ((name: string) => void) | null = null;
  private onServerStopped: ((name: string) => void) | null = null;
  private onRecoveryFailed: RecoveryFailedCallback | null = null;
  private onServerReady: ServerReadyCallback | null = null;

  constructor(config: McpConfig, private readonly storageDir?: string) {
    this.settings = config.settings;
    for (const [name, serverConfig] of Object.entries(config.mcpServers)) {
      this.setServer(name, serverConfig);
    }
  }

  /** Register callback invoked after tool list changes for a server. */
  setToolRefreshCallback(cb: ToolRefreshCallback): void {
    this.onToolRefresh = cb;
  }

  /** Register callback invoked when a server reports `notifications/resources/list_changed`. */
  setResourceListChangedCallback(cb: (name: string) => void): void {
    this.onResourceListChanged = cb;
  }

  setServerStoppedCallback(cb: (name: string) => void): void {
    this.onServerStopped = cb;
  }

  setRecoveryFailedCallback(cb: RecoveryFailedCallback): void {
    this.onRecoveryFailed = cb;
  }

  /** Register callback invoked whenever a server reaches the `ready` state. */
  setServerReadyCallback(cb: ServerReadyCallback): void {
    this.onServerReady = cb;
  }

  /** Effective request timeout for a server: its own override, else the session default. */
  resolveRequestTimeoutMs(server: ManagedServer): number {
    return server.config.requestTimeoutMs ?? this.settings.requestTimeoutMs;
  }

  /**
   * Move a server to a new lifecycle state and record the transition.
   *
   * The transition is reported to the debug log (never to stderr while a UI is
   * active) so maintainers can see when a server started, became ready, or
   * was torn down.
   */
  private _setState(server: ManagedServer, next: ServerState, reason?: string): void {
    const from = server.state;
    if (from === next) return;
    const now = Date.now();
    const ms = server.stateSince !== undefined ? now - server.stateSince : undefined;
    server.state = next;
    server.stateSince = now;
    reportDiagnostic({
      level: "info",
      event: "server.state",
      server: server.name,
      message: `Server ${server.name} state ${from} -> ${next}`,
      fields: { from, to: next, ms, reason },
    });
  }

  createTransport(name: string, config: ServerConfig, cwd?: string): Transport {
    return createTransport(name, config, cwd, this.storageDir, cwd);
  }

  getServer(name: string): ManagedServer | undefined {
    return this.servers.get(name);
  }

  getServerStatus(name: string): ManagedServer | undefined {
    return this.servers.get(name);
  }

  getAllServers(): ManagedServer[] {
    return Array.from(this.servers.values());
  }

  getReadyServers(): ManagedServer[] {
    return this.getAllServers().filter((s) => s.state === "ready");
  }

  private _updateServerCwd(server: ManagedServer, cwd: string): void {
    const prevCwd = server.cwd;
    server.cwd = cwd;
    if (prevCwd !== undefined && prevCwd !== cwd && server.state === "ready" && server.client) {
      void server.client.sendRootsListChanged().catch(() => {
        // WHY: Best-effort notification; some MCP servers ignore or do not support roots/list_changed.
      });
    }
  }

  /** Update workspace cwd for all managed servers and notify ready servers if cwd changed. */
  updateCwd(cwd: string): void {
    for (const server of this.servers.values()) {
      this._updateServerCwd(server, cwd);
    }
  }

  /** Status summary for /mcp command. */
  getStatusSummary(): string {
    const all = this.getAllServers();
    if (all.length === 0) return "pi-mcp: No servers configured (create ~/.pi/agent/mcp.json)";
    const lines = all.map((s) => {
      const icon = s.state === "ready" ? "✓" : s.state === "starting" ? "⟳" : "✗";
      const err = s.lastError ? ` — ${sanitizeLogText(s.lastError.message, 500)}` : "";
      return `  ${icon} ${s.name} (${s.state})${err}`;
    });
    const ready = all.filter((s) => s.state === "ready").length;
    return [`MCP: ${ready}/${all.length} servers ready`, ...lines].join("\n");
  }

  /** Reset OAuth credentials for a server, forcing re-authorization on next connect. */
  async resetServerAuth(name: string): Promise<void> {
    await resetAuth(name, this.storageDir);
  }

  /** Get auth status for a server. */
  async getServerAuthStatus(name: string): Promise<{
    hasTokens: boolean;
    hasClientInfo: boolean;
    savedAt: string | undefined;
    scope: string | undefined;
  } | null> {
    const serverUrl = this.servers.get(name)?.config.url;
    return getAuthStatus(name, this.storageDir, serverUrl);
  }

  /**
    * Mark an OAuth server as requiring re-authorization when a mid-session
    * request (tools/list or tools/call) fails with 401/403 insufficient_scope.
    *
    * Returns the {@link AuthRequiredError} when a server that was already `ready`
    * lost its credentials mid-session, and `undefined` otherwise (unknown server,
    * non-OAuth server, or a server that is not `ready` — see below).
    *
    * Only a `ready` server is handled here. For any other state this is a no-op and
    * must NOT touch the server: a 401/403 during a connect attempt (`initialize` ok,
    * `tools/list` 401 while `starting`) arrives while that attempt owns the
    * lifecycle, and shutting the server down here would abort the attempt's own
    * signal. `classifyConnectFailure` then sees `aborted` first and reports
    * `AuthCancelledError` instead of `auth-required`, which suppresses the startup
    * auth prompt, the interactive authorization, and the recovery failure alike.
    * Left alone, the attempt's `_raiseAuthRequired` invalidates the tokens, records
    * the `AuthRequiredError`, and reports it exactly once.
    */
  async handleUnauthorizedServerError(name: string): Promise<AuthRequiredError | undefined> {
    const server = this.servers.get(name);
    if (!server || !server.config.auth || server.config.transport === "stdio") {
      return undefined;
    }
    if (server.state !== "ready") {
      // WHY: See the JSDoc above — an in-flight connect attempt reports its own 401.
      return undefined;
    }
    const authError = new AuthRequiredError(name);
    // WHY: Shut down immediately before async disk I/O so server.generation increments
    // synchronously and cancels any concurrent _recover() microtask queued by client.onerror.
    await this._shutdown(server);
    server.lastError = authError;
    await invalidateStoredTokens(name, this.storageDir, server.config.url);
    server.lastError = authError;
    return authError;
  }

  /** Get recent stderr output for a server. */
  getServerLogs(name: string): string {
    const server = this.servers.get(name);
    if (!server) return `No server named "${name}"`;
    if (server.stderrLog.length === 0) return `(no stderr output from ${name})`;
    return server.stderrLog.join("\n");
  }

  // ─── Lifecycle ──────────────────────────────────────────────────────────────

  /**
   * Start a server and connect to it.
   * cwd is passed to roots/list — the workspace root exposed to the MCP server.
   */
  async startServer(name: string, cwd: string): Promise<void> {
    const server = this.servers.get(name);
    if (!server) {
      throw new McpError(`Unknown server "${name}"`, name, "config");
    }
    this._updateServerCwd(server, cwd);
    if (server.attempt && !server.controller?.signal.aborted) return server.attempt;
    if (server.state === "ready") return;
    server.retryCount = 0;
    const controller = new AbortController();
    server.controller = controller;
    this._setState(server, "starting");
    const attempt = this._connect(server, cwd, controller.signal);
    server.attempt = attempt;
    try {
      await attempt;
    } finally {
      if (server.attempt === attempt) {
        server.attempt = undefined;
      }
      if (server.controller === controller && (server.state as ServerState) !== "ready") {
        server.controller = undefined;
      }
    }
  }

  async stopServer(name: string): Promise<void> {
    const server = this.servers.get(name);
    if (!server) return;
    await this._shutdown(server);
  }

  async shutdownAll(): Promise<void> {
    await Promise.allSettled(
      Array.from(this.servers.values()).map((s) => this._shutdown(s)),
    );
  }

  /** Add or replace a stopped server entry in the manager. */
  setServer(name: string, serverConfig: ServerConfig): void {
    this.servers.set(name, {
      name,
      config: serverConfig,
      state: "stopped",
      client: null,
      childPid: null,
      retryCount: 0,
      lastError: null,
      stderrLog: [],
      healthCheckTimer: null,
      retryTimer: null,
      stateSince: Date.now(),
    });
  }

  /** Shut down and remove a server entry from the manager. */
  async removeServer(name: string): Promise<void> {
    const server = this.servers.get(name);
    if (!server) return;
    await this._shutdown(server);
    // WHY: _shutdown awaits child-process teardown, and the registration reconcile installs a
    // replacement right after removing the changed entry — it can land in that window. Deleting
    // by name alone would drop that live entry and orphan its stdio child, which nothing can
    // then terminate.
    if (this.servers.get(name) === server) {
      this.servers.delete(name);
    }
  }

  /**
   * Rebuild the server map from a new config.
   * Must only be called after shutdownAll() — old servers are discarded.
   * New servers that didn't exist before are added; servers removed from
   * config are dropped (their tools should already be deactivated).
   */
  rebuildServers(config: McpConfig): void {
    this.settings = config.settings;
    this.servers.clear();
    for (const [name, serverConfig] of Object.entries(config.mcpServers)) {
      this.setServer(name, serverConfig);
    }
  }

  /** Alias for rebuildServers when replacing configuration on session_start. */
  replaceConfig(config: McpConfig): void {
    this.rebuildServers(config);
  }

  // ─── Internal ───────────────────────────────────────────────────────────────

  private _appendStderr(server: ManagedServer, line: string): void {
    // WHY: this buffer is what `/mcp <name>` prints as "Recent output" and what
    // `getServerLogs` returns — text users paste into bug reports. A child server that echoes
    // its own key on stderr must not surface it there, so mask credentials exactly as the
    // debug log does. Redacting twice is a no-op.
    server.stderrLog.push(redactSecrets(line));
    if (server.stderrLog.length > STDERR_BUFFER_SIZE) {
      server.stderrLog.shift();
    }
  }

  private _createAndInstrumentTransport(server: ManagedServer, cwd: string): Transport {
    let transport: Transport;
    try {
      transport = this.createTransport(server.name, server.config, cwd);
    } catch (err) {
      this._setState(server, "stopped", "transport");
      server.lastError = err instanceof Error ? err : new Error(String(err));
      throw new McpError(
        `Failed to create transport: ${server.lastError.message}`,
        server.name,
        "connection",
        err,
      );
    }

    if (transport instanceof StdioClientTransport) {
      instrumentStdioTransport(transport, {
        onPid: (pid) => {
          server.childPid = pid;
        },
        onStderrLine: (rawLine) => {
          const clean = sanitizeStderrLine(rawLine);
          if (!clean) return;
          this._appendStderr(server, clean);
          reportDiagnostic({
            level: "info",
            event: "server.stderr",
            server: server.name,
            message: clean,
            // WHY: child-server stderr historically went to stderr; keep it
            // visible in headless runs, where there is no /mcp ring buffer UI.
            stderr: true,
          });
        },
      });
    }

    return transport;
  }

  private _createClient(
    server: ManagedServer,
    cwd: string,
    signal: AbortSignal,
  ): Client {
    return createMcpClient({
      serverName: server.name,
      // WHY: read `server.cwd` lazily: `updateCwd()` can move the workspace root while the
      // handshake is in flight, and roots/list must answer with the current one.
      getCwd: () => server.cwd ?? cwd,
      onToolsChanged: (client) => this._handleToolListChanged(server, client, signal),
      onResourcesChanged: (client) => this._handleResourceListChanged(server, client),
      onLogLine: (entry) => this._recordServerLog(server, entry),
    });
  }

  /**
   * Record one `notifications/message` entry from a server.
   *
   * WHY: server log messages historically went to stderr; keep them visible in headless runs,
   * where there is no /mcp ring buffer UI, and store the sanitized line for `/mcp <name>`.
   */
  private _recordServerLog(server: ManagedServer, entry: FormattedServerLog): void {
    reportDiagnostic({
      level: "info",
      event: "server.log",
      server: server.name,
      message: entry.formatted,
      fields: { logLevel: entry.safeLevel, logger: entry.safeLogger },
      stderr: true,
    });
    this._appendStderr(server, entry.formatted);
  }

  /**
   * Re-discover tools after a `notifications/tools/list_changed` from a ready server.
   *
   * Extracted from the notification handler so the handler registration stays a single
   * guarded call.
   */
  private async _handleToolListChanged(
    server: ManagedServer,
    client: Client,
    signal: AbortSignal,
  ): Promise<void> {
    const onToolRefresh = this.onToolRefresh;
    if (signal.aborted || !onToolRefresh || server.state !== "ready" || server.client !== client) {
      return;
    }
    try {
      const refreshTimeoutMs = this.resolveRequestTimeoutMs(server);
      await onToolRefresh(server.name, server.client, {
        requestTimeoutMs: refreshTimeoutMs,
        signal,
      });
    } catch (err) {
      const { kind, message } = describeError(err);
      // WHY: Two failures here are already accounted for and must not toast: a cancelled
      // refresh (/mcp:stop, session teardown) is expected, and a 401/403 was just reported
      // by handleUnauthorizedServerError as a re-authorization warning plus the footer hint
      // — a second "Failed to refresh tools" error would read as an unrelated problem.
      const alreadyReported =
        err instanceof AuthCancelledError ||
        signal.aborted ||
        server.lastError instanceof AuthRequiredError;
      reportDiagnostic({
        level: alreadyReported ? "info" : "error",
        event: "server.tools_refresh_failed",
        server: server.name,
        message: `Failed to refresh tools for ${server.name}: ${message}`,
        fields: { kind },
        notify: !alreadyReported,
      });
    }
  }

  /**
   * Invalidate a server's cached resource catalog after a
   * `notifications/resources/list_changed`.
   *
   * Only the current client of a ready server may invalidate: a notification from a
   * superseded client (reconnect race) or from a server that is being torn down must not
   * touch the cache the live connection owns.
   */
  private _handleResourceListChanged(server: ManagedServer, client: Client): void {
    if (server.state !== "ready" || server.client !== client) return;
    this.onResourceListChanged?.(server.name);
  }

  private async _handshakeAndDiscover(params: HandshakeParams): Promise<void> {
    const { server, client, transport, timeoutMs, signal, currentCwd } = params;
    await connectClientWithTimeout({
      client,
      transport,
      timeoutMs,
      signal,
      serverName: server.name,
    });
    if (signal.aborted) throw new AuthCancelledError(server.name);
    if (transport instanceof StdioClientTransport && !server.childPid) {
      server.childPid = transport.pid;
    }
    // Readiness includes tool discovery, not just the MCP handshake.
    if (this.onToolRefresh) {
      await this.onToolRefresh(server.name, client, {
        requestTimeoutMs: timeoutMs,
        signal,
      });
    }
    if (signal.aborted) throw new AuthCancelledError(server.name);
    this._setState(server, "ready");
    server.retryCount = 0;
    server.lastError = null;
    // WHY: Fire only after the server is observably ready with a cleared lastError, so a
    // listener that repaints session state never sees a half-transitioned server.
    this.onServerReady?.(server.name);
    if (server.cwd && server.cwd !== currentCwd) {
      void client.sendRootsListChanged().catch(() => {
        // WHY: Best-effort notification if cwd was updated while handshake was in flight.
      });
    }
  }

  private _startHealthCheck(
    server: ManagedServer,
    timeoutMs: number,
    signal: AbortSignal,
    cwd: string,
  ): void {
    server.healthCheckTimer = startHealthCheck({
      serverName: server.name,
      intervalMs: server.config.healthCheckIntervalMs ?? 0,
      timeoutMs,
      signal,
      getClient: () => server.client,
      isReady: () => server.state === "ready",
      // WHY: Ping failure triggers background recovery if the server is still ready.
      // The ping error is intentionally not kept on lastError because _recover() →
      // _shutdown() synchronously resets lastError during teardown before any reader
      // observes it; if reconnection fails, _recover() records the failure there.
      onPingFailure: () => {
        void this._recover(server, server.cwd ?? cwd);
      },
    });
  }

  // WHY: SSEClientTransport and StreamableHTTPClientTransport report unexpected
  // stream drops via transport.onerror (forwarded to client.onerror) rather than
  // transport.onclose. Crucially, StreamableHTTPClientTransport.send() invokes
  // this.onerror(error) synchronously BEFORE `throw error` rejects the send()
  // promise; deferring _recover() via setImmediate lets Protocol.request() and
  // ToolBridge.execute() observe the original transport/401 error (and run
  // handleUnauthorizedServerError if applicable) before client.close() rejects
  // pending requests with "Connection closed".
  private _scheduleErrorRecovery(
    server: ManagedServer,
    client: Client,
    signal: AbortSignal,
    cwd: string,
  ): void {
    const errorGeneration = server.generation ?? 0;
    setImmediate(() => {
      if (
        (server.generation ?? 0) !== errorGeneration ||
        server.state !== "ready" ||
        signal.aborted ||
        server.client !== client
      ) {
        return;
      }
      void this._recover(server, server.cwd ?? cwd);
    });
  }

  private _handleClientError(
    server: ManagedServer,
    client: Client,
    signal: AbortSignal,
    cwd: string,
    err: unknown,
  ): void {
    if (server.state !== "ready" || signal.aborted || server.client !== client) {
      return;
    }
    const action = classifyClientError(server.config.transport, err);
    if (action.kind === "ignore") {
      this._reportIgnoredClientError(server, action, err);
      return;
    }
    // WHY: deferred recovery — see the note on _scheduleErrorRecovery.
    reportDiagnostic({
      level: "warn",
      event: action.event,
      server: server.name,
      message: `Server ${server.name} ${action.reason}: ${clientErrorMessage(err)}`,
      fields: { transport: action.transport },
    });
    this._scheduleErrorRecovery(server, client, signal, cwd);
  }

  /**
   * Log the one ignorable client error that is worth a line: a server notification whose
   * handler threw (`notifications/progress` with malformed parameters, typically).
   */
  private _reportIgnoredClientError(
    server: ManagedServer,
    action: Extract<ClientErrorAction, { kind: "ignore" }>,
    err: unknown,
  ): void {
    if (action.diagnostic !== "server.notification_error") return;
    reportDiagnostic({
      level: "info",
      event: "server.notification_error",
      server: server.name,
      message: clientErrorMessage(err),
      fields: { transport: server.config.transport },
    });
  }

  private _wireDisconnectRecovery(
    server: ManagedServer,
    client: Client,
    signal: AbortSignal,
    cwd: string,
  ): void {
    // Unexpected transport close (e.g. the server subprocess crashes
    // mid-session). Without this the state would stay "ready" and tools
    // would remain active against a dead server. Deliberate closes abort
    // the signal first, so they are ignored here.
    client.onclose = () => {
      if (signal.aborted || server.client !== client) return;
      reportDiagnostic({
        level: "warn",
        event: "server.disconnected",
        server: server.name,
        message: `Server ${server.name} closed unexpectedly, reconnecting`,
        fields: { transport: server.config.transport },
      });
      void this._recover(server, server.cwd ?? cwd);
    };

    client.onerror = (err) => {
      this._handleClientError(server, client, signal, cwd, err);
    };
  }

  private async _cleanupFailedConnect(
    server: ManagedServer,
    client: Client,
    transport: Transport,
  ): Promise<void> {
    const pid = resolveTransportPid(server.childPid, transport);
    server.childPid = null;
    await Promise.allSettled([
      client.close(),
      terminateStdioTree(pid, transport),
    ]);
    // WHY: Secondary safety check to guarantee PID termination even if transport.close()
    // or child process exit raced during the initial Promise.allSettled.
    await ensurePidTerminated(pid);
    if (server.client === client) {
      server.client = null;
      server.transport = undefined;
    }
  }

  private async _waitRetry(
    server: ManagedServer,
    signal: AbortSignal,
    delayMs: number,
  ): Promise<void> {
    await new Promise<void>((resolve) => {
      if (signal.aborted) {
        resolve();
        return;
      }
      const onAbort = () => {
        if (server.retryTimer) {
          clearTimeout(server.retryTimer);
          server.retryTimer = null;
        }
        resolve();
      };
      const timer = setTimeout(() => {
        signal.removeEventListener("abort", onAbort);
        server.retryTimer = null;
        resolve();
      }, delayMs);
      server.retryTimer = timer;
      signal.addEventListener("abort", onAbort, { once: true });
    });
  }

  private async _connect(
    server: ManagedServer,
    cwd: string,
    signal: AbortSignal,
    corruptStateRecovered = false,
  ): Promise<void> {
    server.lastError = null;
    const currentCwd = server.cwd ?? cwd;
    server.cwd = currentCwd;
    const timeoutMs = this.resolveRequestTimeoutMs(server);
    // Note: retryCount is NOT reset here — it's only reset after successful connect.
    // This ensures the retry limit is respected across reconnection attempts.
    // It IS reset when startServer() is called explicitly (e.g. /mcp:start)
    // via the caller having already set retryCount = 0 before calling _connect.

    const transport = this._createAndInstrumentTransport(server, currentCwd);
    server.transport = transport;

    const client = this._createClient(server, cwd, signal);
    server.client = client;

    try {
      await this._handshakeAndDiscover({ server, client, transport, timeoutMs, signal, currentCwd });
      this._startHealthCheck(server, timeoutMs, signal, cwd);
      this._wireDisconnectRecovery(server, client, signal, cwd);
    } catch (err) {
      await this._handleConnectFailure({
        server,
        client,
        transport,
        cwd,
        signal,
        corruptStateRecovered,
        err,
      });
    }
  }

  private async _handleConnectFailure(context: ConnectFailureContext): Promise<void> {
    const { server, client, transport, signal, corruptStateRecovered, err } = context;
    await this._cleanupFailedConnect(server, client, transport);

    const action = classifyConnectFailure({
      error: err,
      aborted: signal.aborted,
      corruptStateRecovered,
      hasAuth: Boolean(server.config.auth),
      transport: server.config.transport,
      retryCount: server.retryCount,
      maxRetries: this.settings.maxRetries,
      serverName: server.name,
    });

    this._reportConnectFailure(server, action.kind, err);
    return this._applyConnectFailureAction(context, action);
  }

  private _reportConnectFailure(
    server: ManagedServer,
    outcome: ConnectFailureAction["kind"],
    err: unknown,
  ): void {
    const { kind: errorKind, message: errorMessage } = describeError(err);
    reportDiagnostic({
      level: outcome === "cancelled" ? "info" : "warn",
      event: "server.connect_failure",
      server: server.name,
      message: `Connection attempt for ${server.name} failed: ${errorMessage}`,
      fields: { outcome, kind: errorKind, attempt: server.retryCount },
    });
  }

  private async _applyConnectFailureAction(
    context: ConnectFailureContext,
    action: ConnectFailureAction,
  ): Promise<void> {
    const { server, cwd, signal, corruptStateRecovered, err } = context;
    const error = err instanceof Error ? err : new Error(String(err));
    switch (action.kind) {
      case "cancelled":
        throw err instanceof AuthCancelledError ? err : new AuthCancelledError(server.name);
      case "retry-corrupt-state":
        return this._retryCorruptState(server, cwd, signal, err);
      case "auth-required":
        return this._raiseAuthRequired(server, err, action.invalidateTokens);
      case "fail":
        server.lastError = action.error;
        this._setState(server, "stopped", "failed");
        throw action.error;
      case "give-up":
        server.lastError = error;
        this._setState(server, "stopped", "give-up");
        throw action.error;
      case "retry":
        server.lastError = error;
        return this._scheduleReconnect(server, cwd, signal, corruptStateRecovered, action.delayMs);
    }
  }

  private _retryCorruptState(
    server: ManagedServer,
    cwd: string,
    signal: AbortSignal,
    err: unknown,
  ): Promise<void> {
    const { kind: errorKind, message: errorMessage } = describeError(err);
    reportDiagnostic({
      level: "warn",
      event: "oauth.state_corrupt",
      server: server.name,
      message: errorMessage,
      fields: { kind: errorKind },
    });
    if (signal.aborted) throw new AuthCancelledError(server.name);
    return this._connect(server, server.cwd ?? cwd, signal, true);
  }

  private async _raiseAuthRequired(
    server: ManagedServer,
    err: unknown,
    invalidateTokens: boolean,
  ): Promise<never> {
    server.lastError = err instanceof Error ? err : new Error(String(err));
    if (invalidateTokens) {
      await invalidateStoredTokens(server.name, this.storageDir, server.config.url);
    }
    const authError = err instanceof AuthRequiredError ? err : new AuthRequiredError(server.name);
    server.lastError = authError;
    this._setState(server, "stopped", "auth-required");
    throw authError;
  }

  private async _scheduleReconnect(
    server: ManagedServer,
    cwd: string,
    signal: AbortSignal,
    corruptStateRecovered: boolean,
    delayMs: number,
  ): Promise<void> {
    server.retryCount++;
    reportDiagnostic({
      level: "warn",
      event: "server.retry",
      server: server.name,
      message: `Retrying ${server.name} in ${delayMs}ms (attempt ${server.retryCount})`,
      fields: { attempt: server.retryCount, delayMs },
    });
    await this._waitRetry(server, signal, delayMs);
    if (signal.aborted) throw new AuthCancelledError(server.name);
    return this._connect(server, server.cwd ?? cwd, signal, corruptStateRecovered);
  }

  /** Tear down a server that died or stopped responding, then reconnect. */
  private async _recover(server: ManagedServer, cwd: string): Promise<void> {
    const recoveryGeneration = (server.generation ?? 0) + 1;
    await this._shutdown(server);
    if (server.generation === recoveryGeneration) {
      // An external stop during cleanup increments the generation again.
      queueMicrotask(() => {
        if (server.generation !== recoveryGeneration) return;
        void this.startServer(server.name, cwd).catch((err) => {
          if (err instanceof AuthCancelledError) return;
          const error = err instanceof Error ? err : new Error(String(err));
          server.lastError = error;
          const { kind, message } = describeError(err);
          reportDiagnostic({
            level: "error",
            event: "server.recovery_failed",
            server: server.name,
            message: `Background recovery failed for ${server.name}: ${message}`,
            fields: { kind },
          });
          this.onRecoveryFailed?.(server.name, error);
        });
      });
    }
  }

  private async _shutdown(server: ManagedServer): Promise<void> {
    server.controller?.abort();
    server.controller = undefined;
    server.attempt = undefined;
    server.generation = (server.generation ?? 0) + 1;

    // Cancel pending retry
    if (server.retryTimer) {
      clearTimeout(server.retryTimer);
      server.retryTimer = null;
    }

    // Stop health check
    if (server.healthCheckTimer) {
      clearInterval(server.healthCheckTimer);
      server.healthCheckTimer = null;
    }

    this._setState(server, "stopped", "shutdown");
    server.lastError = null;
    this.onServerStopped?.(server.name);
    const client = server.client;
    // WHY: shared with _cleanupFailedConnect so a PID only visible through the transport
    // (childPid already cleared, `_process` gone after close) is still terminated.
    const pid = resolveTransportPid(server.childPid, server.transport);
    const transport = server.transport;
    server.client = null;
    server.transport = undefined;
    server.childPid = null;

    await Promise.allSettled([
      client?.close(),
      terminateStdioTree(pid, transport),
    ]);
    await ensurePidTerminated(pid);
  }
}