pi-mcp-extension 1.11.3 → 1.11.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-mcp-extension",
-  "version": "1.11.3",
+  "version": "1.11.4",
   "description": "MCP (Model Context Protocol) client extension for the Pi coding agent — connect Pi to any MCP server",
   "type": "module",
   "keywords": [

package/src/connection-coordinator.ts

@@ -73,6 +73,24 @@ type InteractiveAuthParams = {
   reset: boolean;
 };
 
+/**
+ * Whether a declined server still deserves the `run /mcp:auth <name>` footer hint.
+ *
+ * WHY: A decline records "this server needed authorization at some point", which stays
+ * true until there is positive evidence otherwise. `ready` is that evidence, and so is a
+ * `stopped` server whose `lastError` is now a non-auth failure (`ECONNREFUSED`, timeout) —
+ * that failure overwrote `AuthRequiredError`, so authorization is no longer the problem
+ * and the hint would send the user down the wrong path. An unknown server drops out too.
+ * Anything ambiguous (no `lastError` yet, still `starting`) keeps the hint, because
+ * clearing it on doubt would hide a pending authorization the user was just told about.
+ */
+function retainsAuthHint(server: Pick<ManagedServer, "state" | "lastError"> | undefined): boolean {
+  if (!server || server.state === "ready") return false;
+  const lastError = server.lastError;
+  if (!lastError || lastError instanceof AuthRequiredError) return true;
+  return server.state !== "stopped";
+}
+
 export class ConnectionCoordinator {
   private readonly deps: ConnectionCoordinatorDeps;
   private readonly connectAttempts = new Map<string, ConnectAttemptEntry>();
@@ -124,7 +142,7 @@ export class ConnectionCoordinator {
   syncStatus(ctx: ExtensionContext | undefined = this.activeCtx): void {
     for (const declinedName of Array.from(this.declinedStartupAuth)) {
       const server = this.deps.manager.getServer(declinedName);
-      if (!server || server.state === "ready") {
+      if (!retainsAuthHint(server)) {
         this.declinedStartupAuth.delete(declinedName);
       }
     }
@@ -147,18 +165,23 @@ export class ConnectionCoordinator {
 
   handleRecoveryFailed(serverName: string, error: unknown): void {
     if (error instanceof AuthRequiredError) {
-      this.declinedStartupAuth.add(serverName);
       this.activeCtx?.ui.notify(
         `pi-mcp: ${serverName} requires re-authorization — run /mcp:auth ${serverName}`,
         "warning",
       );
-      this.syncStatus(this.activeCtx);
-    } else {
-      this.activeCtx?.ui.notify(
-        `pi-mcp: Lost connection to ${serverName} — ${formatErrorMessage(error)}`,
-        "error",
-      );
+      // WHY: Route through markAuthRequired so a mid-session 401 (reported by
+      // ServerManager.handleUnauthorizedServerError) lands in the declined set and
+      // repaints the footer hint exactly like a declined startup prompt.
+      this.markAuthRequired(serverName);
+      return;
     }
+    this.activeCtx?.ui.notify(
+      `pi-mcp: Lost connection to ${serverName} — ${formatErrorMessage(error)}`,
+      "error",
+    );
+    // WHY: The failure just overwrote lastError, so a hint left over from an earlier
+    // decline now points at authorization that is not the problem — clear it.
+    this.syncStatus(this.activeCtx);
   }
 
   cancelServer(serverName: string): void {
@@ -325,6 +348,9 @@ export class ConnectionCoordinator {
         try {
           await this.deps.manager.startServer(name, ctx.cwd);
         } catch (err) {
+          // WHY: A cancelled start means the user stopped the server (/mcp:stop) or the
+          // session is tearing down; reporting it as a start failure would be a bogus error.
+          if (err instanceof AuthCancelledError) return;
           if (err instanceof AuthRequiredError && cfg.auth && ctx.hasUI) {
             authRequiredServers.add(name);
             return;
@@ -335,12 +361,14 @@ export class ConnectionCoordinator {
       }),
     );
 
+    // WHY: Reconcile the footer hint once every eager start has settled — a declined
+    // server may have connected cleanly or failed with a non-auth error, and either
+    // outcome makes "run /mcp:auth <name>" stale.
+    this.syncStatus(ctx);
+
     for (const [name] of serversToStart) {
       if (!authRequiredServers.has(name)) continue;
-      if (this.declinedStartupAuth.has(name)) {
-        this.syncStatus(ctx);
-        continue;
-      }
+      if (this.declinedStartupAuth.has(name)) continue;
       this.enqueueStartupAuthPrompt(name, ctx);
     }
 
@@ -483,8 +511,13 @@ export class ConnectionCoordinator {
       await this._startAndGuardAbort(serverName, ctx, controller.signal);
       return true;
     } catch (err) {
-      if (!(err instanceof AuthRequiredError) || !serverConfig.auth) throw err;
-      return false;
+      if (err instanceof AuthRequiredError && serverConfig.auth) return false;
+      // WHY: A non-auth connect failure (e.g. ECONNREFUSED) overwrites lastError, so a
+      // hint left over from an earlier decline no longer describes this server's problem.
+      if (!(err instanceof AuthCancelledError)) {
+        this.syncStatus(ctx);
+      }
+      throw err;
     }
   }
 
@@ -493,7 +526,17 @@ export class ConnectionCoordinator {
     this._ensureActive(controller.signal, serverName);
     await this.deps.authorize(serverName, ctx, controller.signal, reset);
     this._ensureActive(controller.signal, serverName);
-    await this._startAndGuardAbort(serverName, ctx, controller.signal);
+    try {
+      await this._startAndGuardAbort(serverName, ctx, controller.signal);
+    } catch (err) {
+      // WHY: Authorization can succeed and the reconnect still fail for a non-auth reason
+      // (ECONNREFUSED, timeout). lastError now names that failure, so a hint left over from
+      // the earlier decline would tell the user to authorize a server that just authorized.
+      if (!(err instanceof AuthCancelledError)) {
+        this.syncStatus(ctx);
+      }
+      throw err;
+    }
   }
 
   private async _startAndGuardAbort(

package/src/index.ts

@@ -162,8 +162,20 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
   manager.setRecoveryFailedCallback((serverName, error) => {
     coordinator.handleRecoveryFailed(serverName, error);
   });
+  manager.setServerReadyCallback(() => {
+    // WHY: Eager starts and background _recover() reconnects never pass through
+    // coordinator.connect(), so without this the "run /mcp:auth <name>" footer hint
+    // survives a server that has since connected cleanly.
+    coordinator.syncStatus();
+  });
   bridge.setAuthErrorCallback(async (serverName) => {
-    await manager.handleUnauthorizedServerError(serverName);
+    const authError = await manager.handleUnauthorizedServerError(serverName);
+    if (authError) {
+      // WHY: A mid-session 401/403 stops the server synchronously, which cancels the
+      // deferred _recover() that would have called onRecoveryFailed — so report the
+      // failure here or the re-auth warning and the footer hint never appear.
+      coordinator.handleRecoveryFailed(serverName, authError);
+    }
   });
 
   // ── 3. Session lifecycle ──────────────────────────────────────────────────

package/src/server-manager.ts

@@ -106,6 +106,17 @@ export type ToolRefreshCallback = (
 /** 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;
@@ -206,6 +217,7 @@ export class ServerManager {
   private onToolRefresh: ToolRefreshCallback | 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;
@@ -239,6 +251,11 @@ export class ServerManager {
     this.onRecoveryFailed = cb;
   }
 
+  /** Register callback invoked whenever a server reaches the `ready` state. */
+  setServerReadyCallback(cb: ServerReadyCallback): void {
+    this.onServerReady = cb;
+  }
+
   private _resolveTimeoutMs(server: ManagedServer): number {
     return server.config.requestTimeoutMs ?? this.settings.requestTimeoutMs;
   }
@@ -335,11 +352,29 @@ export class ServerManager {
   /**
     * 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<void> {
+  async handleUnauthorizedServerError(name: string): Promise<AuthRequiredError | undefined> {
     const server = this.servers.get(name);
     if (!server || !server.config.auth || server.config.transport === "stdio") {
-      return;
+      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
@@ -348,6 +383,7 @@ export class ServerManager {
     server.lastError = authError;
     await invalidateStoredTokens(name, this.storageDir, server.config.url);
     server.lastError = authError;
+    return authError;
   }
 
   /** Get recent stderr output for a server. */
@@ -557,30 +593,9 @@ export class ServerManager {
     );
 
     // tools/list_changed: re-discover tools and update Pi registrations
-    client.setNotificationHandler(
-      ToolListChangedNotificationSchema,
-      async () => {
-        if (!signal.aborted && this.onToolRefresh && server.state === "ready" && server.client === client) {
-          try {
-            const refreshTimeoutMs = this._resolveTimeoutMs(server);
-            await this.onToolRefresh(server.name, server.client, {
-              requestTimeoutMs: refreshTimeoutMs,
-              signal,
-            });
-          } catch (err) {
-            const { kind, message } = describeError(err);
-            reportDiagnostic({
-              level: "error",
-              event: "server.tools_refresh_failed",
-              server: server.name,
-              message: `Failed to refresh tools for ${server.name}: ${message}`,
-              fields: { kind },
-              notify: true,
-            });
-          }
-        }
-      },
-    );
+    client.setNotificationHandler(ToolListChangedNotificationSchema, async () => {
+      await this._handleToolListChanged(server, client, signal);
+    });
 
     // notifications/message: structured logging from MCP servers
     client.setNotificationHandler(
@@ -608,6 +623,48 @@ export class ServerManager {
     return client;
   }
 
+  /**
+   * 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._resolveTimeoutMs(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,
+      });
+    }
+  }
+
   private async _handshakeAndDiscover(params: HandshakeParams): Promise<void> {
     const { server, client, transport, timeoutMs, signal, currentCwd } = params;
     await this._connectWithTimeoutAndSignal(client, transport, timeoutMs, signal, server.name);
@@ -626,6 +683,9 @@ export class ServerManager {
     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.