pi-mcp-extension 1.9.3 → 1.9.5

package/package.json

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

package/src/auth-challenge.ts

@@ -17,6 +17,31 @@ export interface ManualAuthChallengeOptions {
   fetchFn?: AuthChallengeFetch | undefined;
   timeoutMs?: number | undefined;
   headers?: Record<string, string> | undefined;
+  signal?: AbortSignal | undefined;
+}
+
+/**
+ * Wrap a fetch implementation so configured server headers are attached as
+ * defaults on every request (e.g. during SDK OAuth metadata discovery and
+ * token exchange) without overriding explicit per-request headers.
+ */
+export function createHeadersFetch(
+  baseHeaders?: Record<string, string>,
+  baseFetch: typeof fetch = fetch,
+): typeof fetch {
+  if (!baseHeaders || Object.keys(baseHeaders).length === 0) {
+    return baseFetch;
+  }
+  return async (input, init) => {
+    const merged = new Headers(baseHeaders);
+    if (init?.headers) {
+      const incoming = new Headers(init.headers);
+      incoming.forEach((value, key) => {
+        merged.set(key, value);
+      });
+    }
+    return baseFetch(input, { ...init, headers: merged });
+  };
 }
 
 /**
@@ -32,14 +57,33 @@ export async function discoverManualAuthChallenge(
   options: ManualAuthChallengeOptions = {},
 ): Promise<ManualAuthChallenge> {
   const fetchFn = options.fetchFn ?? fetch;
+  const timeoutSignal = AbortSignal.timeout(options.timeoutMs ?? AUTH_CHALLENGE_TIMEOUT_MS);
+  const signal = options.signal ? AbortSignal.any([options.signal, timeoutSignal]) : timeoutSignal;
   const headers = new Headers(options.headers);
   headers.set("Accept", "application/json, text/event-stream");
-  const response = await fetchFn(serverUrl, {
+  let response = await fetchFn(serverUrl, {
     method: "GET",
     headers,
-    signal: AbortSignal.timeout(options.timeoutMs ?? AUTH_CHALLENGE_TIMEOUT_MS),
+    signal,
   });
 
+  if (response.status === 405) {
+    // WHY: Streamable HTTP MCP servers may only accept POST and return 405 on GET without
+    // a WWW-Authenticate challenge. Probe with a minimal JSON-RPC ping POST to elicit 401/403.
+    await response.body?.cancel().catch(() => {
+      // WHY: Body cancel errors during challenge probe cleanup are non-fatal.
+    });
+    const postHeaders = new Headers(options.headers);
+    postHeaders.set("Accept", "application/json, text/event-stream");
+    postHeaders.set("Content-Type", "application/json");
+    response = await fetchFn(serverUrl, {
+      method: "POST",
+      headers: postHeaders,
+      body: JSON.stringify({ jsonrpc: "2.0", method: "ping", id: 0 }),
+      signal,
+    });
+  }
+
   try {
     if (response.status !== 401 && response.status !== 403) {
       return {};
@@ -55,6 +99,8 @@ export async function discoverManualAuthChallenge(
       ...(scope ? { scope } : {}),
     };
   } finally {
-    await response.body?.cancel().catch(() => {});
+    await response.body?.cancel().catch(() => {
+      // WHY: Body cancel errors during challenge probe cleanup are non-fatal.
+    });
   }
 }

package/src/auth-lock.ts

@@ -91,7 +91,7 @@ export interface AuthLockTimer {
 export interface AuthLockClock {
   now(): number;
   setInterval(callback: () => void, ms: number): ReturnType<typeof setInterval> | AuthLockTimer;
-  clearInterval(timer: any): void;
+  clearInterval(timer: ReturnType<typeof setInterval> | AuthLockTimer): void;
 }
 
 export interface AuthLockDeps {
@@ -108,6 +108,7 @@ export interface AcquireAuthLockOptions {
 }
 
 export type AuthLockHandle = (() => Promise<void>) & {
+  unlock: () => Promise<void>;
   assertHeld: () => Promise<void>;
   isLost: () => boolean;
 };
@@ -153,7 +154,7 @@ export function createDefaultAuthLockDeps(): AuthLockDeps {
     clock: {
       now: () => Date.now(),
       setInterval: (cb, ms) => setInterval(cb, ms),
-      clearInterval: (timer) => clearInterval(timer),
+      clearInterval: (timer) => clearInterval(timer as ReturnType<typeof setInterval>),
     },
     pid: process.pid,
     hostname: osHostname(),
@@ -204,6 +205,33 @@ function formatHolderDetails(record: LockRecord | null): string | undefined {
 
 const MAX_ACQUIRE_ATTEMPTS = 3;
 
+async function readLockRecordWithRetry(
+  fs: AuthLockFs,
+  path: string,
+  usingRealFs = false,
+): Promise<LockRecord | null> {
+  for (let i = 0; i < 3; i++) {
+    try {
+      const raw = await fs.readFile(path, "utf8");
+      return parseLockRecord(raw);
+    } catch (err) {
+      if ((err as NodeJS.ErrnoException).code === "ENOENT" && i < 2) {
+        // WHY: During stale-lock inspection, a contender briefly renames `path` to
+        // `movedPath` and restores it via `link(movedPath, path)` on nonce mismatch.
+        // Yield across microtasks (and a brief macrotask on real fs) before treating ENOENT as genuine lock loss.
+        await Promise.resolve();
+        await Promise.resolve();
+        if (usingRealFs && i === 1) {
+          await new Promise<void>((resolve) => setTimeout(resolve, 5));
+        }
+        continue;
+      }
+      return null;
+    }
+  }
+  return null;
+}
+
 export async function acquireAuthLock(
   serverName: string,
   storageDir?: string,
@@ -239,8 +267,12 @@ export async function acquireAuthLock(
         await file.writeFile(JSON.stringify(myRecord));
         await file.sync();
       } catch (writeErr) {
-        await file.close().catch(() => {});
-        await deps.fs.unlink(path).catch(() => {});
+        await file.close().catch(() => {
+          // WHY: Best-effort file handle close after write failure.
+        });
+        await deps.fs.unlink(path).catch(() => {
+          // WHY: Best-effort cleanup of partially written lock file before rethrowing writeErr.
+        });
         throw writeErr;
       }
       await file.close();
@@ -257,7 +289,9 @@ export async function acquireAuthLock(
       };
 
       const onExit = () => {
-        if (released || lost || !usingRealFs) return;
+        // WHY: Do not skip cleanup when `lost` is true — if a contender restored our
+        // lock file via `link()` after a transient rename, we must still unlink our own nonce.
+        if (released || !usingRealFs) return;
         try {
           const raw = readFileSync(path, "utf8");
           const current = parseLockRecord(raw);
@@ -265,7 +299,7 @@ export async function acquireAuthLock(
             unlinkSync(path);
           }
         } catch {
-          // Best-effort cleanup on process exit
+          // WHY: Best-effort synchronous cleanup on process exit; lock file may already be unlinked.
         }
       };
 
@@ -273,7 +307,6 @@ export async function acquireAuthLock(
         if (lost || released) return;
         lost = true;
         stopHeartbeat();
-        process.removeListener("exit", onExit);
         options?.onLost?.();
       };
 
@@ -281,13 +314,7 @@ export async function acquireAuthLock(
 
       heartbeatTimer = deps.clock.setInterval(async () => {
         if (released || lost) return;
-        let current: LockRecord | null = null;
-        try {
-          const raw = await deps.fs.readFile(path, "utf8");
-          current = parseLockRecord(raw);
-        } catch {
-          current = null;
-        }
+        const current = await readLockRecordWithRetry(deps.fs, path, usingRealFs);
         if (released || lost) return;
         if (!current || current.nonce !== myRecord.nonce) {
           markLost();
@@ -298,7 +325,15 @@ export async function acquireAuthLock(
           await deps.fs.utimes(path, nowSec, nowSec);
         } catch {
           if (!released && !lost) {
-            markLost();
+            const after = await readLockRecordWithRetry(deps.fs, path, usingRealFs);
+            if (!after || after.nonce !== myRecord.nonce) {
+              markLost();
+            } else {
+              await deps.fs.utimes(path, nowSec, nowSec).catch(() => {
+                // WHY: If utimes still fails after verifying the restored lock nonce, mark lock lost.
+                if (!released && !lost) markLost();
+              });
+            }
           }
         }
       }, HEARTBEAT_MS);
@@ -308,13 +343,7 @@ export async function acquireAuthLock(
         if (lost || released) {
           throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
         }
-        let current: LockRecord | null = null;
-        try {
-          const raw = await deps.fs.readFile(path, "utf8");
-          current = parseLockRecord(raw);
-        } catch {
-          current = null;
-        }
+        const current = await readLockRecordWithRetry(deps.fs, path, usingRealFs);
         if (!current || current.nonce !== myRecord.nonce) {
           markLost();
           throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
@@ -326,20 +355,18 @@ export async function acquireAuthLock(
         released = true;
         stopHeartbeat();
         process.removeListener("exit", onExit);
-        if (lost) return;
-        let current: LockRecord | null = null;
-        try {
-          const raw = await deps.fs.readFile(path, "utf8");
-          current = parseLockRecord(raw);
-        } catch {
-          current = null;
-        }
+        // WHY: Always check and unlink if the file on disk still carries our nonce, even if
+        // `lost` was set during a contender's temporary stale-check rename+link window.
+        const current = await readLockRecordWithRetry(deps.fs, path, usingRealFs);
         if (current && current.nonce === myRecord.nonce) {
-          await deps.fs.unlink(path).catch(() => {});
+          await deps.fs.unlink(path).catch(() => {
+            // WHY: Best-effort lock removal on unlock; ignore ENOENT if concurrently removed.
+          });
         }
       };
 
       return Object.assign(unlock, {
+        unlock,
         assertHeld,
         isLost: () => lost,
       });
@@ -410,12 +437,18 @@ export async function acquireAuthLock(
       }
 
       if (matchesStale) {
-        await deps.fs.unlink(movedPath).catch(() => {});
+        await deps.fs.unlink(movedPath).catch(() => {
+          // WHY: Best-effort removal of reclaimed stale lock file.
+        });
         continue;
       }
 
-      await deps.fs.link(movedPath, path).catch(() => {});
-      await deps.fs.unlink(movedPath).catch(() => {});
+      await deps.fs.link(movedPath, path).catch(() => {
+        // WHY: Best-effort restore of contender's fresh lock file if path was not already recreated.
+      });
+      await deps.fs.unlink(movedPath).catch(() => {
+        // WHY: Best-effort cleanup of temporary stale-check file.
+      });
       throw new AuthLockError(serverName, formatHolderDetails(movedRecord));
     }
   }

package/src/errors.ts

@@ -7,6 +7,8 @@
 export type McpErrorCode =
   | "config"      // Configuration loading or validation failed
   | "connection"  // Transport/connection failed
+  | "auth"        // Authentication/authorization failed
+  | "timeout"     // Request or connection timed out
   | "protocol"    // JSON-RPC protocol violation (server error response, timeout, etc.)
   | "tool";       // Tool execution error (isError: true from server)
 
@@ -33,3 +35,11 @@ export class McpError extends Error {
     return `[${this.server}] ${this.message}`;
   }
 }
+
+export class AuthCancelledError extends Error {
+  constructor(serverName: string) {
+    super(`Authentication cancelled for ${serverName}`);
+    this.name = "AuthCancelledError";
+  }
+}
+

package/src/index.ts

@@ -19,7 +19,7 @@ import { ServerManager } from "./server-manager.js";
 import { AuthLockError, AuthRequiredError } from "./oauth-provider.js";
 import * as Type from "typebox";
 import { ToolBridge } from "./tool-bridge.js";
-import { McpError } from "./errors.js";
+import { McpError, AuthCancelledError } from "./errors.js";
 import { spawn } from "node:child_process";
 
 // OAuth imports
@@ -33,7 +33,7 @@ import {
 import { CallbackUrlValidationError, isSafeAuthorizationUrl, parseCallbackUrl } from "./callback-url.js";
 import { setCallbackPort, McpOAuthProvider, acquireAuthLock, type AuthLockHandle } from "./oauth-provider.js";
 import { auth } from "@modelcontextprotocol/sdk/client/auth.js";
-import { discoverManualAuthChallenge } from "./auth-challenge.js";
+import { discoverManualAuthChallenge, createHeadersFetch } from "./auth-challenge.js";
 import type { ManualAuthChallenge } from "./auth-challenge.js";
 
 export interface BrowserOpenCommand {
@@ -76,12 +76,7 @@ function openBrowser(url: string): Promise<void> {
   });
 }
 
-export class AuthCancelledError extends Error {
-  constructor(serverName: string) {
-    super(`Authentication cancelled for ${serverName}`);
-    this.name = "AuthCancelledError";
-  }
-}
+export { AuthCancelledError } from "./errors.js";
 
 export const authRetryOption = "Retry: open browser again";
 export const authPasteOption = "Paste redirect URL";
@@ -480,16 +475,41 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
     }
   }
 
-  manager.setToolRefreshCallback(async (serverName, client) => {
-    await bridge.refreshTools(serverName, client, () => {
-      const server = manager.getServer(serverName);
-      return server?.client === client && server.state !== "stopped";
-    });
+  let activeCtx: ExtensionContext | undefined;
+
+  manager.setToolRefreshCallback(async (serverName, client, options) => {
+    await bridge.refreshTools(
+      serverName,
+      client,
+      () => {
+        const server = manager.getServer(serverName);
+        return server?.client === client && server.state !== "stopped";
+      },
+      options,
+    );
   });
   manager.setServerStoppedCallback((name) => bridge.deactivateServer(name));
+  manager.setRecoveryFailedCallback((serverName, error) => {
+    if (error instanceof AuthRequiredError) {
+      activeCtx?.ui.notify(
+        `pi-mcp: ${serverName} requires re-authorization — run /mcp:auth ${serverName}`,
+        "warning",
+      );
+      activeCtx?.ui.setStatus?.("pi-mcp", `run /mcp:auth ${serverName}`);
+    } else {
+      activeCtx?.ui.notify(
+        `pi-mcp: Lost connection to ${serverName} — ${formatErrorMessage(error)}`,
+        "error",
+      );
+    }
+  });
+  bridge.setAuthErrorCallback(async (serverName, err) => {
+    await manager.handleUnauthorizedServerError(serverName, err);
+  });
 
   // ── 3. Session lifecycle ──────────────────────────────────────────────────
   pi.on("session_start", async (_event, ctx: ExtensionContext) => {
+    activeCtx = ctx;
     // Only load project-level <cwd>/.pi/mcp.json when the project is trusted.
     const includeProjectConfig =
       typeof (ctx as { isProjectTrusted?: () => boolean }).isProjectTrusted === "function"
@@ -545,6 +565,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
 
   // Connect or disconnect MCP servers registered dynamically via pi.registerMcpServer().
   pi.on("mcp_servers_change", async (event, ctx: ExtensionContext) => {
+    activeCtx = ctx;
     const entries = Array.isArray((event as { servers?: unknown })?.servers)
       ? ((event as { servers: RegisteredMcpServer[] }).servers)
       : getRegisteredEntries();
@@ -594,6 +615,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
   });
 
   pi.on("session_shutdown", async (_event, _ctx: ExtensionContext) => {
+    activeCtx = undefined;
     for (const controller of startupPromptControllers.values()) controller.abort();
     startupPromptControllers.clear();
     for (const attempt of connectAttempts.values()) attempt.controller.abort();
@@ -731,6 +753,8 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
           host: callbackServerConfig.host,
           allowPortFallback: callbackServerConfig.allowPortFallback,
         });
+        if (signal.aborted) throw new AuthCancelledError(serverName);
+        if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
         setCallbackPort(port);
 
         // 2. Generate a cryptographically secure state parameter for CSRF protection
@@ -771,8 +795,28 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
           false,
           () => signal.aborted,
           () => unlock!.assertHeld(),
+          config.url,
         );
 
+        if (!reset) {
+          // WHY: authorize() is only entered after silent connect failed with AuthRequiredError.
+          // If a refresh_token remained on disk (e.g. accepted by the Authorization Server
+          // while the Resource Server rejects the refreshed access token), the SDK's auth()
+          // would silently refresh and return "AUTHORIZED" without opening the browser.
+          // Clear stored tokens while preserving clientInfo and discoveryState.
+          try {
+            await authProvider.invalidateCredentials("tokens");
+          } catch (err) {
+            if (err instanceof McpError && err.code === "config" && /Corrupt OAuth state file/i.test(err.message)) {
+              console.error(`[pi-mcp] ${err.message}`);
+            } else {
+              throw err;
+            }
+          }
+          if (signal.aborted) throw new AuthCancelledError(serverName);
+          if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
+        }
+
         // Why set state before calling auth(): the SDK reads provider.state() while
         // building the authorization URL inside auth(), and omits the state parameter
         // if empty, which the callback server then rejects as a CSRF error.
@@ -781,23 +825,33 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
         // 4. Discover protected-resource challenge metadata from the server.
         const authChallenge: ManualAuthChallenge = await discoverManualAuthChallenge(config.url, {
           headers: config.headers,
+          signal,
         }).catch((): ManualAuthChallenge => {
+          if (signal.aborted) {
+            throw new AuthCancelledError(serverName);
+          }
           console.warn(
             `[pi-mcp] Failed to discover OAuth challenge for "${serverName}"; falling back to standard discovery`,
           );
           return {};
         });
+        if (signal.aborted) throw new AuthCancelledError(serverName);
+        if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
 
         // 5. Register the callback immediately before auth() can open the browser.
         // Attach a catch handler so cancellation after immediate authorization
         // does not produce an unhandled rejection.
         const callbackPromise = waitForCallback(oauthState);
-        callbackPromise.catch(() => {});
+        callbackPromise.catch(() => {
+          // WHY: Suppress unhandled rejection when callbackPromise is cancelled after immediate authorization.
+        });
 
+        const customFetch = config.headers ? createHeadersFetch(config.headers) : undefined;
         const authOptions: {
           serverUrl: string;
           resourceMetadataUrl?: URL;
           scope?: string;
+          fetchFn?: typeof fetch;
         } = { serverUrl: config.url };
         if (authChallenge.resourceMetadataUrl) {
           authOptions.resourceMetadataUrl = authChallenge.resourceMetadataUrl;
@@ -805,13 +859,28 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
         if (authChallenge.scope) {
           authOptions.scope = authChallenge.scope;
         }
+        if (customFetch) {
+          authOptions.fetchFn = customFetch;
+        }
 
         // 6. Start the auth flow.
         // Why branch on the return value: the SDK auth() call returns "AUTHORIZED" when
         // stored or refreshed tokens already succeeded without opening a browser, or
         // "REDIRECT" after invoking redirectToAuthorization(), in which case we must
         // wait for the callback code and call auth() a second time with authorizationCode.
-        const authResult = await auth(authProvider, authOptions);
+        let authResult: "AUTHORIZED" | "REDIRECT";
+        try {
+          authResult = await auth(authProvider, authOptions);
+        } catch (err) {
+          if (err instanceof McpError && err.code === "config" && /Corrupt OAuth state file/i.test(err.message)) {
+            // WHY: McpOAuthProvider._loadState() backed up and removed the corrupt state file.
+            // Retry auth() once so interactive authorization proceeds cleanly.
+            console.error(`[pi-mcp] ${err.message}`);
+            authResult = await auth(authProvider, authOptions);
+          } else {
+            throw err;
+          }
+        }
         if (signal.aborted) throw new AuthCancelledError(serverName);
         if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
 
@@ -852,6 +921,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
             authorizationCode: string;
             resourceMetadataUrl?: URL;
             scope?: string;
+            fetchFn?: typeof fetch;
           } = {
             serverUrl: config.url,
             authorizationCode: code,
@@ -862,8 +932,13 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
           if (authChallenge.scope) {
             finishAuthOptions.scope = authChallenge.scope;
           }
+          if (customFetch) {
+            finishAuthOptions.fetchFn = customFetch;
+          }
 
           const finishAuthResult = await auth(authProvider, finishAuthOptions);
+          if (signal.aborted) throw new AuthCancelledError(serverName);
+          if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
 
           if (finishAuthResult !== "AUTHORIZED") {
             throw new McpError(
@@ -904,6 +979,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
     reset = false,
     alreadyQueued = false,
   ): Promise<void> {
+    activeCtx = ctx;
     const server = manager.getServer(serverName);
     if (!server) throw new McpError(`Unknown server "${serverName}"`, serverName, "config");
     if (signal?.aborted) throw new AuthCancelledError(serverName);
@@ -1027,11 +1103,28 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
     parameters: Type.Object({}),
     async execute() {
       const rows = await Promise.all(manager.getAllServers().map(async (server) => {
-        const auth = server.config.auth ? await manager.getServerAuthStatus(server.name) : null;
+        let auth: Awaited<ReturnType<typeof manager.getServerAuthStatus>> = null;
+        let authCorrupt = false;
+        if (server.config.auth) {
+          try {
+            auth = await manager.getServerAuthStatus(server.name);
+          } catch (err) {
+            // WHY: Never let a single server's corrupt state file abort mcp_status for all servers.
+            // getAuthStatus() already backed up and deleted the corrupt file before throwing.
+            console.error(`[pi-mcp] ${err instanceof Error ? err.message : String(err)}`);
+            authCorrupt = true;
+          }
+        }
         return {
           name: server.name, lifecycle: server.config.lifecycle, state: server.state,
           tools: bridge.getToolActivity(server.name),
-          auth: server.config.auth ? auth?.hasTokens ? "credentials stored" : "authorization may be required" : "not configured",
+          auth: server.config.auth
+            ? authCorrupt
+              ? "corrupt state file backed up (re-auth required)"
+              : auth?.hasTokens
+                ? "credentials stored"
+                : "authorization may be required"
+            : "not configured",
           error: server.lastError
             ? server.lastError instanceof AuthRequiredError
               ? "Authorization required: use mcp_connect in an interactive Pi session"
@@ -1082,6 +1175,12 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
       const { name, reset } = parseAuthArguments(args, manager.getAllServers().map((server) => server.name));
       if (!name) { ctx.ui.notify("Usage: /mcp:auth <name> [--reset]", "error"); return; }
       if (!manager.getServer(name)?.config.auth) { ctx.ui.notify(`No OAuth server named ${name}`, "error"); return; }
+      if (!reset && manager.getServer(name)?.state === "ready") {
+        // WHY: Explicit /mcp:auth on an already-ready server must tear down the existing
+        // connection and force a fresh reconnect/auth check rather than returning a no-op.
+        bridge.deactivateServer(name);
+        await manager.stopServer(name);
+      }
       try { await connectRequested(name, ctx, undefined, reset); ctx.ui.notify(`pi-mcp: ${name} ready`, "info"); }
       catch (err) {
         ctx.ui.notify(`pi-mcp: ${err instanceof Error ? err.message : String(err)}`, err instanceof AuthCancelledError ? "info" : "error");