pi-mcp-extension 1.7.2 → 1.7.3

package/package.json

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

package/src/auth-lock.ts

@@ -0,0 +1,426 @@
+/**
+ * Cross-process per-server interactive OAuth authorization lock.
+ *
+ * Uses an atomic lock file (<hash>.json.lock, mode 0600) containing a
+ * versioned JSON record (PID, hostname, random nonce, creation timestamp)
+ * plus a periodic mtime heartbeat. Automatically reclaims locks left behind
+ * by crashed or SIGKILLed processes using a nonce-verified rename protocol.
+ */
+
+import {
+  open,
+  readFile,
+  stat,
+  rename,
+  unlink,
+  link,
+  utimes,
+  mkdir,
+  chmod,
+} from "node:fs/promises";
+import { readFileSync, unlinkSync } from "node:fs";
+import { hostname as osHostname } from "node:os";
+import { join } from "node:path";
+import { createHash, randomUUID } from "node:crypto";
+import { z } from "zod";
+import { agentDir } from "./config.js";
+
+// ─── Constants & Schema ───────────────────────────────────────────────────────
+
+export const HEARTBEAT_MS = 30_000;
+export const STALE_MS = 90_000; // 3 × 30s heartbeat
+
+export const LockRecordSchema = z.object({
+  v: z.literal(1),
+  pid: z.number().int().positive(),
+  hostname: z.string().min(1),
+  nonce: z.string().min(1),
+  createdAt: z.number(),
+});
+
+export type LockRecord = z.output<typeof LockRecordSchema>;
+
+// ─── Errors ───────────────────────────────────────────────────────────────────
+
+export class AuthRequiredError extends Error {
+  constructor(
+    serverName: string,
+    message = `Authorization required for ${serverName}. Connect in an interactive Pi session or use /mcp:auth ${serverName}.`,
+  ) {
+    super(message);
+    this.name = "AuthRequiredError";
+  }
+}
+
+export class AuthLockError extends AuthRequiredError {
+  constructor(serverName: string, details?: string) {
+    const message = details
+      ? details.startsWith("Authorization lock") || details.startsWith("Another Pi process")
+        ? details
+        : `Another Pi process (${details}) is authorizing ${serverName}. Retry after it finishes.`
+      : `Another Pi process is authorizing ${serverName}. Retry after it finishes.`;
+    super(serverName, message);
+    this.name = "AuthLockError";
+  }
+}
+
+// ─── Dependency Injection Types ───────────────────────────────────────────────
+
+export interface AuthLockFileHandle {
+  writeFile(data: string): Promise<void>;
+  sync(): Promise<void>;
+  close(): Promise<void>;
+}
+
+export interface AuthLockFs {
+  open(path: string, flags: string, mode?: number): Promise<AuthLockFileHandle>;
+  readFile(path: string, encoding: BufferEncoding): Promise<string>;
+  stat(path: string): Promise<{ mtimeMs: number }>;
+  rename(oldPath: string, newPath: string): Promise<void>;
+  unlink(path: string): Promise<void>;
+  link(existingPath: string, newPath: string): Promise<void>;
+  utimes(path: string, atime: number | string | Date, mtime: number | string | Date): Promise<void>;
+  mkdir(path: string, options?: { recursive?: boolean; mode?: number }): Promise<string | undefined | void>;
+  chmod(path: string, mode: number): Promise<void>;
+}
+
+export interface AuthLockTimer {
+  unref?: () => void;
+}
+
+export interface AuthLockClock {
+  now(): number;
+  setInterval(callback: () => void, ms: number): ReturnType<typeof setInterval> | AuthLockTimer;
+  clearInterval(timer: any): void;
+}
+
+export interface AuthLockDeps {
+  fs: AuthLockFs;
+  clock: AuthLockClock;
+  pid: number;
+  hostname: string;
+  isPidAlive: (pid: number) => boolean;
+}
+
+export interface AcquireAuthLockOptions {
+  onLost?: (() => void) | undefined;
+  deps?: Partial<AuthLockDeps> | undefined;
+}
+
+export type AuthLockHandle = (() => Promise<void>) & {
+  assertHeld: () => Promise<void>;
+  isLost: () => boolean;
+};
+
+// ─── Path & Default Dependency Helpers ────────────────────────────────────────
+
+export function authDir(storageDir?: string): string {
+  return storageDir ?? join(agentDir(), "mcp-auth");
+}
+
+export function statePath(serverName: string, storageDir?: string): string {
+  // Hash the server name to avoid filesystem issues with special chars
+  const hash = createHash("sha256").update(serverName).digest("hex").slice(0, 16);
+  return join(authDir(storageDir), `${hash}.json`);
+}
+
+export function lockPath(serverName: string, storageDir?: string): string {
+  return `${statePath(serverName, storageDir)}.lock`;
+}
+
+export function defaultIsPidAlive(pid: number): boolean {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch (err) {
+    return (err as NodeJS.ErrnoException).code === "EPERM";
+  }
+}
+
+export function createDefaultAuthLockDeps(): AuthLockDeps {
+  return {
+    fs: {
+      open: (path, flags, mode) => open(path, flags, mode),
+      readFile: (path, encoding) => readFile(path, encoding),
+      stat: (path) => stat(path),
+      rename: (oldPath, newPath) => rename(oldPath, newPath),
+      unlink: (path) => unlink(path),
+      link: (existingPath, newPath) => link(existingPath, newPath),
+      utimes: (path, atime, mtime) => utimes(path, atime, mtime),
+      mkdir: (path, options) => mkdir(path, options),
+      chmod: (path, mode) => chmod(path, mode),
+    },
+    clock: {
+      now: () => Date.now(),
+      setInterval: (cb, ms) => setInterval(cb, ms),
+      clearInterval: (timer) => clearInterval(timer),
+    },
+    pid: process.pid,
+    hostname: osHostname(),
+    isPidAlive: defaultIsPidAlive,
+  };
+}
+
+// ─── Pure Staleness Policy ────────────────────────────────────────────────────
+
+export function parseLockRecord(raw: string): LockRecord | null {
+  if (!raw) return null;
+  try {
+    const parsed: unknown = JSON.parse(raw);
+    const result = LockRecordSchema.safeParse(parsed);
+    return result.success ? result.data : null;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Pure policy determining whether an existing lock file is stale and safe to reclaim.
+ */
+export function isLockStale(
+  record: LockRecord | null,
+  mtimeMs: number,
+  now: number,
+  self: { hostname: string; isPidAlive: (pid: number) => boolean },
+): boolean {
+  if (now - mtimeMs > STALE_MS) {
+    return true;
+  }
+  if (record === null) {
+    return false;
+  }
+  if (record.hostname === self.hostname && !self.isPidAlive(record.pid)) {
+    return true;
+  }
+  return false;
+}
+
+function formatHolderDetails(record: LockRecord | null): string | undefined {
+  if (!record) return undefined;
+  return `pid ${record.pid} on ${record.hostname}`;
+}
+
+// ─── Lock Acquisition, Heartbeat & Release ────────────────────────────────────
+
+const MAX_ACQUIRE_ATTEMPTS = 3;
+
+export async function acquireAuthLock(
+  serverName: string,
+  storageDir?: string,
+  options?: AcquireAuthLockOptions,
+): Promise<AuthLockHandle> {
+  const defaults = createDefaultAuthLockDeps();
+  const deps: AuthLockDeps = {
+    fs: options?.deps?.fs ?? defaults.fs,
+    clock: options?.deps?.clock ?? defaults.clock,
+    pid: options?.deps?.pid ?? defaults.pid,
+    hostname: options?.deps?.hostname ?? defaults.hostname,
+    isPidAlive: options?.deps?.isPidAlive ?? defaults.isPidAlive,
+  };
+  const usingRealFs = options?.deps?.fs === undefined;
+
+  const directory = authDir(storageDir);
+  await deps.fs.mkdir(directory, { recursive: true, mode: 0o700 });
+  await deps.fs.chmod(directory, 0o700);
+  const path = lockPath(serverName, storageDir);
+
+  for (let attempt = 0; attempt < MAX_ACQUIRE_ATTEMPTS; attempt++) {
+    const myRecord: LockRecord = {
+      v: 1,
+      pid: deps.pid,
+      hostname: deps.hostname,
+      nonce: randomUUID(),
+      createdAt: deps.clock.now(),
+    };
+
+    try {
+      const file = await deps.fs.open(path, "wx", 0o600);
+      try {
+        await file.writeFile(JSON.stringify(myRecord));
+        await file.sync();
+      } catch (writeErr) {
+        await file.close().catch(() => {});
+        await deps.fs.unlink(path).catch(() => {});
+        throw writeErr;
+      }
+      await file.close();
+
+      let released = false;
+      let lost = false;
+      let heartbeatTimer: ReturnType<typeof setInterval> | AuthLockTimer | null = null;
+
+      const stopHeartbeat = () => {
+        if (heartbeatTimer !== null) {
+          deps.clock.clearInterval(heartbeatTimer);
+          heartbeatTimer = null;
+        }
+      };
+
+      const onExit = () => {
+        if (released || lost || !usingRealFs) return;
+        try {
+          const raw = readFileSync(path, "utf8");
+          const current = parseLockRecord(raw);
+          if (current && current.nonce === myRecord.nonce) {
+            unlinkSync(path);
+          }
+        } catch {
+          // Best-effort cleanup on process exit
+        }
+      };
+
+      const markLost = () => {
+        if (lost || released) return;
+        lost = true;
+        stopHeartbeat();
+        process.removeListener("exit", onExit);
+        options?.onLost?.();
+      };
+
+      process.once("exit", onExit);
+
+      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;
+        }
+        if (released || lost) return;
+        if (!current || current.nonce !== myRecord.nonce) {
+          markLost();
+          return;
+        }
+        const nowSec = deps.clock.now() / 1000;
+        try {
+          await deps.fs.utimes(path, nowSec, nowSec);
+        } catch {
+          if (!released && !lost) {
+            markLost();
+          }
+        }
+      }, HEARTBEAT_MS);
+      heartbeatTimer.unref?.();
+
+      const assertHeld = async (): Promise<void> => {
+        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;
+        }
+        if (!current || current.nonce !== myRecord.nonce) {
+          markLost();
+          throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
+        }
+      };
+
+      const unlock = async (): Promise<void> => {
+        if (released) return;
+        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;
+        }
+        if (current && current.nonce === myRecord.nonce) {
+          await deps.fs.unlink(path).catch(() => {});
+        }
+      };
+
+      return Object.assign(unlock, {
+        assertHeld,
+        isLost: () => lost,
+      });
+    } catch (err) {
+      if ((err as NodeJS.ErrnoException).code !== "EEXIST") {
+        throw err;
+      }
+
+      // Lock file already exists — inspect and check staleness
+      let statResult: { mtimeMs: number };
+      try {
+        statResult = await deps.fs.stat(path);
+      } catch (statErr) {
+        if ((statErr as NodeJS.ErrnoException).code === "ENOENT") {
+          continue;
+        }
+        throw statErr;
+      }
+
+      let record: LockRecord | null = null;
+      try {
+        const raw = await deps.fs.readFile(path, "utf8");
+        record = parseLockRecord(raw);
+      } catch (readErr) {
+        if ((readErr as NodeJS.ErrnoException).code === "ENOENT") {
+          continue;
+        }
+        record = null;
+      }
+
+      const stale = isLockStale(record, statResult.mtimeMs, deps.clock.now(), {
+        hostname: deps.hostname,
+        isPidAlive: deps.isPidAlive,
+      });
+      if (!stale) {
+        throw new AuthLockError(serverName, formatHolderDetails(record));
+      }
+
+      // Reclaim protocol: move stale candidate aside, verify identity, then delete or restore
+      const movedPath = `${path}.stale-${randomUUID()}`;
+      try {
+        await deps.fs.rename(path, movedPath);
+      } catch (renameErr) {
+        if ((renameErr as NodeJS.ErrnoException).code === "ENOENT") {
+          continue;
+        }
+        throw renameErr;
+      }
+
+      let movedRecord: LockRecord | null = null;
+      try {
+        const movedRaw = await deps.fs.readFile(movedPath, "utf8");
+        movedRecord = parseLockRecord(movedRaw);
+      } catch {
+        movedRecord = null;
+      }
+
+      let matchesStale = false;
+      if (record !== null) {
+        matchesStale = movedRecord !== null && movedRecord.nonce === record.nonce;
+      } else {
+        try {
+          const movedStat = await deps.fs.stat(movedPath);
+          matchesStale = movedRecord === null && movedStat.mtimeMs === statResult.mtimeMs;
+        } catch {
+          matchesStale = false;
+        }
+      }
+
+      if (matchesStale) {
+        await deps.fs.unlink(movedPath).catch(() => {});
+        continue;
+      }
+
+      await deps.fs.link(movedPath, path).catch(() => {});
+      await deps.fs.unlink(movedPath).catch(() => {});
+      throw new AuthLockError(serverName, formatHolderDetails(movedRecord));
+    }
+  }
+
+  throw new AuthLockError(serverName);
+}
+
+export { acquireAuthLock as acquireLock };

package/src/index.ts

@@ -24,7 +24,7 @@ import {
   stopCallbackServer,
   callbackServerConfigFromRedirectUrl,
 } from "./callback-server.js";
-import { setCallbackPort, McpOAuthProvider, acquireAuthLock } from "./oauth-provider.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 type { ManualAuthChallenge } from "./auth-challenge.js";
@@ -367,14 +367,23 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
     const config = server.config;
     let oauthState: string | undefined;
     let latestAuthorizationUrl: URL | undefined;
+    let lockLost = false;
     const cancel = () => { if (oauthState) cancelCallback(oauthState); };
     signal.addEventListener("abort", cancel, { once: true });
-    let unlock: (() => Promise<void>) | undefined;
+    let unlock: AuthLockHandle | undefined;
     try {
         if (signal.aborted) throw new AuthCancelledError(serverName);
-        unlock = await acquireAuthLock(serverName, paths.authStorageDir);
+        unlock = await acquireAuthLock(serverName, paths.authStorageDir, {
+          onLost: () => {
+            lockLost = true;
+            cancel();
+          },
+        });
         if (signal.aborted) throw new AuthCancelledError(serverName);
-        if (reset) await manager.resetServerAuth(serverName);
+        if (reset) {
+          await unlock.assertHeld();
+          await manager.resetServerAuth(serverName);
+        }
         // Validate that we have a server URL (required for OAuth)
         if (!config.url) {
           throw new McpError(
@@ -411,11 +420,12 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
           config.auth || { type: "oauth" },
           (url: URL) => {
             latestAuthorizationUrl = new URL(url.toString());
-            if (!signal.aborted) (paths.browserOpen ?? openBrowser)(url.toString());
+            if (!signal.aborted && !lockLost) (paths.browserOpen ?? openBrowser)(url.toString());
           },
           paths.authStorageDir,
           false,
           () => signal.aborted,
+          () => unlock!.assertHeld(),
         );
 
         // Why set state before calling auth(): the SDK reads provider.state() while
@@ -458,6 +468,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
         // wait for the callback code and call auth() a second time with authorizationCode.
         const authResult = await auth(authProvider, authOptions);
         if (signal.aborted) throw new AuthCancelledError(serverName);
+        if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
 
         if (authResult === "AUTHORIZED") {
           // Auth succeeded without needing browser interaction (e.g., had valid tokens).
@@ -481,11 +492,12 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
                   "protocol",
                 );
               }
-              if (!signal.aborted) (paths.browserOpen ?? openBrowser)(latestAuthorizationUrl.toString());
+              if (!signal.aborted && !lockLost) (paths.browserOpen ?? openBrowser)(latestAuthorizationUrl.toString());
             },
             ctx,
           );
           if (signal.aborted) throw new AuthCancelledError(serverName);
+          if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
           ctx.ui.notify(
             `pi-mcp: Authorization callback received for ${serverName}. Exchanging token...`,
             "info",
@@ -527,6 +539,12 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
         }
 
       if (signal.aborted) throw new AuthCancelledError(serverName);
+      if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
+    } catch (err) {
+      if (lockLost && !signal.aborted) {
+        throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
+      }
+      throw err;
     } finally {
       signal.removeEventListener("abort", cancel);
       cancel();
@@ -674,7 +692,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
           : !manager.getServer(name)
             ? "Unknown MCP server name. Use mcp_status to list configured servers."
             : err instanceof AuthLockError
-              ? "Another Pi process is authorizing this server. Retry after it finishes or check for a stale auth lock."
+              ? "Another Pi process is authorizing this server. Retry after it finishes."
               : err instanceof AuthRequiredError
                 ? ctx.hasUI
                   ? "MCP authorization required. Retry with mcp_connect or /mcp:auth."

package/src/oauth-provider.ts

@@ -13,10 +13,21 @@
 
 import type { OAuthClientProvider, OAuthDiscoveryState } from "@modelcontextprotocol/sdk/client/auth.js";
 import type { OAuthClientMetadata, OAuthClientInformationMixed, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
-import { readFile, writeFile, mkdir, unlink, chmod, open, stat } from "node:fs/promises";
-import { join } from "node:path";
-import { createHash } from "node:crypto";
-import { agentDir } from "./config.js";
+import { readFile, writeFile, mkdir, unlink, chmod, rename } from "node:fs/promises";
+import { randomUUID } from "node:crypto";
+import { z } from "zod";
+import { McpError } from "./errors.js";
+import {
+  acquireAuthLock,
+  AuthLockError,
+  AuthRequiredError,
+  authDir,
+  statePath,
+  type AcquireAuthLockOptions,
+  type AuthLockHandle,
+} from "./auth-lock.js";
+
+export { acquireAuthLock, AuthLockError, AuthRequiredError, type AcquireAuthLockOptions, type AuthLockHandle };
 
 // ─── Types ────────────────────────────────────────────────────────────────────
 
@@ -57,7 +68,7 @@ export function getCallbackPort(): number {
   return callbackPort;
 }
 
-// ─── Persistent State Types ───────────────────────────────────────────────────
+// ─── Persistent State Types & Schema ──────────────────────────────────────────
 
 interface StoredClientInfo {
   client_id: string;
@@ -77,74 +88,113 @@ interface StoredTokens {
 interface StoredState {
   clientInfo: StoredClientInfo | undefined;
   tokens: StoredTokens | undefined;
-  codeVerifier: string | undefined;
   discoveryState: OAuthDiscoveryState | undefined;
 }
 
-// ─── File helpers ─────────────────────────────────────────────────────────────
-
-function authDir(storageDir?: string): string {
-  return storageDir ?? join(agentDir(), "mcp-auth");
-}
+const StoredClientInfoSchema = z.object({
+  client_id: z.string(),
+  client_secret: z.string().optional(),
+});
+
+const StoredTokensSchema = z.object({
+  access_token: z.string(),
+  token_type: z.string().optional(),
+  refresh_token: z.string().optional(),
+  expires_in: z.number().optional(),
+  scope: z.string().optional(),
+  saved_at: z.string().optional(),
+});
+
+export const StoredStateSchema = z.object({
+  clientInfo: StoredClientInfoSchema.optional(),
+  tokens: StoredTokensSchema.optional(),
+  discoveryState: z.record(z.unknown()).optional(),
+});
 
-function statePath(serverName: string, storageDir?: string): string {
-  // Hash the server name to avoid filesystem issues with special chars
-  const hash = createHash("sha256").update(serverName).digest("hex").slice(0, 16);
-  return join(authDir(storageDir), `${hash}.json`);
-}
+// ─── File helpers ─────────────────────────────────────────────────────────────
 
-// Exclusive per-server authorization lock shared by independent Pi processes.
-export async function acquireAuthLock(serverName: string, storageDir?: string): Promise<() => Promise<void>> {
-  const directory = authDir(storageDir);
-  await mkdir(directory, { recursive: true, mode: 0o700 });
-  await chmod(directory, 0o700);
-  const path = `${statePath(serverName, storageDir)}.lock`;
+async function loadState(serverName: string, storageDir?: string): Promise<StoredState> {
+  const path = statePath(serverName, storageDir);
+  let raw: string;
   try {
-    const file = await open(path, "wx", 0o600);
-    try { await file.writeFile(String(process.pid)); }
-    catch (err) { await unlink(path).catch(() => {}); throw err; }
-    finally { await file.close(); }
-    return async () => { await unlink(path).catch(() => {}); };
+    raw = await readFile(path, "utf8");
   } catch (err) {
-    if ((err as NodeJS.ErrnoException).code !== "EEXIST") throw err;
-    throw new AuthLockError(serverName);
+    if ((err as NodeJS.ErrnoException).code === "ENOENT") {
+      return {
+        clientInfo: undefined,
+        tokens: undefined,
+        discoveryState: undefined,
+      };
+    }
+    throw err;
   }
-}
 
-async function loadState(serverName: string, storageDir?: string): Promise<StoredState> {
+  let parsed: z.output<typeof StoredStateSchema>;
   try {
-    const raw = await readFile(statePath(serverName, storageDir), "utf8");
-    const parsed = JSON.parse(raw) as StoredState;
-    return {
-      clientInfo: parsed.clientInfo ?? undefined,
-      tokens: parsed.tokens ?? undefined,
-      codeVerifier: parsed.codeVerifier ?? undefined,
-      discoveryState: parsed.discoveryState ?? undefined,
-    };
-  } catch {
-    return {
-      clientInfo: undefined,
-      tokens: undefined,
-      codeVerifier: undefined,
-      discoveryState: undefined,
-    };
+    const json: unknown = JSON.parse(raw);
+    parsed = StoredStateSchema.parse(json);
+  } catch (err) {
+    const backupPath = `${path}.corrupt-${Date.now()}`;
+    await rename(path, backupPath);
+    throw new McpError(
+      `Corrupt OAuth state file for "${serverName}" backed up to ${backupPath}`,
+      serverName,
+      "config",
+      err,
+    );
   }
+
+  return {
+    clientInfo: parsed.clientInfo
+      ? {
+          client_id: parsed.clientInfo.client_id,
+          client_secret: parsed.clientInfo.client_secret,
+        }
+      : undefined,
+    tokens: parsed.tokens
+      ? {
+          access_token: parsed.tokens.access_token,
+          token_type: parsed.tokens.token_type,
+          refresh_token: parsed.tokens.refresh_token,
+          expires_in: parsed.tokens.expires_in,
+          scope: parsed.tokens.scope,
+          saved_at: parsed.tokens.saved_at,
+        }
+      : undefined,
+    discoveryState: parsed.discoveryState as OAuthDiscoveryState | undefined,
+  };
 }
 
-async function saveState(serverName: string, state: StoredState, storageDir?: string, isCancelled: () => boolean = () => false): Promise<void> {
+async function saveState(
+  serverName: string,
+  state: StoredState,
+  storageDir?: string,
+  isCancelled: () => boolean = () => false,
+): Promise<void> {
   const directory = authDir(storageDir);
   await mkdir(directory, { recursive: true, mode: 0o700 });
   await chmod(directory, 0o700);
-  // Only write defined fields
+  // Only write defined fields (never write codeVerifier to shared state)
   const toWrite: Record<string, unknown> = {};
   if (state.clientInfo !== undefined) toWrite.clientInfo = state.clientInfo;
   if (state.tokens !== undefined) toWrite.tokens = state.tokens;
-  if (state.codeVerifier !== undefined) toWrite.codeVerifier = state.codeVerifier;
   if (state.discoveryState !== undefined) toWrite.discoveryState = state.discoveryState;
   if (isCancelled()) throw new AuthRequiredError(serverName, "Authorization cancelled");
   const path = statePath(serverName, storageDir);
-  await writeFile(path, JSON.stringify(toWrite, null, 2), { encoding: "utf8", mode: 0o600 });
-  await chmod(path, 0o600);
+  const tmpPath = `${path}.tmp-${randomUUID()}`;
+  try {
+    await writeFile(tmpPath, JSON.stringify(toWrite, null, 2), { encoding: "utf8", mode: 0o600 });
+    await chmod(tmpPath, 0o600);
+    if (isCancelled()) {
+      await unlink(tmpPath).catch(() => {});
+      throw new AuthRequiredError(serverName, "Authorization cancelled");
+    }
+    await rename(tmpPath, path);
+    await chmod(path, 0o600);
+  } catch (err) {
+    await unlink(tmpPath).catch(() => {});
+    throw err;
+  }
 }
 
 // ─── OAuthClientProvider Implementation ────────────────────────────────────────
@@ -155,6 +205,8 @@ export class McpOAuthProvider implements OAuthClientProvider {
   private _redirectUrl: string | undefined;
   private _onAuthRequired: ((url: URL) => void | Promise<void>) | undefined;
   private _oauthState: string | undefined;
+  private _codeVerifier: string | undefined;
+  private _cachedTokens: StoredTokens | undefined;
 
   constructor(
     serverName: string,
@@ -163,6 +215,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
     private readonly storageDir?: string,
     private readonly silent = false,
     private readonly isCancelled: () => boolean = () => false,
+    private readonly assertLockHeld?: (() => Promise<void>) | undefined,
   ) {
     this.serverName = serverName;
     this.authConfig = authConfig;
@@ -217,19 +270,21 @@ export class McpOAuthProvider implements OAuthClientProvider {
 
   async saveClientInformation(clientInformation: OAuthClientInformationMixed): Promise<void> {
     if (this.silent) throw new AuthRequiredError(this.serverName);
+    await this.assertLockHeld?.();
     const state = await loadState(this.serverName, this.storageDir);
     state.clientInfo = {
       client_id: clientInformation.client_id,
       client_secret: clientInformation.client_secret,
     };
+    await this.assertLockHeld?.();
     await saveState(this.serverName, state, this.storageDir, this.isCancelled);
   }
 
   // --- tokens ---
 
   async tokens(): Promise<OAuthTokens | undefined> {
-    const state = await loadState(this.serverName, this.storageDir);
-    if (!state.tokens) return undefined;
+    const storedTokens = this._cachedTokens ?? (await loadState(this.serverName, this.storageDir)).tokens;
+    if (!storedTokens) return undefined;
 
     // Always return stored tokens — even if expired.
     // The SDK's auth() function checks tokens?.refresh_token and attempts
@@ -247,19 +302,17 @@ export class McpOAuthProvider implements OAuthClientProvider {
 
     // Build OAuthTokens — only include defined fields
     const tokens: Record<string, string> = {
-      access_token: state.tokens.access_token,
+      access_token: storedTokens.access_token,
     };
-    if (state.tokens.token_type !== undefined) tokens.token_type = state.tokens.token_type;
-    if (state.tokens.refresh_token !== undefined) tokens.refresh_token = state.tokens.refresh_token;
+    if (storedTokens.token_type !== undefined) tokens.token_type = storedTokens.token_type;
+    if (storedTokens.refresh_token !== undefined) tokens.refresh_token = storedTokens.refresh_token;
     return tokens as unknown as OAuthTokens;
   }
 
   // --- saveTokens ---
 
   async saveTokens(tokens: OAuthTokens): Promise<void> {
-    if (this.silent) await this.checkInteractiveLock();
-    const state = await loadState(this.serverName, this.storageDir);
-    state.tokens = {
+    const storedTokens: StoredTokens = {
       access_token: tokens.access_token,
       token_type: tokens.token_type,
       refresh_token: tokens.refresh_token,
@@ -267,8 +320,28 @@ export class McpOAuthProvider implements OAuthClientProvider {
       scope: tokens.scope,
       saved_at: new Date().toISOString(),
     };
+
+    if (this.silent) {
+      try {
+        const state = await loadState(this.serverName, this.storageDir);
+        state.tokens = storedTokens;
+        await saveState(this.serverName, state, this.storageDir, this.isCancelled);
+        this._cachedTokens = undefined;
+      } catch (err) {
+        this._cachedTokens = storedTokens;
+        console.error(`[pi-mcp] Failed to persist refreshed OAuth tokens for "${this.serverName}":`, err);
+      }
+      return;
+    }
+
+    if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
+    await this.assertLockHeld?.();
+    const state = await loadState(this.serverName, this.storageDir);
+    state.tokens = storedTokens;
     if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
+    await this.assertLockHeld?.();
     await saveState(this.serverName, state, this.storageDir, this.isCancelled);
+    this._cachedTokens = undefined;
   }
 
   // --- redirectToAuthorization ---
@@ -285,30 +358,30 @@ export class McpOAuthProvider implements OAuthClientProvider {
     }
   }
 
-  // --- PKCE code verifier ---
+  // --- PKCE code verifier (in-memory per provider instance) ---
 
   async saveCodeVerifier(codeVerifier: string): Promise<void> {
-    // The SDK saves PKCE before redirecting. Silent attempts must not overwrite an active flow.
+    // The SDK saves PKCE before redirecting. Silent attempts must not start a browser flow.
     if (this.silent) throw new AuthRequiredError(this.serverName);
-    const state = await loadState(this.serverName, this.storageDir);
-    state.codeVerifier = codeVerifier;
-    await saveState(this.serverName, state, this.storageDir, this.isCancelled);
+    if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
+    this._codeVerifier = codeVerifier;
   }
 
   async codeVerifier(): Promise<string> {
-    const state = await loadState(this.serverName, this.storageDir);
-    if (!state.codeVerifier) {
+    if (!this._codeVerifier) {
       throw new Error(`[pi-mcp] No PKCE code verifier found for "${this.serverName}"`);
     }
-    return state.codeVerifier;
+    return this._codeVerifier;
   }
 
   // --- Discovery state caching ---
 
   async saveDiscoveryState(discState: OAuthDiscoveryState): Promise<void> {
     if (this.silent) return; // No shared state writes from background discovery.
+    await this.assertLockHeld?.();
     const state = await loadState(this.serverName, this.storageDir);
     state.discoveryState = discState;
+    await this.assertLockHeld?.();
     await saveState(this.serverName, state, this.storageDir, this.isCancelled);
   }
 
@@ -338,35 +411,33 @@ export class McpOAuthProvider implements OAuthClientProvider {
 
   // --- Credential invalidation ---
 
-  private async checkInteractiveLock(): Promise<void> {
-    try { await stat(`${statePath(this.serverName, this.storageDir)}.lock`); }
-    catch (err) { if ((err as NodeJS.ErrnoException).code === "ENOENT") return; throw err; }
-    throw new AuthRequiredError(this.serverName);
-  }
-
   async invalidateCredentials(scope: "all" | "client" | "tokens" | "verifier" | "discovery"): Promise<void> {
     if (this.silent) throw new AuthRequiredError(this.serverName);
+    await this.assertLockHeld?.();
     const state = await loadState(this.serverName, this.storageDir);
     switch (scope) {
       case "all":
         state.clientInfo = undefined;
         state.tokens = undefined;
-        state.codeVerifier = undefined;
         state.discoveryState = undefined;
+        this._codeVerifier = undefined;
+        this._cachedTokens = undefined;
         break;
       case "client":
         state.clientInfo = undefined;
         break;
       case "tokens":
         state.tokens = undefined;
+        this._cachedTokens = undefined;
         break;
       case "verifier":
-        state.codeVerifier = undefined;
+        this._codeVerifier = undefined;
         break;
       case "discovery":
         state.discoveryState = undefined;
         break;
     }
+    await this.assertLockHeld?.();
     await saveState(this.serverName, state, this.storageDir, this.isCancelled);
   }
 }
@@ -387,7 +458,6 @@ export async function getAuthStatus(serverName: string, storageDir?: string): Pr
   if (
     state.clientInfo === undefined &&
     state.tokens === undefined &&
-    state.codeVerifier === undefined &&
     state.discoveryState === undefined
   ) {
     return null;
@@ -412,17 +482,3 @@ export async function resetAuth(serverName: string, storageDir?: string): Promis
     throw err;
   }
 }
-
-export class AuthRequiredError extends Error {
-  constructor(serverName: string, message = `Authorization required for ${serverName}. Connect in an interactive Pi session or use /mcp:auth ${serverName}.`) {
-    super(message);
-    this.name = "AuthRequiredError";
-  }
-}
-
-export class AuthLockError extends AuthRequiredError {
-  constructor(serverName: string) {
-    super(serverName, `Another Pi process is authorizing ${serverName}; if it exited, remove its stale auth lock before retrying`);
-    this.name = "AuthLockError";
-  }
-}