pi-mcp-extension 1.7.1 → 1.7.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -63,7 +63,7 @@ Config files are loaded from two locations. **Project config overrides global co
63
63
 
64
64
  | Location | Scope |
65
65
  |---|---|
66
- | `~/.pi/agent/mcp.json` | Global — applies to all projects |
66
+ | `~/.pi/agent/mcp.json` | Global — applies to all projects (honors `$PI_CODING_AGENT_DIR` when set: `<agent dir>/mcp.json`) |
67
67
  | `.pi/mcp.json` | Project — overrides global per-server |
68
68
 
69
69
  ### Full Example
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "pi-mcp-extension",
3
- "version": "1.7.1",
3
+ "version": "1.7.3",
4
4
  "description": "MCP (Model Context Protocol) client extension for the Pi coding agent — connect Pi to any MCP server",
5
5
  "type": "module",
6
6
  "keywords": [
@@ -0,0 +1,426 @@
1
+ /**
2
+ * Cross-process per-server interactive OAuth authorization lock.
3
+ *
4
+ * Uses an atomic lock file (<hash>.json.lock, mode 0600) containing a
5
+ * versioned JSON record (PID, hostname, random nonce, creation timestamp)
6
+ * plus a periodic mtime heartbeat. Automatically reclaims locks left behind
7
+ * by crashed or SIGKILLed processes using a nonce-verified rename protocol.
8
+ */
9
+
10
+ import {
11
+ open,
12
+ readFile,
13
+ stat,
14
+ rename,
15
+ unlink,
16
+ link,
17
+ utimes,
18
+ mkdir,
19
+ chmod,
20
+ } from "node:fs/promises";
21
+ import { readFileSync, unlinkSync } from "node:fs";
22
+ import { hostname as osHostname } from "node:os";
23
+ import { join } from "node:path";
24
+ import { createHash, randomUUID } from "node:crypto";
25
+ import { z } from "zod";
26
+ import { agentDir } from "./config.js";
27
+
28
+ // ─── Constants & Schema ───────────────────────────────────────────────────────
29
+
30
+ export const HEARTBEAT_MS = 30_000;
31
+ export const STALE_MS = 90_000; // 3 × 30s heartbeat
32
+
33
+ export const LockRecordSchema = z.object({
34
+ v: z.literal(1),
35
+ pid: z.number().int().positive(),
36
+ hostname: z.string().min(1),
37
+ nonce: z.string().min(1),
38
+ createdAt: z.number(),
39
+ });
40
+
41
+ export type LockRecord = z.output<typeof LockRecordSchema>;
42
+
43
+ // ─── Errors ───────────────────────────────────────────────────────────────────
44
+
45
+ export class AuthRequiredError extends Error {
46
+ constructor(
47
+ serverName: string,
48
+ message = `Authorization required for ${serverName}. Connect in an interactive Pi session or use /mcp:auth ${serverName}.`,
49
+ ) {
50
+ super(message);
51
+ this.name = "AuthRequiredError";
52
+ }
53
+ }
54
+
55
+ export class AuthLockError extends AuthRequiredError {
56
+ constructor(serverName: string, details?: string) {
57
+ const message = details
58
+ ? details.startsWith("Authorization lock") || details.startsWith("Another Pi process")
59
+ ? details
60
+ : `Another Pi process (${details}) is authorizing ${serverName}. Retry after it finishes.`
61
+ : `Another Pi process is authorizing ${serverName}. Retry after it finishes.`;
62
+ super(serverName, message);
63
+ this.name = "AuthLockError";
64
+ }
65
+ }
66
+
67
+ // ─── Dependency Injection Types ───────────────────────────────────────────────
68
+
69
+ export interface AuthLockFileHandle {
70
+ writeFile(data: string): Promise<void>;
71
+ sync(): Promise<void>;
72
+ close(): Promise<void>;
73
+ }
74
+
75
+ export interface AuthLockFs {
76
+ open(path: string, flags: string, mode?: number): Promise<AuthLockFileHandle>;
77
+ readFile(path: string, encoding: BufferEncoding): Promise<string>;
78
+ stat(path: string): Promise<{ mtimeMs: number }>;
79
+ rename(oldPath: string, newPath: string): Promise<void>;
80
+ unlink(path: string): Promise<void>;
81
+ link(existingPath: string, newPath: string): Promise<void>;
82
+ utimes(path: string, atime: number | string | Date, mtime: number | string | Date): Promise<void>;
83
+ mkdir(path: string, options?: { recursive?: boolean; mode?: number }): Promise<string | undefined | void>;
84
+ chmod(path: string, mode: number): Promise<void>;
85
+ }
86
+
87
+ export interface AuthLockTimer {
88
+ unref?: () => void;
89
+ }
90
+
91
+ export interface AuthLockClock {
92
+ now(): number;
93
+ setInterval(callback: () => void, ms: number): ReturnType<typeof setInterval> | AuthLockTimer;
94
+ clearInterval(timer: any): void;
95
+ }
96
+
97
+ export interface AuthLockDeps {
98
+ fs: AuthLockFs;
99
+ clock: AuthLockClock;
100
+ pid: number;
101
+ hostname: string;
102
+ isPidAlive: (pid: number) => boolean;
103
+ }
104
+
105
+ export interface AcquireAuthLockOptions {
106
+ onLost?: (() => void) | undefined;
107
+ deps?: Partial<AuthLockDeps> | undefined;
108
+ }
109
+
110
+ export type AuthLockHandle = (() => Promise<void>) & {
111
+ assertHeld: () => Promise<void>;
112
+ isLost: () => boolean;
113
+ };
114
+
115
+ // ─── Path & Default Dependency Helpers ────────────────────────────────────────
116
+
117
+ export function authDir(storageDir?: string): string {
118
+ return storageDir ?? join(agentDir(), "mcp-auth");
119
+ }
120
+
121
+ export function statePath(serverName: string, storageDir?: string): string {
122
+ // Hash the server name to avoid filesystem issues with special chars
123
+ const hash = createHash("sha256").update(serverName).digest("hex").slice(0, 16);
124
+ return join(authDir(storageDir), `${hash}.json`);
125
+ }
126
+
127
+ export function lockPath(serverName: string, storageDir?: string): string {
128
+ return `${statePath(serverName, storageDir)}.lock`;
129
+ }
130
+
131
+ export function defaultIsPidAlive(pid: number): boolean {
132
+ try {
133
+ process.kill(pid, 0);
134
+ return true;
135
+ } catch (err) {
136
+ return (err as NodeJS.ErrnoException).code === "EPERM";
137
+ }
138
+ }
139
+
140
+ export function createDefaultAuthLockDeps(): AuthLockDeps {
141
+ return {
142
+ fs: {
143
+ open: (path, flags, mode) => open(path, flags, mode),
144
+ readFile: (path, encoding) => readFile(path, encoding),
145
+ stat: (path) => stat(path),
146
+ rename: (oldPath, newPath) => rename(oldPath, newPath),
147
+ unlink: (path) => unlink(path),
148
+ link: (existingPath, newPath) => link(existingPath, newPath),
149
+ utimes: (path, atime, mtime) => utimes(path, atime, mtime),
150
+ mkdir: (path, options) => mkdir(path, options),
151
+ chmod: (path, mode) => chmod(path, mode),
152
+ },
153
+ clock: {
154
+ now: () => Date.now(),
155
+ setInterval: (cb, ms) => setInterval(cb, ms),
156
+ clearInterval: (timer) => clearInterval(timer),
157
+ },
158
+ pid: process.pid,
159
+ hostname: osHostname(),
160
+ isPidAlive: defaultIsPidAlive,
161
+ };
162
+ }
163
+
164
+ // ─── Pure Staleness Policy ────────────────────────────────────────────────────
165
+
166
+ export function parseLockRecord(raw: string): LockRecord | null {
167
+ if (!raw) return null;
168
+ try {
169
+ const parsed: unknown = JSON.parse(raw);
170
+ const result = LockRecordSchema.safeParse(parsed);
171
+ return result.success ? result.data : null;
172
+ } catch {
173
+ return null;
174
+ }
175
+ }
176
+
177
+ /**
178
+ * Pure policy determining whether an existing lock file is stale and safe to reclaim.
179
+ */
180
+ export function isLockStale(
181
+ record: LockRecord | null,
182
+ mtimeMs: number,
183
+ now: number,
184
+ self: { hostname: string; isPidAlive: (pid: number) => boolean },
185
+ ): boolean {
186
+ if (now - mtimeMs > STALE_MS) {
187
+ return true;
188
+ }
189
+ if (record === null) {
190
+ return false;
191
+ }
192
+ if (record.hostname === self.hostname && !self.isPidAlive(record.pid)) {
193
+ return true;
194
+ }
195
+ return false;
196
+ }
197
+
198
+ function formatHolderDetails(record: LockRecord | null): string | undefined {
199
+ if (!record) return undefined;
200
+ return `pid ${record.pid} on ${record.hostname}`;
201
+ }
202
+
203
+ // ─── Lock Acquisition, Heartbeat & Release ────────────────────────────────────
204
+
205
+ const MAX_ACQUIRE_ATTEMPTS = 3;
206
+
207
+ export async function acquireAuthLock(
208
+ serverName: string,
209
+ storageDir?: string,
210
+ options?: AcquireAuthLockOptions,
211
+ ): Promise<AuthLockHandle> {
212
+ const defaults = createDefaultAuthLockDeps();
213
+ const deps: AuthLockDeps = {
214
+ fs: options?.deps?.fs ?? defaults.fs,
215
+ clock: options?.deps?.clock ?? defaults.clock,
216
+ pid: options?.deps?.pid ?? defaults.pid,
217
+ hostname: options?.deps?.hostname ?? defaults.hostname,
218
+ isPidAlive: options?.deps?.isPidAlive ?? defaults.isPidAlive,
219
+ };
220
+ const usingRealFs = options?.deps?.fs === undefined;
221
+
222
+ const directory = authDir(storageDir);
223
+ await deps.fs.mkdir(directory, { recursive: true, mode: 0o700 });
224
+ await deps.fs.chmod(directory, 0o700);
225
+ const path = lockPath(serverName, storageDir);
226
+
227
+ for (let attempt = 0; attempt < MAX_ACQUIRE_ATTEMPTS; attempt++) {
228
+ const myRecord: LockRecord = {
229
+ v: 1,
230
+ pid: deps.pid,
231
+ hostname: deps.hostname,
232
+ nonce: randomUUID(),
233
+ createdAt: deps.clock.now(),
234
+ };
235
+
236
+ try {
237
+ const file = await deps.fs.open(path, "wx", 0o600);
238
+ try {
239
+ await file.writeFile(JSON.stringify(myRecord));
240
+ await file.sync();
241
+ } catch (writeErr) {
242
+ await file.close().catch(() => {});
243
+ await deps.fs.unlink(path).catch(() => {});
244
+ throw writeErr;
245
+ }
246
+ await file.close();
247
+
248
+ let released = false;
249
+ let lost = false;
250
+ let heartbeatTimer: ReturnType<typeof setInterval> | AuthLockTimer | null = null;
251
+
252
+ const stopHeartbeat = () => {
253
+ if (heartbeatTimer !== null) {
254
+ deps.clock.clearInterval(heartbeatTimer);
255
+ heartbeatTimer = null;
256
+ }
257
+ };
258
+
259
+ const onExit = () => {
260
+ if (released || lost || !usingRealFs) return;
261
+ try {
262
+ const raw = readFileSync(path, "utf8");
263
+ const current = parseLockRecord(raw);
264
+ if (current && current.nonce === myRecord.nonce) {
265
+ unlinkSync(path);
266
+ }
267
+ } catch {
268
+ // Best-effort cleanup on process exit
269
+ }
270
+ };
271
+
272
+ const markLost = () => {
273
+ if (lost || released) return;
274
+ lost = true;
275
+ stopHeartbeat();
276
+ process.removeListener("exit", onExit);
277
+ options?.onLost?.();
278
+ };
279
+
280
+ process.once("exit", onExit);
281
+
282
+ heartbeatTimer = deps.clock.setInterval(async () => {
283
+ if (released || lost) return;
284
+ let current: LockRecord | null = null;
285
+ try {
286
+ const raw = await deps.fs.readFile(path, "utf8");
287
+ current = parseLockRecord(raw);
288
+ } catch {
289
+ current = null;
290
+ }
291
+ if (released || lost) return;
292
+ if (!current || current.nonce !== myRecord.nonce) {
293
+ markLost();
294
+ return;
295
+ }
296
+ const nowSec = deps.clock.now() / 1000;
297
+ try {
298
+ await deps.fs.utimes(path, nowSec, nowSec);
299
+ } catch {
300
+ if (!released && !lost) {
301
+ markLost();
302
+ }
303
+ }
304
+ }, HEARTBEAT_MS);
305
+ heartbeatTimer.unref?.();
306
+
307
+ const assertHeld = async (): Promise<void> => {
308
+ if (lost || released) {
309
+ throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
310
+ }
311
+ let current: LockRecord | null = null;
312
+ try {
313
+ const raw = await deps.fs.readFile(path, "utf8");
314
+ current = parseLockRecord(raw);
315
+ } catch {
316
+ current = null;
317
+ }
318
+ if (!current || current.nonce !== myRecord.nonce) {
319
+ markLost();
320
+ throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
321
+ }
322
+ };
323
+
324
+ const unlock = async (): Promise<void> => {
325
+ if (released) return;
326
+ released = true;
327
+ stopHeartbeat();
328
+ process.removeListener("exit", onExit);
329
+ if (lost) return;
330
+ let current: LockRecord | null = null;
331
+ try {
332
+ const raw = await deps.fs.readFile(path, "utf8");
333
+ current = parseLockRecord(raw);
334
+ } catch {
335
+ current = null;
336
+ }
337
+ if (current && current.nonce === myRecord.nonce) {
338
+ await deps.fs.unlink(path).catch(() => {});
339
+ }
340
+ };
341
+
342
+ return Object.assign(unlock, {
343
+ assertHeld,
344
+ isLost: () => lost,
345
+ });
346
+ } catch (err) {
347
+ if ((err as NodeJS.ErrnoException).code !== "EEXIST") {
348
+ throw err;
349
+ }
350
+
351
+ // Lock file already exists — inspect and check staleness
352
+ let statResult: { mtimeMs: number };
353
+ try {
354
+ statResult = await deps.fs.stat(path);
355
+ } catch (statErr) {
356
+ if ((statErr as NodeJS.ErrnoException).code === "ENOENT") {
357
+ continue;
358
+ }
359
+ throw statErr;
360
+ }
361
+
362
+ let record: LockRecord | null = null;
363
+ try {
364
+ const raw = await deps.fs.readFile(path, "utf8");
365
+ record = parseLockRecord(raw);
366
+ } catch (readErr) {
367
+ if ((readErr as NodeJS.ErrnoException).code === "ENOENT") {
368
+ continue;
369
+ }
370
+ record = null;
371
+ }
372
+
373
+ const stale = isLockStale(record, statResult.mtimeMs, deps.clock.now(), {
374
+ hostname: deps.hostname,
375
+ isPidAlive: deps.isPidAlive,
376
+ });
377
+ if (!stale) {
378
+ throw new AuthLockError(serverName, formatHolderDetails(record));
379
+ }
380
+
381
+ // Reclaim protocol: move stale candidate aside, verify identity, then delete or restore
382
+ const movedPath = `${path}.stale-${randomUUID()}`;
383
+ try {
384
+ await deps.fs.rename(path, movedPath);
385
+ } catch (renameErr) {
386
+ if ((renameErr as NodeJS.ErrnoException).code === "ENOENT") {
387
+ continue;
388
+ }
389
+ throw renameErr;
390
+ }
391
+
392
+ let movedRecord: LockRecord | null = null;
393
+ try {
394
+ const movedRaw = await deps.fs.readFile(movedPath, "utf8");
395
+ movedRecord = parseLockRecord(movedRaw);
396
+ } catch {
397
+ movedRecord = null;
398
+ }
399
+
400
+ let matchesStale = false;
401
+ if (record !== null) {
402
+ matchesStale = movedRecord !== null && movedRecord.nonce === record.nonce;
403
+ } else {
404
+ try {
405
+ const movedStat = await deps.fs.stat(movedPath);
406
+ matchesStale = movedRecord === null && movedStat.mtimeMs === statResult.mtimeMs;
407
+ } catch {
408
+ matchesStale = false;
409
+ }
410
+ }
411
+
412
+ if (matchesStale) {
413
+ await deps.fs.unlink(movedPath).catch(() => {});
414
+ continue;
415
+ }
416
+
417
+ await deps.fs.link(movedPath, path).catch(() => {});
418
+ await deps.fs.unlink(movedPath).catch(() => {});
419
+ throw new AuthLockError(serverName, formatHolderDetails(movedRecord));
420
+ }
421
+ }
422
+
423
+ throw new AuthLockError(serverName);
424
+ }
425
+
426
+ export { acquireAuthLock as acquireLock };
@@ -193,7 +193,11 @@ export function callbackServerConfigFromRedirectUrl(redirectUrl?: string): Callb
193
193
  throw new Error(`Local OAuth redirect URL must use path ${CALLBACK_PATH}`);
194
194
  }
195
195
 
196
- const preferredPort = Number(url.port || "80");
196
+ if (!url.port) {
197
+ throw new Error("Manual OAuth redirect URL must include an explicit port (e.g. http://127.0.0.1:19876/callback)");
198
+ }
199
+
200
+ const preferredPort = Number(url.port);
197
201
  if (preferredPort === 0) {
198
202
  throw new Error("Manual OAuth redirect URL must use a fixed non-zero port");
199
203
  }
package/src/index.ts CHANGED
@@ -24,7 +24,7 @@ import {
24
24
  stopCallbackServer,
25
25
  callbackServerConfigFromRedirectUrl,
26
26
  } from "./callback-server.js";
27
- import { setCallbackPort, McpOAuthProvider, acquireAuthLock } from "./oauth-provider.js";
27
+ import { setCallbackPort, McpOAuthProvider, acquireAuthLock, type AuthLockHandle } from "./oauth-provider.js";
28
28
  import { auth } from "@modelcontextprotocol/sdk/client/auth.js";
29
29
  import { discoverManualAuthChallenge } from "./auth-challenge.js";
30
30
  import type { ManualAuthChallenge } from "./auth-challenge.js";
@@ -367,14 +367,23 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
367
367
  const config = server.config;
368
368
  let oauthState: string | undefined;
369
369
  let latestAuthorizationUrl: URL | undefined;
370
+ let lockLost = false;
370
371
  const cancel = () => { if (oauthState) cancelCallback(oauthState); };
371
372
  signal.addEventListener("abort", cancel, { once: true });
372
- let unlock: (() => Promise<void>) | undefined;
373
+ let unlock: AuthLockHandle | undefined;
373
374
  try {
374
375
  if (signal.aborted) throw new AuthCancelledError(serverName);
375
- unlock = await acquireAuthLock(serverName, paths.authStorageDir);
376
+ unlock = await acquireAuthLock(serverName, paths.authStorageDir, {
377
+ onLost: () => {
378
+ lockLost = true;
379
+ cancel();
380
+ },
381
+ });
376
382
  if (signal.aborted) throw new AuthCancelledError(serverName);
377
- if (reset) await manager.resetServerAuth(serverName);
383
+ if (reset) {
384
+ await unlock.assertHeld();
385
+ await manager.resetServerAuth(serverName);
386
+ }
378
387
  // Validate that we have a server URL (required for OAuth)
379
388
  if (!config.url) {
380
389
  throw new McpError(
@@ -411,16 +420,20 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
411
420
  config.auth || { type: "oauth" },
412
421
  (url: URL) => {
413
422
  latestAuthorizationUrl = new URL(url.toString());
414
- if (!signal.aborted) (paths.browserOpen ?? openBrowser)(url.toString());
423
+ if (!signal.aborted && !lockLost) (paths.browserOpen ?? openBrowser)(url.toString());
415
424
  },
416
425
  paths.authStorageDir,
417
426
  false,
418
427
  () => signal.aborted,
428
+ () => unlock!.assertHeld(),
419
429
  );
420
430
 
421
- // Set the state before auth() builds the authorization URL.
431
+ // Why set state before calling auth(): the SDK reads provider.state() while
432
+ // building the authorization URL inside auth(), and omits the state parameter
433
+ // if empty, which the callback server then rejects as a CSRF error.
422
434
  authProvider.setState(oauthState);
423
435
 
436
+ // 4. Discover protected-resource challenge metadata from the server.
424
437
  const authChallenge: ManualAuthChallenge = await discoverManualAuthChallenge(config.url, {
425
438
  headers: config.headers,
426
439
  }).catch((): ManualAuthChallenge => {
@@ -430,7 +443,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
430
443
  return {};
431
444
  });
432
445
 
433
- // Register the callback immediately before auth() can open the browser.
446
+ // 5. Register the callback immediately before auth() can open the browser.
434
447
  // Attach a catch handler so cancellation after immediate authorization
435
448
  // does not produce an unhandled rejection.
436
449
  const callbackPromise = waitForCallback(oauthState);
@@ -448,9 +461,14 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
448
461
  authOptions.scope = authChallenge.scope;
449
462
  }
450
463
 
451
- // Start the auth flow. REDIRECT means browser interaction is required.
464
+ // 6. Start the auth flow.
465
+ // Why branch on the return value: the SDK auth() call returns "AUTHORIZED" when
466
+ // stored or refreshed tokens already succeeded without opening a browser, or
467
+ // "REDIRECT" after invoking redirectToAuthorization(), in which case we must
468
+ // wait for the callback code and call auth() a second time with authorizationCode.
452
469
  const authResult = await auth(authProvider, authOptions);
453
470
  if (signal.aborted) throw new AuthCancelledError(serverName);
471
+ if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
454
472
 
455
473
  if (authResult === "AUTHORIZED") {
456
474
  // Auth succeeded without needing browser interaction (e.g., had valid tokens).
@@ -462,7 +480,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
462
480
  "info",
463
481
  );
464
482
 
465
- // 6. Wait for the callback while giving the user an escape hatch.
483
+ // 7. Wait for the callback while giving the user an escape hatch.
466
484
  const code = await waitForOAuthCallback(
467
485
  serverName,
468
486
  callbackPromise,
@@ -474,16 +492,17 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
474
492
  "protocol",
475
493
  );
476
494
  }
477
- if (!signal.aborted) (paths.browserOpen ?? openBrowser)(latestAuthorizationUrl.toString());
495
+ if (!signal.aborted && !lockLost) (paths.browserOpen ?? openBrowser)(latestAuthorizationUrl.toString());
478
496
  },
479
497
  ctx,
480
498
  );
481
499
  if (signal.aborted) throw new AuthCancelledError(serverName);
500
+ if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
482
501
  ctx.ui.notify(
483
502
  `pi-mcp: Authorization callback received for ${serverName}. Exchanging token...`,
484
503
  "info",
485
504
  );
486
- // 7. Complete the OAuth flow with the authorization code
505
+ // 8. Complete the OAuth flow with the authorization code
487
506
  const finishAuthOptions: {
488
507
  serverUrl: string;
489
508
  authorizationCode: string;
@@ -520,6 +539,12 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
520
539
  }
521
540
 
522
541
  if (signal.aborted) throw new AuthCancelledError(serverName);
542
+ if (lockLost) throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
543
+ } catch (err) {
544
+ if (lockLost && !signal.aborted) {
545
+ throw new AuthLockError(serverName, `Authorization lock lost for ${serverName}; retry`);
546
+ }
547
+ throw err;
523
548
  } finally {
524
549
  signal.removeEventListener("abort", cancel);
525
550
  cancel();
@@ -667,7 +692,7 @@ export default async function (pi: ExtensionAPI, paths: { bootstrapCwd?: string;
667
692
  : !manager.getServer(name)
668
693
  ? "Unknown MCP server name. Use mcp_status to list configured servers."
669
694
  : err instanceof AuthLockError
670
- ? "Another Pi process is authorizing this server. Retry after it finishes or check for a stale auth lock."
695
+ ? "Another Pi process is authorizing this server. Retry after it finishes."
671
696
  : err instanceof AuthRequiredError
672
697
  ? ctx.hasUI
673
698
  ? "MCP authorization required. Retry with mcp_connect or /mcp:auth."
@@ -13,10 +13,21 @@
13
13
 
14
14
  import type { OAuthClientProvider, OAuthDiscoveryState } from "@modelcontextprotocol/sdk/client/auth.js";
15
15
  import type { OAuthClientMetadata, OAuthClientInformationMixed, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
16
- import { readFile, writeFile, mkdir, unlink, chmod, open, stat } from "node:fs/promises";
17
- import { join } from "node:path";
18
- import { createHash } from "node:crypto";
19
- import { agentDir } from "./config.js";
16
+ import { readFile, writeFile, mkdir, unlink, chmod, rename } from "node:fs/promises";
17
+ import { randomUUID } from "node:crypto";
18
+ import { z } from "zod";
19
+ import { McpError } from "./errors.js";
20
+ import {
21
+ acquireAuthLock,
22
+ AuthLockError,
23
+ AuthRequiredError,
24
+ authDir,
25
+ statePath,
26
+ type AcquireAuthLockOptions,
27
+ type AuthLockHandle,
28
+ } from "./auth-lock.js";
29
+
30
+ export { acquireAuthLock, AuthLockError, AuthRequiredError, type AcquireAuthLockOptions, type AuthLockHandle };
20
31
 
21
32
  // ─── Types ────────────────────────────────────────────────────────────────────
22
33
 
@@ -57,7 +68,7 @@ export function getCallbackPort(): number {
57
68
  return callbackPort;
58
69
  }
59
70
 
60
- // ─── Persistent State Types ───────────────────────────────────────────────────
71
+ // ─── Persistent State Types & Schema ──────────────────────────────────────────
61
72
 
62
73
  interface StoredClientInfo {
63
74
  client_id: string;
@@ -77,74 +88,113 @@ interface StoredTokens {
77
88
  interface StoredState {
78
89
  clientInfo: StoredClientInfo | undefined;
79
90
  tokens: StoredTokens | undefined;
80
- codeVerifier: string | undefined;
81
91
  discoveryState: OAuthDiscoveryState | undefined;
82
92
  }
83
93
 
84
- // ─── File helpers ─────────────────────────────────────────────────────────────
85
-
86
- function authDir(storageDir?: string): string {
87
- return storageDir ?? join(agentDir(), "mcp-auth");
88
- }
94
+ const StoredClientInfoSchema = z.object({
95
+ client_id: z.string(),
96
+ client_secret: z.string().optional(),
97
+ });
98
+
99
+ const StoredTokensSchema = z.object({
100
+ access_token: z.string(),
101
+ token_type: z.string().optional(),
102
+ refresh_token: z.string().optional(),
103
+ expires_in: z.number().optional(),
104
+ scope: z.string().optional(),
105
+ saved_at: z.string().optional(),
106
+ });
107
+
108
+ export const StoredStateSchema = z.object({
109
+ clientInfo: StoredClientInfoSchema.optional(),
110
+ tokens: StoredTokensSchema.optional(),
111
+ discoveryState: z.record(z.unknown()).optional(),
112
+ });
89
113
 
90
- function statePath(serverName: string, storageDir?: string): string {
91
- // Hash the server name to avoid filesystem issues with special chars
92
- const hash = createHash("sha256").update(serverName).digest("hex").slice(0, 16);
93
- return join(authDir(storageDir), `${hash}.json`);
94
- }
114
+ // ─── File helpers ─────────────────────────────────────────────────────────────
95
115
 
96
- // Exclusive per-server authorization lock shared by independent Pi processes.
97
- export async function acquireAuthLock(serverName: string, storageDir?: string): Promise<() => Promise<void>> {
98
- const directory = authDir(storageDir);
99
- await mkdir(directory, { recursive: true, mode: 0o700 });
100
- await chmod(directory, 0o700);
101
- const path = `${statePath(serverName, storageDir)}.lock`;
116
+ async function loadState(serverName: string, storageDir?: string): Promise<StoredState> {
117
+ const path = statePath(serverName, storageDir);
118
+ let raw: string;
102
119
  try {
103
- const file = await open(path, "wx", 0o600);
104
- try { await file.writeFile(String(process.pid)); }
105
- catch (err) { await unlink(path).catch(() => {}); throw err; }
106
- finally { await file.close(); }
107
- return async () => { await unlink(path).catch(() => {}); };
120
+ raw = await readFile(path, "utf8");
108
121
  } catch (err) {
109
- if ((err as NodeJS.ErrnoException).code !== "EEXIST") throw err;
110
- throw new AuthLockError(serverName);
122
+ if ((err as NodeJS.ErrnoException).code === "ENOENT") {
123
+ return {
124
+ clientInfo: undefined,
125
+ tokens: undefined,
126
+ discoveryState: undefined,
127
+ };
128
+ }
129
+ throw err;
111
130
  }
112
- }
113
131
 
114
- async function loadState(serverName: string, storageDir?: string): Promise<StoredState> {
132
+ let parsed: z.output<typeof StoredStateSchema>;
115
133
  try {
116
- const raw = await readFile(statePath(serverName, storageDir), "utf8");
117
- const parsed = JSON.parse(raw) as StoredState;
118
- return {
119
- clientInfo: parsed.clientInfo ?? undefined,
120
- tokens: parsed.tokens ?? undefined,
121
- codeVerifier: parsed.codeVerifier ?? undefined,
122
- discoveryState: parsed.discoveryState ?? undefined,
123
- };
124
- } catch {
125
- return {
126
- clientInfo: undefined,
127
- tokens: undefined,
128
- codeVerifier: undefined,
129
- discoveryState: undefined,
130
- };
134
+ const json: unknown = JSON.parse(raw);
135
+ parsed = StoredStateSchema.parse(json);
136
+ } catch (err) {
137
+ const backupPath = `${path}.corrupt-${Date.now()}`;
138
+ await rename(path, backupPath);
139
+ throw new McpError(
140
+ `Corrupt OAuth state file for "${serverName}" backed up to ${backupPath}`,
141
+ serverName,
142
+ "config",
143
+ err,
144
+ );
131
145
  }
146
+
147
+ return {
148
+ clientInfo: parsed.clientInfo
149
+ ? {
150
+ client_id: parsed.clientInfo.client_id,
151
+ client_secret: parsed.clientInfo.client_secret,
152
+ }
153
+ : undefined,
154
+ tokens: parsed.tokens
155
+ ? {
156
+ access_token: parsed.tokens.access_token,
157
+ token_type: parsed.tokens.token_type,
158
+ refresh_token: parsed.tokens.refresh_token,
159
+ expires_in: parsed.tokens.expires_in,
160
+ scope: parsed.tokens.scope,
161
+ saved_at: parsed.tokens.saved_at,
162
+ }
163
+ : undefined,
164
+ discoveryState: parsed.discoveryState as OAuthDiscoveryState | undefined,
165
+ };
132
166
  }
133
167
 
134
- async function saveState(serverName: string, state: StoredState, storageDir?: string, isCancelled: () => boolean = () => false): Promise<void> {
168
+ async function saveState(
169
+ serverName: string,
170
+ state: StoredState,
171
+ storageDir?: string,
172
+ isCancelled: () => boolean = () => false,
173
+ ): Promise<void> {
135
174
  const directory = authDir(storageDir);
136
175
  await mkdir(directory, { recursive: true, mode: 0o700 });
137
176
  await chmod(directory, 0o700);
138
- // Only write defined fields
177
+ // Only write defined fields (never write codeVerifier to shared state)
139
178
  const toWrite: Record<string, unknown> = {};
140
179
  if (state.clientInfo !== undefined) toWrite.clientInfo = state.clientInfo;
141
180
  if (state.tokens !== undefined) toWrite.tokens = state.tokens;
142
- if (state.codeVerifier !== undefined) toWrite.codeVerifier = state.codeVerifier;
143
181
  if (state.discoveryState !== undefined) toWrite.discoveryState = state.discoveryState;
144
182
  if (isCancelled()) throw new AuthRequiredError(serverName, "Authorization cancelled");
145
183
  const path = statePath(serverName, storageDir);
146
- await writeFile(path, JSON.stringify(toWrite, null, 2), { encoding: "utf8", mode: 0o600 });
147
- await chmod(path, 0o600);
184
+ const tmpPath = `${path}.tmp-${randomUUID()}`;
185
+ try {
186
+ await writeFile(tmpPath, JSON.stringify(toWrite, null, 2), { encoding: "utf8", mode: 0o600 });
187
+ await chmod(tmpPath, 0o600);
188
+ if (isCancelled()) {
189
+ await unlink(tmpPath).catch(() => {});
190
+ throw new AuthRequiredError(serverName, "Authorization cancelled");
191
+ }
192
+ await rename(tmpPath, path);
193
+ await chmod(path, 0o600);
194
+ } catch (err) {
195
+ await unlink(tmpPath).catch(() => {});
196
+ throw err;
197
+ }
148
198
  }
149
199
 
150
200
  // ─── OAuthClientProvider Implementation ────────────────────────────────────────
@@ -155,6 +205,8 @@ export class McpOAuthProvider implements OAuthClientProvider {
155
205
  private _redirectUrl: string | undefined;
156
206
  private _onAuthRequired: ((url: URL) => void | Promise<void>) | undefined;
157
207
  private _oauthState: string | undefined;
208
+ private _codeVerifier: string | undefined;
209
+ private _cachedTokens: StoredTokens | undefined;
158
210
 
159
211
  constructor(
160
212
  serverName: string,
@@ -163,6 +215,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
163
215
  private readonly storageDir?: string,
164
216
  private readonly silent = false,
165
217
  private readonly isCancelled: () => boolean = () => false,
218
+ private readonly assertLockHeld?: (() => Promise<void>) | undefined,
166
219
  ) {
167
220
  this.serverName = serverName;
168
221
  this.authConfig = authConfig;
@@ -217,19 +270,21 @@ export class McpOAuthProvider implements OAuthClientProvider {
217
270
 
218
271
  async saveClientInformation(clientInformation: OAuthClientInformationMixed): Promise<void> {
219
272
  if (this.silent) throw new AuthRequiredError(this.serverName);
273
+ await this.assertLockHeld?.();
220
274
  const state = await loadState(this.serverName, this.storageDir);
221
275
  state.clientInfo = {
222
276
  client_id: clientInformation.client_id,
223
277
  client_secret: clientInformation.client_secret,
224
278
  };
279
+ await this.assertLockHeld?.();
225
280
  await saveState(this.serverName, state, this.storageDir, this.isCancelled);
226
281
  }
227
282
 
228
283
  // --- tokens ---
229
284
 
230
285
  async tokens(): Promise<OAuthTokens | undefined> {
231
- const state = await loadState(this.serverName, this.storageDir);
232
- if (!state.tokens) return undefined;
286
+ const storedTokens = this._cachedTokens ?? (await loadState(this.serverName, this.storageDir)).tokens;
287
+ if (!storedTokens) return undefined;
233
288
 
234
289
  // Always return stored tokens — even if expired.
235
290
  // The SDK's auth() function checks tokens?.refresh_token and attempts
@@ -247,19 +302,17 @@ export class McpOAuthProvider implements OAuthClientProvider {
247
302
 
248
303
  // Build OAuthTokens — only include defined fields
249
304
  const tokens: Record<string, string> = {
250
- access_token: state.tokens.access_token,
305
+ access_token: storedTokens.access_token,
251
306
  };
252
- if (state.tokens.token_type !== undefined) tokens.token_type = state.tokens.token_type;
253
- if (state.tokens.refresh_token !== undefined) tokens.refresh_token = state.tokens.refresh_token;
307
+ if (storedTokens.token_type !== undefined) tokens.token_type = storedTokens.token_type;
308
+ if (storedTokens.refresh_token !== undefined) tokens.refresh_token = storedTokens.refresh_token;
254
309
  return tokens as unknown as OAuthTokens;
255
310
  }
256
311
 
257
312
  // --- saveTokens ---
258
313
 
259
314
  async saveTokens(tokens: OAuthTokens): Promise<void> {
260
- if (this.silent) await this.checkInteractiveLock();
261
- const state = await loadState(this.serverName, this.storageDir);
262
- state.tokens = {
315
+ const storedTokens: StoredTokens = {
263
316
  access_token: tokens.access_token,
264
317
  token_type: tokens.token_type,
265
318
  refresh_token: tokens.refresh_token,
@@ -267,8 +320,28 @@ export class McpOAuthProvider implements OAuthClientProvider {
267
320
  scope: tokens.scope,
268
321
  saved_at: new Date().toISOString(),
269
322
  };
323
+
324
+ if (this.silent) {
325
+ try {
326
+ const state = await loadState(this.serverName, this.storageDir);
327
+ state.tokens = storedTokens;
328
+ await saveState(this.serverName, state, this.storageDir, this.isCancelled);
329
+ this._cachedTokens = undefined;
330
+ } catch (err) {
331
+ this._cachedTokens = storedTokens;
332
+ console.error(`[pi-mcp] Failed to persist refreshed OAuth tokens for "${this.serverName}":`, err);
333
+ }
334
+ return;
335
+ }
336
+
337
+ if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
338
+ await this.assertLockHeld?.();
339
+ const state = await loadState(this.serverName, this.storageDir);
340
+ state.tokens = storedTokens;
270
341
  if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
342
+ await this.assertLockHeld?.();
271
343
  await saveState(this.serverName, state, this.storageDir, this.isCancelled);
344
+ this._cachedTokens = undefined;
272
345
  }
273
346
 
274
347
  // --- redirectToAuthorization ---
@@ -285,30 +358,30 @@ export class McpOAuthProvider implements OAuthClientProvider {
285
358
  }
286
359
  }
287
360
 
288
- // --- PKCE code verifier ---
361
+ // --- PKCE code verifier (in-memory per provider instance) ---
289
362
 
290
363
  async saveCodeVerifier(codeVerifier: string): Promise<void> {
291
- // The SDK saves PKCE before redirecting. Silent attempts must not overwrite an active flow.
364
+ // The SDK saves PKCE before redirecting. Silent attempts must not start a browser flow.
292
365
  if (this.silent) throw new AuthRequiredError(this.serverName);
293
- const state = await loadState(this.serverName, this.storageDir);
294
- state.codeVerifier = codeVerifier;
295
- await saveState(this.serverName, state, this.storageDir, this.isCancelled);
366
+ if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
367
+ this._codeVerifier = codeVerifier;
296
368
  }
297
369
 
298
370
  async codeVerifier(): Promise<string> {
299
- const state = await loadState(this.serverName, this.storageDir);
300
- if (!state.codeVerifier) {
371
+ if (!this._codeVerifier) {
301
372
  throw new Error(`[pi-mcp] No PKCE code verifier found for "${this.serverName}"`);
302
373
  }
303
- return state.codeVerifier;
374
+ return this._codeVerifier;
304
375
  }
305
376
 
306
377
  // --- Discovery state caching ---
307
378
 
308
379
  async saveDiscoveryState(discState: OAuthDiscoveryState): Promise<void> {
309
380
  if (this.silent) return; // No shared state writes from background discovery.
381
+ await this.assertLockHeld?.();
310
382
  const state = await loadState(this.serverName, this.storageDir);
311
383
  state.discoveryState = discState;
384
+ await this.assertLockHeld?.();
312
385
  await saveState(this.serverName, state, this.storageDir, this.isCancelled);
313
386
  }
314
387
 
@@ -338,35 +411,33 @@ export class McpOAuthProvider implements OAuthClientProvider {
338
411
 
339
412
  // --- Credential invalidation ---
340
413
 
341
- private async checkInteractiveLock(): Promise<void> {
342
- try { await stat(`${statePath(this.serverName, this.storageDir)}.lock`); }
343
- catch (err) { if ((err as NodeJS.ErrnoException).code === "ENOENT") return; throw err; }
344
- throw new AuthRequiredError(this.serverName);
345
- }
346
-
347
414
  async invalidateCredentials(scope: "all" | "client" | "tokens" | "verifier" | "discovery"): Promise<void> {
348
415
  if (this.silent) throw new AuthRequiredError(this.serverName);
416
+ await this.assertLockHeld?.();
349
417
  const state = await loadState(this.serverName, this.storageDir);
350
418
  switch (scope) {
351
419
  case "all":
352
420
  state.clientInfo = undefined;
353
421
  state.tokens = undefined;
354
- state.codeVerifier = undefined;
355
422
  state.discoveryState = undefined;
423
+ this._codeVerifier = undefined;
424
+ this._cachedTokens = undefined;
356
425
  break;
357
426
  case "client":
358
427
  state.clientInfo = undefined;
359
428
  break;
360
429
  case "tokens":
361
430
  state.tokens = undefined;
431
+ this._cachedTokens = undefined;
362
432
  break;
363
433
  case "verifier":
364
- state.codeVerifier = undefined;
434
+ this._codeVerifier = undefined;
365
435
  break;
366
436
  case "discovery":
367
437
  state.discoveryState = undefined;
368
438
  break;
369
439
  }
440
+ await this.assertLockHeld?.();
370
441
  await saveState(this.serverName, state, this.storageDir, this.isCancelled);
371
442
  }
372
443
  }
@@ -387,7 +458,6 @@ export async function getAuthStatus(serverName: string, storageDir?: string): Pr
387
458
  if (
388
459
  state.clientInfo === undefined &&
389
460
  state.tokens === undefined &&
390
- state.codeVerifier === undefined &&
391
461
  state.discoveryState === undefined
392
462
  ) {
393
463
  return null;
@@ -405,19 +475,10 @@ export async function getAuthStatus(serverName: string, storageDir?: string): Pr
405
475
  * Used to force re-authorization on next connection.
406
476
  */
407
477
  export async function resetAuth(serverName: string, storageDir?: string): Promise<void> {
408
- await unlink(statePath(serverName, storageDir)).catch(() => {});
409
- }
410
-
411
- export class AuthRequiredError extends Error {
412
- constructor(serverName: string, message = `Authorization required for ${serverName}. Connect in an interactive Pi session or use /mcp:auth ${serverName}.`) {
413
- super(message);
414
- this.name = "AuthRequiredError";
415
- }
416
- }
417
-
418
- export class AuthLockError extends AuthRequiredError {
419
- constructor(serverName: string) {
420
- super(serverName, `Another Pi process is authorizing ${serverName}; if it exited, remove its stale auth lock before retrying`);
421
- this.name = "AuthLockError";
478
+ try {
479
+ await unlink(statePath(serverName, storageDir));
480
+ } catch (err) {
481
+ if ((err as NodeJS.ErrnoException).code === "ENOENT") return;
482
+ throw err;
422
483
  }
423
484
  }
@@ -25,7 +25,16 @@ import {
25
25
  import { McpError } from "./errors.js";
26
26
  import type { McpConfig, ServerConfig, Settings } from "./config.js";
27
27
  import { McpOAuthProvider, AuthRequiredError, getAuthStatus, resetAuth } from "./oauth-provider.js";
28
- import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js";
28
+ import { UnauthorizedError, type OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js";
29
+
30
+ export function isOAuthUnauthorizedError(error: unknown): boolean {
31
+ return (
32
+ error instanceof AuthRequiredError ||
33
+ error instanceof UnauthorizedError ||
34
+ (typeof (error as { code?: unknown })?.code === "number" && (error as { code: number }).code === 401) ||
35
+ /\b(401|Unauthorized)\b/i.test(error instanceof Error ? error.message : String(error))
36
+ );
37
+ }
29
38
 
30
39
  // ─── Types ────────────────────────────────────────────────────────────────────
31
40
 
@@ -356,7 +365,7 @@ export class ServerManager {
356
365
  try {
357
366
  await client.connect(transport);
358
367
  if (signal.aborted) throw new Error("Connection cancelled");
359
- if (server.config.transport === "stdio") server.childPid = (transport as any).process?.pid ?? null;
368
+ server.childPid = transport instanceof StdioClientTransport ? transport.pid : null;
360
369
  // Readiness includes tool discovery, not just the MCP handshake.
361
370
  if (this.onToolRefresh) await this.onToolRefresh(server.name, client);
362
371
  if (signal.aborted) throw new Error("Connection cancelled");
@@ -388,10 +397,16 @@ export class ServerManager {
388
397
  if (signal.aborted) throw new Error("Connection cancelled");
389
398
  const error = err instanceof Error ? err : new Error(String(err));
390
399
  server.lastError = error;
391
- if (error instanceof AuthRequiredError || error instanceof McpError && error.code === "protocol") {
400
+ if (error instanceof AuthRequiredError || (error instanceof McpError && error.code === "protocol")) {
392
401
  server.state = "stopped";
393
402
  throw error;
394
403
  }
404
+ if (Boolean(server.config.auth) && server.config.transport !== "stdio" && isOAuthUnauthorizedError(error)) {
405
+ const authError = new AuthRequiredError(server.name);
406
+ server.lastError = authError;
407
+ server.state = "stopped";
408
+ throw authError;
409
+ }
395
410
  if (server.retryCount >= this.settings.maxRetries) {
396
411
  server.state = "stopped";
397
412
  throw new McpError(`Connection failed after ${server.retryCount} retries: ${error.message}`, server.name, "connection", err);