pi-mcp-extension
Version:
Full-featured, lightweight MCP extension for the Pi coding agent: stdio, SSE & Streamable HTTP transports, OAuth 2.1 (PKCE + DCR), lazy/eager lifecycles — minimal dependencies, zero bloat.
545 lines (489 loc) • 21.1 kB
text/typescript
/**
* File-based OAuthClientProvider for MCP servers.
*
* Implements the full OAuth2 Authorization Code flow with PKCE and
* Dynamic Client Registration (RFC 7591) as required by the MCP spec.
*
* Token and client state are persisted per-server under ~/.pi/agent/mcp-auth/
* so they survive pi restarts without requiring re-authorization.
*
* Usage: config adds `auth: { ... }` to a server config. This module
* constructs an OAuthClientProvider and the transport factory wires it in.
*/
import type { OAuthClientProvider, OAuthDiscoveryState } from "@modelcontextprotocol/sdk/client/auth.js";
import type { OAuthClientMetadata, OAuthClientInformationMixed, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
import { AuthLockError, AuthRequiredError, McpError } from "./errors.js";
import {
acquireAuthLock,
type AcquireAuthLockOptions,
type AuthLockHandle,
} from "./auth-lock.js";
import {
applyStoredStateInvalidation,
assertIssuerBindingCompatible,
buildInteractiveStoredClientInfo,
buildNewSilentStoredClientInfo,
buildOAuthClientMetadata,
buildStaticClientInformation,
DEFAULT_CALLBACK_PORT,
DEFAULT_REDIRECT_URL,
hasRedirectUriMismatch,
isClientSecretExpired,
isPeerRotatedTokens,
mergeSilentStoredClientInfo,
resolveProviderRedirectUrl,
stampStateServerUrl,
toOAuthClientInformation,
toOAuthTokens,
toStoredTokens,
type AuthConfig,
type InvalidationScope,
} from "./oauth-policy.js";
import {
loadState,
loadStoredRedirectUrisSync,
saveState,
type StoredClientInfo,
type StoredState,
type StoredTokens,
} from "./oauth-state-store.js";
import { describeError } from "./debug-log-format.js";
import { reportDiagnostic } from "./debug-log.js";
export { acquireAuthLock, AuthLockError, AuthRequiredError, type AcquireAuthLockOptions, type AuthLockHandle };
export { DEFAULT_CALLBACK_PORT, DEFAULT_REDIRECT_URL, type AuthConfig, type InvalidationScope };
export interface McpOAuthProviderOptions {
serverName: string;
authConfig: AuthConfig;
serverUrl?: string | undefined;
storageDir?: string | undefined;
mode?: "silent" | "interactive" | undefined;
redirectUrl?: string | URL | undefined;
onAuthRequired?: ((url: URL) => void | Promise<void>) | undefined;
isCancelled?: (() => boolean) | undefined;
assertLockHeld?: (() => Promise<void>) | undefined;
}
// ─── OAuthClientProvider Implementation ────────────────────────────────────────
export class McpOAuthProvider implements OAuthClientProvider {
private readonly serverName: string;
private readonly authConfig: AuthConfig;
private readonly serverUrl: string | undefined;
private readonly storageDir: string | undefined;
private readonly mode: "silent" | "interactive";
private readonly explicitRedirectUrl: string | URL | undefined;
private readonly _onAuthRequired: ((url: URL) => void | Promise<void>) | undefined;
private readonly isCancelled: () => boolean;
private readonly assertLockHeld: (() => Promise<void>) | undefined;
private _oauthState: string | undefined;
private _codeVerifier: string | undefined;
private _cachedTokens: StoredTokens | undefined;
private _lastReadTokens: StoredTokens | undefined;
private _reloadedPeerTokens = false;
private _storedRedirectUris: string[] | null | undefined;
constructor(options: McpOAuthProviderOptions) {
this.serverName = options.serverName;
this.authConfig = options.authConfig;
this.serverUrl = options.serverUrl;
this.storageDir = options.storageDir;
this.mode = options.mode ?? "interactive";
this.explicitRedirectUrl = options.redirectUrl;
this._onAuthRequired = options.onAuthRequired;
this.isCancelled = options.isCancelled ?? (() => false);
this.assertLockHeld = options.assertLockHeld;
}
private get silent(): boolean {
return this.mode === "silent";
}
private getStoredRedirectUrisSync(): string[] | undefined {
if (this._storedRedirectUris !== undefined) {
return this._storedRedirectUris ?? undefined;
}
const stored = loadStoredRedirectUrisSync(this.serverName, this.storageDir, this.serverUrl);
this._storedRedirectUris = stored ?? null;
return stored;
}
private async _mutateInteractiveState(mutator: (state: StoredState) => boolean): Promise<void> {
if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
await this.assertLockHeld?.();
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
const shouldSave = mutator(state);
if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
await this.assertLockHeld?.();
if (!shouldSave) return;
// WHY: Only write back to disk when stored state for this serverUrl actually had
// credentials to clear, preventing a cancelled or unauthenticated project override
// from overwriting another server URL's state file before new credentials are saved.
stampStateServerUrl(state, this.serverUrl);
await saveState(this.serverName, state, this.storageDir, this.isCancelled);
}
// --- redirectUrl ---
get redirectUrl(): string | URL {
return resolveProviderRedirectUrl(
this.mode,
this.explicitRedirectUrl,
this.authConfig,
this.silent ? this.getStoredRedirectUrisSync() : undefined,
);
}
// --- clientMetadata ---
get clientMetadata(): OAuthClientMetadata {
return buildOAuthClientMetadata(
this.serverName,
this.authConfig,
this.mode,
this.explicitRedirectUrl,
this.silent ? this.getStoredRedirectUrisSync() : undefined,
);
}
// --- clientInformation ---
private async _staticClientInformation(): Promise<OAuthClientInformationMixed> {
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
return buildStaticClientInformation(this.authConfig, state.clientInfo);
}
private _cacheStoredRedirectUris(info: StoredClientInfo): void {
if (Array.isArray(info.redirect_uris) && info.redirect_uris.length > 0) {
this._storedRedirectUris = info.redirect_uris;
}
}
private _resolveStoredClientInfo(info: StoredClientInfo): OAuthClientInformationMixed | undefined {
this._cacheStoredRedirectUris(info);
if (isClientSecretExpired(info)) {
if (this.silent) {
throw new AuthRequiredError(
this.serverName,
"Stored OAuth client secret has expired; interactive re-authentication required",
);
}
return undefined;
}
if (!this.silent && hasRedirectUriMismatch(info, String(this.redirectUrl))) {
// WHY: If the callback port or configured redirectUrl changed since DCR,
// return undefined in interactive mode so the SDK re-registers with the new redirect_uri.
return undefined;
}
return toOAuthClientInformation(info);
}
async clientInformation(): Promise<OAuthClientInformationMixed | undefined> {
// If static credentials provided, use those
if (this.authConfig.clientId) {
return this._staticClientInformation();
}
// Otherwise load from persisted DCR state
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
if (state.clientInfo) {
return this._resolveStoredClientInfo(state.clientInfo);
}
if (this.silent) throw new AuthRequiredError(this.serverName);
return undefined;
}
// --- saveClientInformation (DCR) ---
async saveClientInformation(clientInformation: OAuthClientInformationMixed): Promise<void> {
const issuerToBind =
typeof clientInformation.issuer === "string" ? clientInformation.issuer : undefined;
if (this.silent) {
await this._saveClientInformationSilent(clientInformation, issuerToBind);
return;
}
await this._saveClientInformationInteractive(clientInformation, issuerToBind);
}
private _reportClientInfoSaved(): void {
reportDiagnostic({
level: "info",
event: "oauth.client_info_saved",
server: this.serverName,
message: `Stored OAuth client information for "${this.serverName}"`,
fields: { outcome: "saved" },
});
}
private _reportClientInfoPersistFailure(err: unknown): void {
if (err instanceof AuthRequiredError || err instanceof McpError) {
throw err;
}
// WHY: When another process holds the interactive auth lock (e.g. during an interactive
// /mcp:auth or --reset), acquireAuthLock fails fast. We skip this background binding
// without failing the refresh — the SDK will bind client information again on the next refresh.
if (err instanceof AuthLockError) return;
const { kind, message } = describeError(err);
reportDiagnostic({
level: "error",
event: "oauth.client_info_persist_failed",
server: this.serverName,
message: `Failed to persist OAuth client info for "${this.serverName}": ${message}`,
fields: { kind },
});
}
private async _saveClientInformationSilent(
clientInformation: OAuthClientInformationMixed,
issuerToBind: string | undefined,
): Promise<void> {
let handle: AuthLockHandle | undefined;
try {
handle = await acquireAuthLock(this.serverName, this.storageDir);
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
const storedClient = state.clientInfo;
const matchingId = this.authConfig.clientId ?? storedClient?.client_id;
if (!matchingId || clientInformation.client_id !== matchingId) {
// WHY: Silent background connections must never execute Dynamic Client Registration
// for a new or different client_id without interactive user consent.
throw new AuthRequiredError(this.serverName);
}
assertIssuerBindingCompatible(this.serverName, storedClient?.issuer, issuerToBind);
stampStateServerUrl(state, this.serverUrl);
state.clientInfo = storedClient
? mergeSilentStoredClientInfo(storedClient, this.authConfig, issuerToBind)
: buildNewSilentStoredClientInfo(clientInformation, this.authConfig, issuerToBind);
await saveState(this.serverName, state, this.storageDir, this.isCancelled);
this._reportClientInfoSaved();
} catch (err) {
this._reportClientInfoPersistFailure(err);
} finally {
await handle?.();
}
}
private async _saveClientInformationInteractive(
clientInformation: OAuthClientInformationMixed,
issuerToBind: string | undefined,
): Promise<void> {
await this._mutateInteractiveState((state) => {
const storedClient = state.clientInfo;
if (storedClient?.client_id === clientInformation.client_id) {
assertIssuerBindingCompatible(this.serverName, storedClient.issuer, issuerToBind);
}
const storedClientInfo = buildInteractiveStoredClientInfo(
clientInformation,
Boolean(this.authConfig.clientId),
issuerToBind,
);
state.clientInfo = storedClientInfo;
this._cacheStoredRedirectUris(storedClientInfo);
return true;
});
}
// --- tokens ---
async tokens(): Promise<OAuthTokens | undefined> {
const storedTokens =
this._cachedTokens ?? (await loadState(this.serverName, this.storageDir, this.serverUrl)).tokens;
if (!storedTokens) return undefined;
this._lastReadTokens = storedTokens;
// Always return stored tokens — even if expired.
// The SDK's auth() function checks tokens?.refresh_token and attempts
// silent refresh before falling back to a new authorization flow.
// Returning undefined for expired tokens would prevent that silent refresh
// and force the user to re-authenticate via browser every time.
//
// Flow when we return expired tokens:
// transport sends expired access_token → 401
// → auth() sees refresh_token → silent refresh → success → retry
//
// Flow when we return undefined (WRONG):
// transport has no token → auth() → no refresh possible → REDIRECT
// → user must re-authenticate in browser
return toOAuthTokens(storedTokens);
}
// --- saveTokens ---
private async _saveTokensSilent(storedTokens: StoredTokens): Promise<void> {
let handle: AuthLockHandle | undefined;
try {
handle = await acquireAuthLock(this.serverName, this.storageDir);
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
stampStateServerUrl(state, this.serverUrl);
state.tokens = storedTokens;
await saveState(this.serverName, state, this.storageDir, this.isCancelled);
this._cachedTokens = undefined;
this._lastReadTokens = storedTokens;
this._reloadedPeerTokens = false;
reportDiagnostic({
level: "info",
event: "oauth.refresh",
server: this.serverName,
message: `Refreshed OAuth tokens for "${this.serverName}"`,
fields: { outcome: "rotated", peerReloaded: false },
});
} catch (err) {
// WHY: When another process holds the interactive auth lock (e.g. during --reset),
// keep the refreshed tokens in memory for this session without clobbering disk state.
this._cachedTokens = storedTokens;
this._lastReadTokens = storedTokens;
this._reloadedPeerTokens = false;
if (!(err instanceof AuthLockError)) {
const { kind, message } = describeError(err);
reportDiagnostic({
level: "error",
event: "oauth.tokens_persist_failed",
server: this.serverName,
message: `Failed to persist refreshed OAuth tokens for "${this.serverName}": ${message}`,
fields: { kind },
});
}
} finally {
await handle?.();
}
}
private async _saveTokensInteractive(storedTokens: StoredTokens): Promise<void> {
await this._mutateInteractiveState((state) => {
state.tokens = storedTokens;
return true;
});
this._cachedTokens = undefined;
this._lastReadTokens = storedTokens;
this._reloadedPeerTokens = false;
reportDiagnostic({
level: "info",
event: "oauth.tokens_saved",
server: this.serverName,
message: `Stored OAuth tokens for "${this.serverName}"`,
fields: { outcome: "authorized", peerReloaded: false },
});
}
async saveTokens(tokens: OAuthTokens): Promise<void> {
const storedTokens = toStoredTokens(tokens);
if (this.silent) {
await this._saveTokensSilent(storedTokens);
return;
}
await this._saveTokensInteractive(storedTokens);
}
// --- redirectToAuthorization ---
async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
if (this._onAuthRequired) {
await this._onAuthRequired(authorizationUrl);
} else {
// Fallback: no interactive handler is wired, so the authorization URL is
// surfaced through the diagnostic sink. The URL carries no credential
// (no authorization code, no PKCE verifier) and is redacted like any
// other message before it is written to the log file.
reportDiagnostic({
level: "info",
event: "oauth.authorization_required",
server: this.serverName,
message: `OAuth authorization required for "${this.serverName}". Open this URL in your browser: ${authorizationUrl.toString()}`,
fields: { outcome: "redirect" },
notify: true,
// WHY: this fallback fires only headless (no interactive handler is
// wired); the URL is the only way to proceed, so it must reach stderr.
stderr: true,
});
}
}
// --- PKCE code verifier (in-memory per provider instance) ---
async saveCodeVerifier(codeVerifier: string): Promise<void> {
// The SDK saves PKCE before redirecting. Silent attempts must not start a browser flow.
if (this.silent) throw new AuthRequiredError(this.serverName);
if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
this._codeVerifier = codeVerifier;
}
async codeVerifier(): Promise<string> {
if (!this._codeVerifier) {
throw new Error(`[pi-mcp] No PKCE code verifier found for "${this.serverName}"`);
}
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._mutateInteractiveState((state) => {
state.discoveryState = discState;
return true;
});
}
async discoveryState(): Promise<OAuthDiscoveryState | undefined> {
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
return state.discoveryState;
}
// --- OAuth state parameter (CSRF protection) ---
/**
* Set the OAuth state parameter before calling auth().
* This should be called with a cryptographically random value.
*/
setState(state: string): void {
this._oauthState = state;
}
/**
* Returns the OAuth state parameter for CSRF protection.
* This is called by the SDK's auth() function when building the authorization URL.
* Returns empty string if no state has been set (no CSRF protection).
*/
async state(): Promise<string> {
// The SDK reads state() before building the authorization URL. Silent attempts must not start a browser flow.
if (this.silent) throw new AuthRequiredError(this.serverName);
if (this.isCancelled()) throw new AuthRequiredError(this.serverName, "Authorization cancelled");
return this._oauthState || "";
}
// --- Credential invalidation ---
private _applyInvalidation(state: StoredState, scope: InvalidationScope): boolean {
if (scope === "all" || scope === "client") {
this._storedRedirectUris = undefined;
}
if (scope === "all" || scope === "tokens") {
this._cachedTokens = undefined;
}
if (scope === "all" || scope === "verifier") {
this._codeVerifier = undefined;
}
return applyStoredStateInvalidation(state, scope);
}
private async _readDiskStateForPeerReload(): Promise<StoredState | undefined> {
try {
return await loadState(this.serverName, this.storageDir, this.serverUrl);
} catch {
// WHY: If disk state is unreadable, fall through to invalidation and AuthRequiredError.
return undefined;
}
}
private async _tryReloadPeerRotatedTokens(): Promise<boolean> {
if (this._reloadedPeerTokens) return false;
const diskTokens = (await this._readDiskStateForPeerReload())?.tokens;
if (!isPeerRotatedTokens(diskTokens, this._lastReadTokens)) return false;
// WHY: Another Pi process rotated the single-use refresh token on disk while
// our refresh request was in flight. Reload the peer's tokens instead of deleting
// them or throwing AuthRequiredError so the SDK's auth() retry succeeds.
this._cachedTokens = diskTokens;
this._lastReadTokens = diskTokens;
this._reloadedPeerTokens = true;
reportDiagnostic({
level: "info",
event: "oauth.refresh",
server: this.serverName,
message: `Reloaded OAuth tokens rotated by another Pi process for "${this.serverName}"`,
fields: { outcome: "peer_reloaded", peerReloaded: true },
});
return true;
}
private async _removeCredentialsFromDisk(scope: InvalidationScope): Promise<void> {
let handle: AuthLockHandle | undefined;
try {
handle = await acquireAuthLock(this.serverName, this.storageDir);
const state = await loadState(this.serverName, this.storageDir, this.serverUrl);
const changedDisk = this._applyInvalidation(state, scope);
if (!changedDisk) return;
stampStateServerUrl(state, this.serverUrl);
await saveState(this.serverName, state, this.storageDir, this.isCancelled);
} catch {
// WHY: Best-effort cleanup of rejected credentials on disk; if another process
// holds the lock or disk is unavailable, still throw AuthRequiredError.
} finally {
await handle?.();
}
}
private async _invalidateCredentialsSilent(scope: InvalidationScope): Promise<void> {
if (scope === "tokens" && (await this._tryReloadPeerRotatedTokens())) return;
this._cachedTokens = undefined;
this._reloadedPeerTokens = false;
if (scope === "verifier") {
this._codeVerifier = undefined;
throw new AuthRequiredError(this.serverName);
}
await this._removeCredentialsFromDisk(scope);
throw new AuthRequiredError(this.serverName);
}
private async _invalidateCredentialsInteractive(scope: InvalidationScope): Promise<void> {
await this._mutateInteractiveState((state) => this._applyInvalidation(state, scope));
}
async invalidateCredentials(scope: InvalidationScope): Promise<void> {
if (this.silent) {
await this._invalidateCredentialsSilent(scope);
return;
}
await this._invalidateCredentialsInteractive(scope);
}
}
export { getAuthStatus, invalidateStoredTokens, resetAuth } from "./oauth-status.js";