diff --git a/shared/glean/mcp/src/auth-provider.ts b/shared/glean/mcp/src/auth-provider.ts index 5b68b99..2e204bb 100644 --- a/shared/glean/mcp/src/auth-provider.ts +++ b/shared/glean/mcp/src/auth-provider.ts @@ -9,10 +9,18 @@ import { randomUUID } from "node:crypto"; import { platform } from "node:os"; import { getCallbackUrl, setExpectedState } from "./auth-callback-server.js"; import { + acquireDataFileLockSync, + releaseDataFileLock, + type DataLockHandle, +} from "./data-dir.js"; +import { + acquireClientRegistrationLock, clearCredentials, credentialsMtimeMs, loadCredentials, + releaseClientRegistrationLock, saveCredentials, + type CredentialMetadata, } from "./token-store.js"; export type InvalidationScope = "all" | "client" | "tokens" | "verifier"; @@ -34,6 +42,11 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } +export function normalizeAccountEmail(email: string | undefined): string | undefined { + const normalized = email?.trim().toLowerCase(); + return normalized || undefined; +} + /** * Open `url` in the user's default browser. Used for the self-open sign-in * path when the client does not support URL-mode elicitation (where the client @@ -72,6 +85,14 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { private _authUrlPending = false; // mtime at last read; detects sibling rewrites of the shared store. private _credentialsMtimeMs: number | undefined; + // Consecutive abandoned sign-ins with the current client. + private _abandonedSignIns = 0; + // Held from clientInformation() returning undefined until the SDK either + // persists the registration or the connection attempt fails. + private _clientRegistrationLock: DataLockHandle | undefined; + private _accountEmail: string | undefined; + private _clientServerUrl: string | undefined; + private _tokenUpdatedAt: number | undefined; authorizationUrl: string | undefined; @@ -86,8 +107,12 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { constructor() { const stored = loadCredentials(); if (stored) { - this._tokens = stored.tokens as OAuthTokens | undefined; - this._clientInfo = stored.clientInfo as OAuthClientInformationMixed | undefined; + this._tokens = (stored.tokens ?? undefined) as OAuthTokens | undefined; + this._clientInfo = (stored.clientInfo ?? undefined) as OAuthClientInformationMixed | undefined; + this._accountEmail = stored.accountEmail ?? undefined; + this._clientServerUrl = stored.clientServerUrl ?? undefined; + this._abandonedSignIns = stored.abandonedSignIns ?? 0; + this._tokenUpdatedAt = stored.tokenUpdatedAt ?? undefined; } this._credentialsMtimeMs = credentialsMtimeMs(); } @@ -108,9 +133,35 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { if (!stored) return; if (stored.tokens) { this._tokens = stored.tokens as OAuthTokens; + } else if ( + stored.tokenUpdatedAt !== undefined && + stored.tokenUpdatedAt !== null && + (this._tokenUpdatedAt === undefined || + stored.tokenUpdatedAt > this._tokenUpdatedAt) + ) { + // A newer explicit reset/invalidation is different from a transient + // missing file or a client-only rewrite; propagate the token tombstone. + this._tokens = undefined; } if (stored.clientInfo) { this._clientInfo = stored.clientInfo as OAuthClientInformationMixed; + // A sibling completed the registration while this process was waiting. + if (this._clientRegistrationLock) { + releaseClientRegistrationLock(this._clientRegistrationLock); + this._clientRegistrationLock = undefined; + } + } + if (stored.accountEmail !== undefined) { + this._accountEmail = stored.accountEmail ?? undefined; + } + if (stored.clientServerUrl !== undefined) { + this._clientServerUrl = stored.clientServerUrl ?? undefined; + } + if (stored.abandonedSignIns !== undefined) { + this._abandonedSignIns = stored.abandonedSignIns ?? 0; + } + if (stored.tokenUpdatedAt !== undefined) { + this._tokenUpdatedAt = stored.tokenUpdatedAt ?? undefined; } } @@ -134,6 +185,7 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { return false; } this._tokens = diskTokens; + this._tokenUpdatedAt = stored?.tokenUpdatedAt ?? Date.now(); this._credentialsMtimeMs = diskMtime; if (stored?.clientInfo) { this._clientInfo = stored.clientInfo as OAuthClientInformationMixed; @@ -184,13 +236,35 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { } clientInformation(): OAuthClientInformationMixed | undefined { + // Pick up a sibling's registration; the SDK registers whenever this is undefined. + this.syncTokensFromDisk(); + if (this._clientInfo) return this._clientInfo; + + // The SDK's provider API is synchronous here. Acquire a lock before + // returning undefined and hold it until saveClientInformation() or the + // connection error path. A second process therefore waits for the first + // registration to land instead of registering a second client. + if (!this._clientRegistrationLock) { + this._clientRegistrationLock = acquireClientRegistrationLock(); + // The lock may have been held by a sibling that just completed DCR. + // Force a post-lock read before deciding that registration is needed. + this._credentialsMtimeMs = undefined; + this.syncTokensFromDisk(); + if (this._clientInfo) { + releaseClientRegistrationLock(this._clientRegistrationLock); + this._clientRegistrationLock = undefined; + } + } return this._clientInfo; } saveClientInformation(info: OAuthClientInformationMixed): void { + console.error(`[auth] Registered OAuth client: ${info.client_id}`); this._clientInfo = info; - saveCredentials(this._tokens, this._clientInfo); + saveCredentials(this._tokens, this._clientInfo, this.credentialMetadata()); this._credentialsMtimeMs = credentialsMtimeMs(); + releaseClientRegistrationLock(this._clientRegistrationLock); + this._clientRegistrationLock = undefined; } tokens(): OAuthTokens | undefined { @@ -200,33 +274,96 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { saveTokens(tokens: OAuthTokens): void { this._tokens = tokens; + this._tokenUpdatedAt = Date.now(); this._authUrlPending = false; - saveCredentials(this._tokens, this._clientInfo); + this._abandonedSignIns = 0; + saveCredentials(this._tokens, this._clientInfo, this.credentialMetadata()); // Own write must not look like a sibling change. this._credentialsMtimeMs = credentialsMtimeMs(); this.onTokensChanged?.(tokens); } + /** Persist the account/server context without changing the grant. */ + setAccountContext(accountEmail: string | undefined, serverUrl: string): void { + this._accountEmail = normalizeAccountEmail(accountEmail); + this._clientServerUrl = serverUrl; + saveCredentials(this._tokens, this._clientInfo, this.credentialMetadata()); + this._credentialsMtimeMs = credentialsMtimeMs(); + } + + /** The account associated with the currently cached grant, if known. */ + accountEmail(): string | undefined { + return this._accountEmail; + } + + /** The server for which the cached DCR client was registered, if known. */ + clientServerUrl(): string | undefined { + return this._clientServerUrl; + } + + /** + * Force a new user sign-in while retaining a client registered for the same + * server. This is used for account switching; full setup reset clears the + * provider and registered client in the setup tool. + */ + resetAuthentication(accountEmail: string | undefined, serverUrl: string): void { + const hadTokens = this._tokens !== undefined; + this._tokens = undefined; + this._tokenUpdatedAt = Date.now(); + this._accountEmail = normalizeAccountEmail(accountEmail); + this._clientServerUrl = serverUrl; + this._codeVerifier = ""; + this._pendingAuthCode = undefined; + this.authorizationUrl = undefined; + this._authUrlPending = false; + this._abandonedSignIns = 0; + saveCredentials(undefined, this._clientInfo, this.credentialMetadata(), { + forceTokenUpdate: true, + }); + this._credentialsMtimeMs = credentialsMtimeMs(); + if (hadTokens) this.onTokensChanged?.(undefined); + } + + private credentialMetadata(): CredentialMetadata { + return { + accountEmail: this._accountEmail ?? null, + clientServerUrl: this._clientServerUrl ?? null, + abandonedSignIns: this._abandonedSignIns, + tokenUpdatedAt: this._tokenUpdatedAt ?? null, + }; + } + async invalidateCredentials(scope: InvalidationScope): Promise { console.error(`[auth] Invalidating credentials: scope=${scope}`); const tokensClearedBefore = this._tokens === undefined; switch (scope) { case "all": + this.releaseClientRegistrationLock(); this._tokens = undefined; this._clientInfo = undefined; + this._accountEmail = undefined; + this._clientServerUrl = undefined; this._codeVerifier = ""; this._authUrlPending = false; + // Fresh client → fresh retry budget. + this._abandonedSignIns = 0; clearCredentials(); break; case "client": + this.releaseClientRegistrationLock(); this._clientInfo = undefined; - saveCredentials(this._tokens, undefined); + saveCredentials(this._tokens, undefined, this.credentialMetadata(), { + forceClientUpdate: true, + }); break; case "tokens": // Usually a sibling's rotation — try adopting before clearing. if (await this.adoptNewerTokenWithGrace()) return; this._tokens = undefined; - saveCredentials(undefined, this._clientInfo); + this._tokenUpdatedAt = Date.now(); + saveCredentials(undefined, this._clientInfo, this.credentialMetadata(), { + forceTokenUpdate: true, + }); break; case "verifier": this._codeVerifier = ""; @@ -240,6 +377,11 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { } } + releaseClientRegistrationLock(): void { + releaseClientRegistrationLock(this._clientRegistrationLock); + this._clientRegistrationLock = undefined; + } + // True if we previously issued an authorize URL but never received tokens — // implying the URL was likely rejected by the server (e.g. stale client_id). needsFreshClient(): boolean { @@ -250,6 +392,32 @@ export class GleanOAuthClientProvider implements OAuthClientProvider { ); } + // Abandoned sign-in: keep the client and clear the pending flow. Returns + // false after two consecutive failures (client likely dead → re-register). + abandonPendingSignIn(clientRejected = false): boolean { + const lock = acquireDataFileLockSync("mcp-auth-recovery", { + waitMs: 5_000, + staleMs: 30_000, + }); + try { + // Re-read first so a sibling's token/client metadata is not clobbered by + // this process's recovery bookkeeping. + this.syncTokensFromDisk(); + this._abandonedSignIns = clientRejected + ? 2 + : this._abandonedSignIns + 1; + this._codeVerifier = ""; + this._pendingAuthCode = undefined; + this.authorizationUrl = undefined; + this._authUrlPending = false; + saveCredentials(this._tokens, this._clientInfo, this.credentialMetadata()); + this._credentialsMtimeMs = credentialsMtimeMs(); + return !clientRejected && this._abandonedSignIns < 2; + } finally { + releaseDataFileLock(lock); + } + } + get pendingAuthCode(): string | undefined { return this._pendingAuthCode; } diff --git a/shared/glean/mcp/src/data-dir.ts b/shared/glean/mcp/src/data-dir.ts new file mode 100644 index 0000000..b872934 --- /dev/null +++ b/shared/glean/mcp/src/data-dir.ts @@ -0,0 +1,249 @@ +import fs from "node:fs"; +import path from "node:path"; +import { randomUUID } from "node:crypto"; +import { homedir } from "node:os"; + +const DIR_MODE = 0o700; +const FILE_MODE = 0o600; +const MIGRATION_MARKER = ".legacy-store-migrated-v1"; +const MIGRATION_LOCK = ".legacy-store-migration.lock"; + +// These are the files that historically followed PLUGIN_DATA_DIR. Auth and +// setup state now live in the stable per-user directory so terminal, VS Code, +// Cursor, and managed plugin launches share one state store. +const MIGRATED_FILES = [ + "mcp-credentials.json", + "mcp-server-url.json", + "remote-tools-cache.json", + "glean-server.log", +]; + +export interface DataLockOptions { + waitMs?: number; + staleMs?: number; +} + +export interface DataLockHandle { + path: string; + token: string; +} + +function envValue(name: string): string | undefined { + const value = process.env[name]?.trim(); + if (!value || value.startsWith("${")) return undefined; + return value; +} + +/** + * Stable auth/config directory shared by every host and launcher. + * + * GLEAN_AUTH_DATA_DIR is intentionally an explicit escape hatch for tests and + * controlled deployments. PLUGIN_DATA_DIR is deliberately not used here: it + * is host-managed and was the source of the terminal/plugin store split. + */ +export function resolveDataDir(): string { + return ( + envValue("GLEAN_AUTH_DATA_DIR") ?? + path.join(homedir() || process.env.TMPDIR || "/tmp", ".glean") + ); +} + +function legacyDataDirs(): string[] { + const canonical = path.resolve(resolveDataDir()); + const candidates = [envValue("PLUGIN_DATA_DIR"), envValue("CLAUDE_PLUGIN_DATA")]; + return [...new Set(candidates.filter((dir): dir is string => !!dir))] + .map((dir) => path.resolve(dir)) + .filter((dir) => dir !== canonical); +} + +function sleepSync(ms: number): void { + const signal = new Int32Array(new SharedArrayBuffer(4)); + Atomics.wait(signal, 0, 0, ms); +} + +function lockPath(name: string): string { + return path.join(resolveDataDir(), `.${name}.lock`); +} + +function isStale(lockFile: string, staleMs: number): boolean { + try { + return Date.now() - fs.statSync(lockFile).mtimeMs > staleMs; + } catch { + return false; + } +} + +/** + * Acquire a small cross-process lock using O_EXCL semantics. The returned + * path must be released by the caller. Stale locks are recoverable after the + * configured timeout so a crashed plugin cannot permanently block auth. + */ +export function acquireDataFileLockSync( + name: string, + options: DataLockOptions = {}, +): DataLockHandle | undefined { + const waitMs = options.waitMs ?? 30_000; + const staleMs = options.staleMs ?? 120_000; + const filePath = lockPath(name); + const dir = path.dirname(filePath); + + try { + fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); + fs.chmodSync(dir, DIR_MODE); + } catch { + return undefined; + } + + const deadline = Date.now() + Math.max(0, waitMs); + do { + try { + const token = randomUUID(); + const fd = fs.openSync(filePath, "wx", FILE_MODE); + try { + fs.writeFileSync( + fd, + JSON.stringify({ + pid: process.pid, + token, + createdAt: new Date().toISOString(), + }), + { encoding: "utf-8" }, + ); + } finally { + fs.closeSync(fd); + } + return { path: filePath, token }; + } catch (err) { + const code = err && typeof err === "object" && "code" in err + ? (err as { code?: string }).code + : undefined; + if (code !== "EEXIST") return undefined; + if (isStale(filePath, staleMs)) { + try { + fs.rmSync(filePath, { force: true }); + } catch { + /* another process may have replaced or removed the lock */ + } + continue; + } + if (Date.now() >= deadline) return undefined; + sleepSync(Math.min(25, Math.max(1, deadline - Date.now()))); + } + } while (Date.now() <= deadline); + + return undefined; +} + +export function releaseDataFileLock(lock: DataLockHandle | undefined): void { + if (!lock) return; + try { + const contents = JSON.parse(fs.readFileSync(lock.path, "utf-8")) as { + token?: string; + }; + if (contents.token !== lock.token) return; + fs.rmSync(lock.path, { force: true }); + } catch { + /* best effort; stale-lock recovery handles a crashed owner */ + } +} + +function copyFileAtomically(source: string, target: string): void { + const dir = path.dirname(target); + fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); + fs.chmodSync(dir, DIR_MODE); + const tmp = `${target}.${process.pid}.migration.tmp`; + try { + fs.copyFileSync(source, tmp); + fs.chmodSync(tmp, FILE_MODE); + fs.renameSync(tmp, target); + fs.chmodSync(target, FILE_MODE); + } finally { + try { + fs.rmSync(tmp, { force: true }); + } catch { + /* ignore cleanup failures */ + } + } +} + +/** + * Migrate files from the old host-managed store(s) once. If both stores exist, + * the newest source file wins; after migration all current code writes only to + * resolveDataDir(). The marker prevents an old, already-running process from + * repeatedly overwriting the canonical store on every request. + */ +export function migrateLegacyData(): void { + const canonical = resolveDataDir(); + const sources = legacyDataDirs(); + if (sources.length === 0) return; + + const marker = path.join(canonical, MIGRATION_MARKER); + try { + if (fs.existsSync(marker)) return; + } catch { + return; + } + + const lock = acquireDataFileLockSync( + MIGRATION_LOCK.slice(1), + { waitMs: 5_000, staleMs: 60_000 }, + ); + if (!lock) return; + + try { + if (fs.existsSync(marker)) return; + fs.mkdirSync(canonical, { recursive: true, mode: DIR_MODE }); + fs.chmodSync(canonical, DIR_MODE); + + for (const filename of MIGRATED_FILES) { + const candidates = sources + .map((dir) => path.join(dir, filename)) + .filter((file) => { + try { + return fs.statSync(file).isFile(); + } catch { + return false; + } + }) + .sort((a, b) => fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs); + if (candidates.length === 0) continue; + + const target = path.join(canonical, filename); + const newestSource = candidates[0]; + let targetIsOlder = true; + try { + targetIsOlder = fs.statSync(target).mtimeMs < fs.statSync(newestSource).mtimeMs; + } catch { + /* target does not exist */ + } + if (!fs.existsSync(target) || targetIsOlder) { + copyFileAtomically(newestSource, target); + console.error(`[auth] Migrated legacy data: ${filename}`); + } + } + + fs.writeFileSync( + marker, + JSON.stringify({ migratedAt: new Date().toISOString(), sources }), + { encoding: "utf-8", mode: FILE_MODE }, + ); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + console.error(`[auth] Failed to migrate legacy data: ${msg}`); + } finally { + releaseDataFileLock(lock); + } +} + +export function ensureDataDir(): string { + const dir = resolveDataDir(); + fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); + fs.chmodSync(dir, DIR_MODE); + return dir; +} + +export const dataDirConstants = { + DIR_MODE, + FILE_MODE, + MIGRATION_MARKER, +}; diff --git a/shared/glean/mcp/src/index.ts b/shared/glean/mcp/src/index.ts index e524b7f..4229f45 100644 --- a/shared/glean/mcp/src/index.ts +++ b/shared/glean/mcp/src/index.ts @@ -8,13 +8,17 @@ import { } from "@modelcontextprotocol/sdk/types.js"; import path from "node:path"; import fs from "node:fs"; -import { homedir, tmpdir } from "node:os"; +import { tmpdir } from "node:os"; import { AuthRequiredError, createRemoteClient, type RemoteClientOptions, } from "./remote-client.js"; -import { GleanOAuthClientProvider, openBrowser } from "./auth-provider.js"; +import { + GleanOAuthClientProvider, + normalizeAccountEmail, + openBrowser, +} from "./auth-provider.js"; import { startCallbackServer, closeCallbackServer, @@ -31,7 +35,7 @@ import { saveServerUrl, clearServerUrl, } from "./url-config-store.js"; -import { clearCredentials } from "./token-store.js"; +import { clearCredentials, loadCredentials } from "./token-store.js"; import { loadRemoteTools, saveRemoteTools, @@ -45,6 +49,8 @@ import { } from "./tools/remote-passthrough.js"; import { resolveSessionId } from "./session-id.js"; import { resolveServerUrlFromEmail } from "./config-search.js"; +import { resolveDataDir } from "./data-dir.js"; +import { normalizeServerUrl } from "./server-url.js"; import { PLUGIN_VERSION } from "./version.js"; function readEnv(...keys: string[]): string | undefined { @@ -63,11 +69,6 @@ function resolveServerUrl(): string | undefined { return loadServerUrl(); } -function normalizeServerUrl(raw: string): string { - const parsed = new URL(raw); - return `${parsed.origin}/mcp/gateway/proxy`; -} - const SETUP_REQUIRED_TEXT = `[SETUP_REQUIRED]\n\n` + `To connect, enter your work email (e.g. you@acme.com) and we'll find ` + @@ -93,8 +94,7 @@ const AUTH_REDIRECT_TO_SETUP_TEXT = "(no arguments) to sign in to Glean, then retry this tool."; function resolveLogPath(): string { - const base = process.env.PLUGIN_DATA_DIR || path.join(homedir(), ".glean"); - return path.join(base, "glean-server.log"); + return path.join(resolveDataDir(), "glean-server.log"); } const LOG_PATH = resolveLogPath(); @@ -522,6 +522,11 @@ async function advanceSetup(): Promise { return { content: [{ type: "text", text: SETUP_REQUIRED_TEXT }] }; } + const provider = getOAuthProvider(); + if (provider.clientServerUrl() !== serverUrl) { + provider.setAccountContext(provider.accountEmail(), serverUrl); + } + const conn = await connectWithSignIn(serverUrl); if (!conn.ok) return conn.result; const remoteClient = conn.client; @@ -796,6 +801,20 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => { }; } + const previousUrl = loadServerUrl(); + const stored = loadCredentials(); + const requestedEmail = normalizeAccountEmail(email); + const urlChanged = previousUrl !== normalized; + const reusableClient = + !!stored?.clientInfo && stored.clientServerUrl === normalized; + // A new email on the same server is an account switch, not a new + // instance: clear only the grant and retain the server-scoped DCR + // client. This lets users change accounts without adding a client. + const accountChanged = + !!requestedEmail && + !!stored?.tokens && + stored.accountEmail !== requestedEmail; + try { saveServerUrl(normalized); } catch (err) { @@ -808,14 +827,30 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => { }; } - // New instance — clear stale auth state. The on-disk remote-tool - // cache for the previous URL is left intact (so switching back is - // instant); we just rehydrate from whatever cache exists for the - // newly configured URL — empty for a first-time server. - clearCredentials(); - oauthProvider = undefined; + // A DCR client is scoped to its server. Reuse a retained client after + // setup changes only when its recorded server matches; otherwise + // remove it so the SDK registers for the new instance. + if (urlChanged && !reusableClient) { + clearCredentials(); + oauthProvider = undefined; + } + + const provider = getOAuthProvider(); + if (accountChanged) { + provider.resetAuthentication(requestedEmail, normalized); + } else { + provider.setAccountContext( + requestedEmail ?? provider.accountEmail(), + normalized, + ); + } cachedRemoteTools = loadRemoteTools(normalized); - logLine("setup.configured", { serverUrl: normalized }); + logLine("setup.configured", { + serverUrl: normalized, + urlChanged, + reusableClient, + accountChanged, + }); // Fall through to advanceSetup, which will now find URL ✓ and try // to drive auth + tool fetch in the same call. } diff --git a/shared/glean/mcp/src/remote-client.ts b/shared/glean/mcp/src/remote-client.ts index 568acab..b58f4db 100644 --- a/shared/glean/mcp/src/remote-client.ts +++ b/shared/glean/mcp/src/remote-client.ts @@ -112,6 +112,14 @@ export class AuthRequiredError extends Error { } let pendingTransport: StreamableHTTPClientTransport | undefined; +// The provider API is synchronous, so two connect calls in one MCP process can +// both observe an empty client before the SDK's async DCR request completes. +// Serialize that first registration per provider in addition to the OS lock +// used by token-store for sibling processes. +const pendingRegistrations = new WeakMap< + GleanOAuthClientProvider, + Promise +>(); // Serialize the SDK operations that can drive an OAuth token refresh // (client.connect and finishAuth). getOAuthProvider() is a process-wide @@ -194,20 +202,42 @@ export async function createRemoteClient( } catch (err) { const msg = err instanceof Error ? err.message : String(err); console.error(`[auth] Code exchange failed: ${msg} — discarding stale auth state`); - authProvider.clearPendingAuth(); pendingTransport = undefined; - await authProvider.invalidateCredentials("all"); + // Keep the registration for code/verifier failures. An explicit + // invalid_client response escalates immediately; generic repeated + // failures use the shared two-attempt recovery budget. + if (!authProvider.abandonPendingSignIn(isInvalidClientError(err))) { + await authProvider.invalidateCredentials("all"); + } return createRemoteClient(serverUrl, opts, chatSessionId); } } - // DCR recovery: we previously issued an authorize URL but never received - // tokens. The URL was likely rejected by the server (most commonly: the - // cached DCR client was deleted server-side). Force a fresh DCR so the next - // URL we generate uses a valid, server-known client_id. + // Unfinished sign-in: reuse the existing registration first; fresh DCR is + // the escalation path. if (authProvider?.needsFreshClient()) { - console.error("[auth] Previous auth URL didn't complete — forcing fresh DCR"); - await authProvider.invalidateCredentials("all"); + if (authProvider.abandonPendingSignIn()) { + console.error( + "[auth] Previous sign-in didn't complete — retrying with the existing client", + ); + } else { + console.error( + "[auth] Sign-in failed twice with this client — forcing fresh DCR", + ); + await authProvider.invalidateCredentials("all"); + } + } + + let registrationNeeded = false; + if (authProvider?.clientInformation) { + registrationNeeded = !authProvider.clientInformation(); + if (registrationNeeded) { + const pending = pendingRegistrations.get(authProvider); + if (pending) { + await pending; + registrationNeeded = !authProvider.clientInformation(); + } + } } const client = new Client( @@ -219,10 +249,25 @@ export async function createRemoteClient( const accessTokenAtConnect = authProvider?.tokens()?.access_token; const transport = buildTransport(serverUrl, opts, chatSessionId); + let trackedRegistration: Promise | undefined; + let connectPromise: Promise; + if (registrationNeeded && authProvider) { + connectPromise = withConnectLock(() => client.connect(transport)); + trackedRegistration = connectPromise.then( + () => undefined, + () => undefined, + ); + pendingRegistrations.set(authProvider, trackedRegistration); + } else { + connectPromise = withConnectLock(() => client.connect(transport)); + } try { - await withConnectLock(() => client.connect(transport)); + await connectPromise; } catch (error) { + // If registration failed before saveClientInformation() ran, do not leave + // the cross-process registration lock behind until stale-lock expiry. + authProvider?.releaseClientRegistrationLock?.(); if (error instanceof UnauthorizedError && authProvider) { const refreshedAccessToken = authProvider.tokens()?.access_token; if ( @@ -255,11 +300,24 @@ export async function createRemoteClient( return createRemoteClient(serverUrl, opts, chatSessionId, true); } throw error; + } finally { + if ( + trackedRegistration && + authProvider && + pendingRegistrations.get(authProvider) === trackedRegistration + ) { + pendingRegistrations.delete(authProvider); + } } return client; } +function isInvalidClientError(error: unknown): boolean { + const msg = error instanceof Error ? error.message : String(error); + return /invalid[_ -]?client|unknown client|client authentication failed/i.test(msg); +} + // Match broadly; the caller's disk re-check gates the actual retry. function isLikelyRefreshFailure(error: unknown): boolean { const msg = error instanceof Error ? error.message : String(error); diff --git a/shared/glean/mcp/src/remote-tools-cache-store.ts b/shared/glean/mcp/src/remote-tools-cache-store.ts index fa2900b..e3ad5cc 100644 --- a/shared/glean/mcp/src/remote-tools-cache-store.ts +++ b/shared/glean/mcp/src/remote-tools-cache-store.ts @@ -1,14 +1,18 @@ import fs from "node:fs"; import path from "node:path"; -import { homedir } from "node:os"; import type { Tool } from "@modelcontextprotocol/sdk/types.js"; +import { + ensureDataDir, + migrateLegacyData, + resolveDataDir, +} from "./data-dir.js"; const CACHE_FILENAME = "remote-tools-cache.json"; const DIR_MODE = 0o700; const FILE_MODE = 0o600; function resolveCacheDir(): string { - return process.env.PLUGIN_DATA_DIR || path.join(homedir(), ".glean"); + return resolveDataDir(); } function cacheFile(): string { @@ -23,6 +27,7 @@ interface CacheEntry { type Store = Record; function readStore(): Store { + migrateLegacyData(); try { const raw = fs.readFileSync(cacheFile(), "utf-8"); const data = JSON.parse(raw); @@ -37,8 +42,7 @@ function readStore(): Store { function writeStore(store: Store): void { const filePath = cacheFile(); - const dir = path.dirname(filePath); - fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); + const dir = ensureDataDir(); fs.chmodSync(dir, DIR_MODE); fs.writeFileSync(filePath, JSON.stringify(store, null, 2), { encoding: "utf-8", @@ -68,6 +72,7 @@ export function saveRemoteTools(serverUrl: string, tools: Tool[]): void { } export function clearRemoteTools(serverUrl?: string): void { + migrateLegacyData(); try { if (!serverUrl) { fs.rmSync(cacheFile(), { force: true }); diff --git a/shared/glean/mcp/src/server-url.ts b/shared/glean/mcp/src/server-url.ts new file mode 100644 index 0000000..82ba0e0 --- /dev/null +++ b/shared/glean/mcp/src/server-url.ts @@ -0,0 +1,16 @@ +const MCP_GATEWAY_PATH = "/mcp/gateway/proxy"; + +/** + * Normalize either a normal QE origin or a path-prefixed experimental QE URL + * to the MCP gateway endpoint. + */ +export function normalizeServerUrl(raw: string): string { + const parsed = new URL(raw); + const pathname = parsed.pathname.replace(/\/+$/, ""); + const prefix = pathname.endsWith(MCP_GATEWAY_PATH) + ? pathname.slice(0, -MCP_GATEWAY_PATH.length) + : pathname === "/" + ? "" + : pathname; + return `${parsed.origin}${prefix}${MCP_GATEWAY_PATH}${parsed.search}`; +} diff --git a/shared/glean/mcp/src/token-store.ts b/shared/glean/mcp/src/token-store.ts index 535a07d..476bc77 100644 --- a/shared/glean/mcp/src/token-store.ts +++ b/shared/glean/mcp/src/token-store.ts @@ -1,38 +1,64 @@ import fs from "node:fs"; import path from "node:path"; -import { homedir } from "node:os"; +import { + acquireDataFileLockSync, + ensureDataDir, + type DataLockHandle, + migrateLegacyData, + releaseDataFileLock, + resolveDataDir, +} from "./data-dir.js"; const CREDENTIALS_FILENAME = "mcp-credentials.json"; const DIR_MODE = 0o700; const FILE_MODE = 0o600; +const CLIENT_REGISTRATION_LOCK = "mcp-client-registration"; -function resolveCredentialsDir(): string { - return process.env.PLUGIN_DATA_DIR || path.join(homedir(), ".glean"); +export interface CredentialMetadata { + /** Work account used for the last successful/initiated sign-in. */ + accountEmail?: string | null; + /** OAuth clients are server-specific. */ + clientServerUrl?: string | null; + /** Shared recovery budget across sibling plugin processes. */ + abandonedSignIns?: number | null; + /** Local monotonic-ish timestamp for conflict-free token merging. */ + tokenUpdatedAt?: number | null; } -function credentialsFile(): string { - return path.join(resolveCredentialsDir(), CREDENTIALS_FILENAME); -} - -interface StoredCredentials { +export interface StoredCredentials extends CredentialMetadata { tokens?: unknown; clientInfo?: unknown; } -export function loadCredentials(): StoredCredentials | undefined { +export interface SaveCredentialsOptions { + forceTokenUpdate?: boolean; + forceClientUpdate?: boolean; +} + +function credentialsFile(): string { + return path.join(resolveDataDir(), CREDENTIALS_FILENAME); +} + +function readCredentialsFile(filePath: string): StoredCredentials | undefined { try { - const raw = fs.readFileSync(credentialsFile(), "utf-8"); + const raw = fs.readFileSync(filePath, "utf-8"); return JSON.parse(raw) as StoredCredentials; } catch { return undefined; } } +export function loadCredentials(): StoredCredentials | undefined { + migrateLegacyData(); + return readCredentialsFile(credentialsFile()); +} + /** * mtime of the credentials file (epoch ms), or undefined if unreadable. * Cheap change probe: a single stat, no read + parse. */ export function credentialsMtimeMs(): number | undefined { + migrateLegacyData(); try { return fs.statSync(credentialsFile()).mtimeMs; } catch { @@ -40,13 +66,71 @@ export function credentialsMtimeMs(): number | undefined { } } -export function saveCredentials(tokens: unknown, clientInfo: unknown): void { +/** + * Acquire the registration lock before returning undefined from + * clientInformation(). The MCP SDK interprets undefined as permission to call + * /oauth/register, so the lock must span that asynchronous SDK operation and + * is released by saveClientInformation() or the connection error path. + */ +export function acquireClientRegistrationLock(): DataLockHandle | undefined { + return acquireDataFileLockSync(CLIENT_REGISTRATION_LOCK, { + waitMs: 30_000, + staleMs: 120_000, + }); +} + +export function releaseClientRegistrationLock(lock: DataLockHandle | undefined): void { + releaseDataFileLock(lock); +} + +export function saveCredentials( + tokens: unknown, + clientInfo: unknown, + metadata: CredentialMetadata = {}, + options: SaveCredentialsOptions = {}, +): void { try { + migrateLegacyData(); const filePath = credentialsFile(); - const dir = path.dirname(filePath); - fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); - fs.chmodSync(dir, DIR_MODE); - const data: StoredCredentials = { tokens, clientInfo }; + const dir = ensureDataDir(); + const existing = readCredentialsFile(filePath); + const incomingTokenTime = metadata.tokenUpdatedAt ?? undefined; + const existingTokenTime = existing?.tokenUpdatedAt ?? undefined; + const existingTokenIsNewer = + !options.forceTokenUpdate && + existing?.tokens !== undefined && + (tokens === undefined || + (existingTokenTime !== undefined && + (incomingTokenTime === undefined || + existingTokenTime > incomingTokenTime))); + const effectiveTokens = existingTokenIsNewer ? existing?.tokens : tokens; + const effectiveTokenTime = existingTokenIsNewer + ? existingTokenTime + : metadata.tokenUpdatedAt !== undefined + ? metadata.tokenUpdatedAt + : existing?.tokenUpdatedAt; + const effectiveClientInfo = + !options.forceClientUpdate && clientInfo === undefined + ? existing?.clientInfo + : clientInfo; + const data: StoredCredentials = { + tokens: effectiveTokens, + clientInfo: effectiveClientInfo, + // Omitted metadata means preserve it; null explicitly clears it. + accountEmail: + metadata.accountEmail !== undefined + ? metadata.accountEmail + : existing?.accountEmail, + clientServerUrl: + metadata.clientServerUrl !== undefined + ? metadata.clientServerUrl + : existing?.clientServerUrl, + abandonedSignIns: + metadata.abandonedSignIns !== undefined + ? metadata.abandonedSignIns + : existing?.abandonedSignIns, + tokenUpdatedAt: effectiveTokenTime, + }; // Temp-file + rename: concurrent readers never see a half-written store. const tmpPath = `${filePath}.${process.pid}.tmp`; fs.writeFileSync(tmpPath, JSON.stringify(data, null, 2), { @@ -54,6 +138,8 @@ export function saveCredentials(tokens: unknown, clientInfo: unknown): void { mode: FILE_MODE, }); fs.renameSync(tmpPath, filePath); + fs.chmodSync(dir, DIR_MODE); + fs.chmodSync(filePath, FILE_MODE); } catch (err) { const msg = err instanceof Error ? err.message : String(err); console.error(`[auth] Failed to persist credentials: ${msg}`); @@ -62,6 +148,7 @@ export function saveCredentials(tokens: unknown, clientInfo: unknown): void { export function clearCredentials(): void { try { + migrateLegacyData(); fs.rmSync(credentialsFile(), { force: true }); } catch (err) { const msg = err instanceof Error ? err.message : String(err); diff --git a/shared/glean/mcp/src/url-config-store.ts b/shared/glean/mcp/src/url-config-store.ts index df46d7e..6a7d9e6 100644 --- a/shared/glean/mcp/src/url-config-store.ts +++ b/shared/glean/mcp/src/url-config-store.ts @@ -1,13 +1,17 @@ import fs from "node:fs"; import path from "node:path"; -import { homedir } from "node:os"; +import { + ensureDataDir, + migrateLegacyData, + resolveDataDir, +} from "./data-dir.js"; const CONFIG_FILENAME = "mcp-server-url.json"; const DIR_MODE = 0o700; const FILE_MODE = 0o600; function resolveConfigDir(): string { - return process.env.PLUGIN_DATA_DIR || path.join(homedir(), ".glean"); + return resolveDataDir(); } function configFile(): string { @@ -19,6 +23,7 @@ interface StoredConfig { } export function loadServerUrl(): string | undefined { + migrateLegacyData(); try { const raw = fs.readFileSync(configFile(), "utf-8"); const data = JSON.parse(raw) as StoredConfig; @@ -30,9 +35,9 @@ export function loadServerUrl(): string | undefined { } export function saveServerUrl(url: string): void { + migrateLegacyData(); const filePath = configFile(); - const dir = path.dirname(filePath); - fs.mkdirSync(dir, { recursive: true, mode: DIR_MODE }); + const dir = ensureDataDir(); fs.chmodSync(dir, DIR_MODE); const data: StoredConfig = { serverUrl: url }; fs.writeFileSync(filePath, JSON.stringify(data, null, 2), { @@ -43,6 +48,7 @@ export function saveServerUrl(url: string): void { } export function clearServerUrl(): void { + migrateLegacyData(); try { fs.rmSync(configFile(), { force: true }); } catch { diff --git a/shared/glean/mcp/tests/auth-provider.test.ts b/shared/glean/mcp/tests/auth-provider.test.ts index 3e4e032..3e2365e 100644 --- a/shared/glean/mcp/tests/auth-provider.test.ts +++ b/shared/glean/mcp/tests/auth-provider.test.ts @@ -28,6 +28,8 @@ describe("GleanOAuthClientProvider", () => { beforeEach(() => { delete process.env.PLUGIN_DATA_DIR; + delete process.env.CLAUDE_PLUGIN_DATA; + delete process.env.GLEAN_AUTH_DATA_DIR; // Skip the rotation grace window by default so invalidation tests don't // wait out the real 2s poll; the grace test overrides this explicitly. process.env.GLEAN_ROTATION_GRACE_MS = "0"; @@ -215,6 +217,82 @@ describe("GleanOAuthClientProvider", () => { expect(provider.tokens()).toBeUndefined(); }); + // --- Client reuse: an abandoned sign-in must not burn the DCR client. + // Every registration permanently adds a client server-side, so the + // existing one is retried first and a fresh DCR is the escalation path. --- + + it("clientInformation() adopts a sibling's registration from disk", () => { + // Constructed with no credentials — this process would otherwise register. + const provider = new GleanOAuthClientProvider(); + expect(provider.clientInformation()).toBeUndefined(); + + // Sibling process wins the registration race and persists its client. + writeCredFileNewer(undefined, { client_id: "cid_sibling" }); + + expect(provider.clientInformation()).toEqual({ client_id: "cid_sibling" }); + }); + + it("abandonPendingSignIn keeps the registered client and clears the pending flow", async () => { + const provider = new GleanOAuthClientProvider(); + provider.saveClientInformation({ client_id: "cid" } as any); + provider.saveCodeVerifier("v1"); + await provider.redirectToAuthorization( + new URL("https://example.com/oauth/authorize?state=s1"), + ); + expect(provider.needsFreshClient()).toBe(true); + + expect(provider.abandonPendingSignIn()).toBe(true); + + expect(provider.clientInformation()).toEqual({ client_id: "cid" }); + expect(provider.codeVerifier()).toBe(""); + expect(provider.authorizationUrl).toBeUndefined(); + expect(provider.needsFreshClient()).toBe(false); + // Registration untouched on disk — nothing was wiped. + const raw = JSON.parse(fs.readFileSync(credFile, "utf-8")); + expect(raw.clientInfo.client_id).toBe("cid"); + }); + + it("abandonPendingSignIn exhausts after two consecutive abandonments", () => { + const provider = new GleanOAuthClientProvider(); + expect(provider.abandonPendingSignIn()).toBe(true); + expect(provider.abandonPendingSignIn()).toBe(false); + }); + + it("a completed sign-in resets the abandonment budget", () => { + const provider = new GleanOAuthClientProvider(); + expect(provider.abandonPendingSignIn()).toBe(true); + provider.saveTokens({ access_token: "tok" } as any); + // Fresh budget: the next abandonment retries the client again. + expect(provider.abandonPendingSignIn()).toBe(true); + }); + + it("invalidateCredentials('all') resets the budget for the next registration", async () => { + const provider = new GleanOAuthClientProvider(); + expect(provider.abandonPendingSignIn()).toBe(true); + expect(provider.abandonPendingSignIn()).toBe(false); + await provider.invalidateCredentials("all"); + // A freshly registered client gets a fresh retry budget. + expect(provider.abandonPendingSignIn()).toBe(true); + }); + + it("resetAuthentication keeps the client for account switching", () => { + const provider = new GleanOAuthClientProvider(); + provider.saveTokens({ access_token: "tok", refresh_token: "refresh" } as any); + provider.saveClientInformation({ client_id: "cid" } as any); + + provider.resetAuthentication( + "new-account@example.com", + "https://example.com/mcp/gateway/proxy", + ); + + expect(provider.tokens()).toBeUndefined(); + expect(provider.clientInformation()).toEqual({ client_id: "cid" }); + expect(provider.accountEmail()).toBe("new-account@example.com"); + const raw = JSON.parse(fs.readFileSync(credFile, "utf-8")); + expect(raw.tokens).toBeUndefined(); + expect(raw.clientInfo.client_id).toBe("cid"); + }); + it("saveClientInformation persists to disk", () => { const provider = new GleanOAuthClientProvider(); const info = { client_id: "cid", client_secret: "sec" } as any; diff --git a/shared/glean/mcp/tests/remote-client-auth-retry.test.ts b/shared/glean/mcp/tests/remote-client-auth-retry.test.ts index f730cc7..2508234 100644 --- a/shared/glean/mcp/tests/remote-client-auth-retry.test.ts +++ b/shared/glean/mcp/tests/remote-client-auth-retry.test.ts @@ -90,6 +90,51 @@ describe("createRemoteClient sibling-refresh retry", () => { }); }); +describe("createRemoteClient abandoned sign-in client reuse", () => { + beforeEach(() => { + connectMock.mockReset(); + }); + + function makeAbandonedProvider(budgetLeft: boolean) { + return { + tokens: () => ({ access_token: "T0" }), + authorizationUrl: undefined, + pendingAuthCode: undefined, + needsFreshClient: () => true, + abandonPendingSignIn: vi.fn(() => budgetLeft), + invalidateCredentials: vi.fn(), + } as any; + } + + it("reuses the existing client on the first abandoned sign-in", async () => { + connectMock.mockResolvedValueOnce(undefined); + const provider = makeAbandonedProvider(true); + + await createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-3", + ); + + expect(provider.abandonPendingSignIn).toHaveBeenCalledTimes(1); + // The registration survives — no wipe, no fresh DCR. + expect(provider.invalidateCredentials).not.toHaveBeenCalled(); + }); + + it("falls back to a fresh DCR once the retry budget is exhausted", async () => { + connectMock.mockResolvedValueOnce(undefined); + const provider = makeAbandonedProvider(false); + + await createRemoteClient( + "https://acme-be.glean.com/mcp/gateway/proxy", + { authProvider: provider }, + "sess-4", + ); + + expect(provider.invalidateCredentials).toHaveBeenCalledWith("all"); + }); +}); + describe("createRemoteClient refresh-collision retry", () => { beforeEach(() => { connectMock.mockReset(); diff --git a/shared/glean/mcp/tests/remote-tools-cache-store.test.ts b/shared/glean/mcp/tests/remote-tools-cache-store.test.ts index 0c0987f..4638400 100644 --- a/shared/glean/mcp/tests/remote-tools-cache-store.test.ts +++ b/shared/glean/mcp/tests/remote-tools-cache-store.test.ts @@ -8,7 +8,7 @@ const tmpDir = fs.mkdtempSync( path.join(os.tmpdir(), "remote-tools-cache-store-test-"), ); -vi.stubEnv("PLUGIN_DATA_DIR", tmpDir); +vi.stubEnv("GLEAN_AUTH_DATA_DIR", tmpDir); const { loadRemoteTools, saveRemoteTools, clearRemoteTools } = await import( "../src/remote-tools-cache-store.js" diff --git a/shared/glean/mcp/tests/url-config-store.test.ts b/shared/glean/mcp/tests/url-config-store.test.ts index b7c9a5f..a4856af 100644 --- a/shared/glean/mcp/tests/url-config-store.test.ts +++ b/shared/glean/mcp/tests/url-config-store.test.ts @@ -5,7 +5,7 @@ import os from "node:os"; const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "url-config-store-test-")); -vi.stubEnv("PLUGIN_DATA_DIR", tmpDir); +vi.stubEnv("GLEAN_AUTH_DATA_DIR", tmpDir); const { loadServerUrl, saveServerUrl, clearServerUrl } = await import( "../src/url-config-store.js"