@oh-my-pi/pi-coding-agent 11.8.3 → 11.9.0

package/src/mcp/manager.ts

@@ -6,7 +6,10 @@
  */
 import { logger } from "@oh-my-pi/pi-utils";
 import type { TSchema } from "@sinclair/typebox";
+import type { SourceMeta } from "../capability/types";
+import { resolveConfigValue } from "../config/resolve-config-value";
 import type { CustomTool } from "../extensibility/custom-tools/types";
+import type { AuthStorage } from "../session/auth-storage";
 import { connectToServer, disconnectServer, listTools } from "./client";
 import { loadAllMCPConfigs, validateServerConfig } from "./config";
 import type { MCPToolDetails } from "./tool-bridge";
@@ -14,8 +17,6 @@ import { DeferredMCPTool, MCPTool } from "./tool-bridge";
 import type { MCPToolCache } from "./tool-cache";
 import type { MCPServerConfig, MCPServerConnection, MCPToolDefinition } from "./types";
 
-type SourceMeta = import("../capability/types").SourceMeta;
-
 type ToolLoadResult = {
 	connection: MCPServerConnection;
 	serverTools: MCPToolDefinition[];
@@ -82,12 +83,20 @@ export class MCPManager {
 	#pendingConnections = new Map<string, Promise<MCPServerConnection>>();
 	#pendingToolLoads = new Map<string, Promise<ToolLoadResult>>();
 	#sources = new Map<string, SourceMeta>();
+	#authStorage: AuthStorage | null = null;
 
 	constructor(
 		private cwd: string,
 		private toolCache: MCPToolCache | null = null,
 	) {}
 
+	/**
+	 * Set the auth storage for resolving OAuth credentials.
+	 */
+	setAuthStorage(authStorage: AuthStorage): void {
+		this.#authStorage = authStorage;
+	}
+
 	/**
 	 * Discover and connect to all MCP servers from .mcp.json files.
 	 * Returns tools and any connection errors.
@@ -154,8 +163,14 @@ export class MCPManager {
 				continue;
 			}
 
-			const connectionPromise = connectToServer(name, config).then(
+			// Resolve auth config before connecting
+			const resolvedConfig = await this.#resolveAuthConfig(config);
+
+			const connectionPromise = connectToServer(name, resolvedConfig).then(
 				connection => {
+					// Store original config (without resolved tokens) to keep
+					// cache keys stable and avoid leaking rotating credentials.
+					connection.config = config;
 					if (sources[name]) {
 						connection._source = sources[name];
 					}
@@ -286,6 +301,15 @@ export class MCPManager {
 		return this.#connections.get(name);
 	}
 
+	/**
+	 * Get current connection status for a server.
+	 */
+	getConnectionStatus(name: string): "connected" | "connecting" | "disconnected" {
+		if (this.#connections.has(name)) return "connected";
+		if (this.#pendingConnections.has(name) || this.#pendingToolLoads.has(name)) return "connecting";
+		return "disconnected";
+	}
+
 	/**
 	 * Get the source metadata for a server.
 	 */
@@ -304,6 +328,13 @@ export class MCPManager {
 		throw new Error(`MCP server not connected: ${name}`);
 	}
 
+	/**
+	 * Resolve auth and shell-command substitutions in config before connecting.
+	 */
+	async prepareConfig(config: MCPServerConfig): Promise<MCPServerConfig> {
+		return this.#resolveAuthConfig(config);
+	}
+
 	/**
 	 * Get all connected server names.
 	 */
@@ -369,6 +400,64 @@ export class MCPManager {
 		const promises = Array.from(this.#connections.keys()).map(name => this.refreshServerTools(name));
 		await Promise.allSettled(promises);
 	}
+
+	/**
+	 * Resolve OAuth credentials and shell commands in config.
+	 */
+	async #resolveAuthConfig(config: MCPServerConfig): Promise<MCPServerConfig> {
+		let resolved: MCPServerConfig = { ...config };
+
+		const auth = config.auth;
+		if (auth?.type === "oauth" && auth.credentialId && this.#authStorage) {
+			const credentialId = auth.credentialId;
+			try {
+				const credential = this.#authStorage.get(credentialId);
+				if (credential?.type === "oauth") {
+					if (resolved.type === "http" || resolved.type === "sse") {
+						resolved = {
+							...resolved,
+							headers: {
+								...resolved.headers,
+								Authorization: `Bearer ${credential.access}`,
+							},
+						};
+					} else {
+						resolved = {
+							...resolved,
+							env: {
+								...resolved.env,
+								OAUTH_ACCESS_TOKEN: credential.access,
+							},
+						};
+					}
+				}
+			} catch (error) {
+				logger.warn("Failed to resolve OAuth credential", { credentialId, error });
+			}
+		}
+
+		if (resolved.type !== "http" && resolved.type !== "sse") {
+			if (resolved.env) {
+				const nextEnv: Record<string, string> = {};
+				for (const [key, value] of Object.entries(resolved.env)) {
+					const resolvedValue = await resolveConfigValue(value);
+					if (resolvedValue) nextEnv[key] = resolvedValue;
+				}
+				resolved = { ...resolved, env: nextEnv };
+			}
+		} else {
+			if (resolved.headers) {
+				const nextHeaders: Record<string, string> = {};
+				for (const [key, value] of Object.entries(resolved.headers)) {
+					const resolvedValue = await resolveConfigValue(value);
+					if (resolvedValue) nextHeaders[key] = resolvedValue;
+				}
+				resolved = { ...resolved, headers: nextHeaders };
+			}
+		}
+
+		return resolved;
+	}
 }
 
 /**

package/src/mcp/oauth-discovery.ts

@@ -0,0 +1,274 @@
+/**
+ * MCP OAuth Auto-Discovery
+ *
+ * Automatically detects OAuth requirements from MCP server responses
+ * and extracts authentication endpoints.
+ */
+
+export interface OAuthEndpoints {
+	authorizationUrl: string;
+	tokenUrl: string;
+	clientId?: string;
+	scopes?: string;
+}
+
+export interface AuthDetectionResult {
+	requiresAuth: boolean;
+	authType?: "oauth" | "apikey" | "unknown";
+	oauth?: OAuthEndpoints;
+	message?: string;
+}
+
+/**
+ * Detect if an error indicates authentication is required.
+ * Checks for common auth error patterns.
+ */
+export function detectAuthError(error: Error): boolean {
+	const errorMsg = error.message.toLowerCase();
+
+	// Check for HTTP auth status codes
+	if (
+		errorMsg.includes("401") ||
+		errorMsg.includes("403") ||
+		errorMsg.includes("unauthorized") ||
+		errorMsg.includes("forbidden") ||
+		errorMsg.includes("authentication required") ||
+		errorMsg.includes("authentication failed")
+	) {
+		return true;
+	}
+
+	return false;
+}
+
+/**
+ * Extract OAuth endpoints from error response.
+ * Looks for WWW-Authenticate header format or JSON error bodies.
+ */
+export function extractOAuthEndpoints(error: Error): OAuthEndpoints | null {
+	const errorMsg = error.message;
+
+	const readEndpointsFromObject = (obj: Record<string, unknown>): OAuthEndpoints | null => {
+		const authorizationUrl =
+			(obj.authorization_url as string | undefined) ||
+			(obj.authorizationUrl as string | undefined) ||
+			(obj.authorization_endpoint as string | undefined) ||
+			(obj.authorizationEndpoint as string | undefined) ||
+			(obj.authorization_uri as string | undefined) ||
+			(obj.authorizationUri as string | undefined);
+		const tokenUrl =
+			(obj.token_url as string | undefined) ||
+			(obj.tokenUrl as string | undefined) ||
+			(obj.token_endpoint as string | undefined) ||
+			(obj.tokenEndpoint as string | undefined) ||
+			(obj.token_uri as string | undefined) ||
+			(obj.tokenUri as string | undefined);
+
+		if (!authorizationUrl || !tokenUrl) return null;
+
+		const scopeFromArray = Array.isArray(obj.scopes_supported)
+			? (obj.scopes_supported as unknown[]).filter(v => typeof v === "string").join(" ")
+			: undefined;
+		const scopes = (obj.scopes as string | undefined) || (obj.scope as string | undefined) || scopeFromArray;
+		const clientId =
+			(obj.client_id as string | undefined) ||
+			(obj.clientId as string | undefined) ||
+			(obj.default_client_id as string | undefined) ||
+			(obj.public_client_id as string | undefined);
+
+		return { authorizationUrl, tokenUrl, clientId, scopes };
+	};
+
+	const clientIdFromAuthUrl = (authorizationUrl: string): string | undefined => {
+		try {
+			return new URL(authorizationUrl).searchParams.get("client_id") ?? undefined;
+		} catch {
+			return undefined;
+		}
+	};
+
+	const scopeFromAuthUrl = (authorizationUrl: string): string | undefined => {
+		try {
+			return new URL(authorizationUrl).searchParams.get("scope") ?? undefined;
+		} catch {
+			return undefined;
+		}
+	};
+
+	try {
+		// Try to parse as JSON error response
+		// Many MCP servers return JSON with OAuth endpoints in error body
+		const jsonMatch = errorMsg.match(/\{[\s\S]*\}/);
+		if (jsonMatch) {
+			const errorBody = JSON.parse(jsonMatch[0]) as Record<string, unknown>;
+
+			// Check for OAuth endpoints in error body
+			if (errorBody.oauth || errorBody.authorization || errorBody.auth) {
+				const oauthData = (errorBody.oauth || errorBody.authorization || errorBody.auth) as Record<string, unknown>;
+				const endpoints = readEndpointsFromObject(oauthData);
+				if (endpoints) {
+					return {
+						...endpoints,
+						clientId: endpoints.clientId || clientIdFromAuthUrl(endpoints.authorizationUrl),
+						scopes: endpoints.scopes || scopeFromAuthUrl(endpoints.authorizationUrl),
+					};
+				}
+			}
+
+			const topLevelEndpoints = readEndpointsFromObject(errorBody);
+			if (topLevelEndpoints) {
+				return {
+					...topLevelEndpoints,
+					clientId: topLevelEndpoints.clientId || clientIdFromAuthUrl(topLevelEndpoints.authorizationUrl),
+					scopes: topLevelEndpoints.scopes || scopeFromAuthUrl(topLevelEndpoints.authorizationUrl),
+				};
+			}
+		}
+	} catch {
+		// Not JSON, continue with other detection methods
+	}
+
+	const challengeEntries = Array.from(errorMsg.matchAll(/([a-zA-Z_][a-zA-Z0-9_-]*)="([^"]+)"/g));
+	if (challengeEntries.length > 0) {
+		const challengeValues = new Map<string, string>();
+		for (const [, rawKey, value] of challengeEntries) {
+			challengeValues.set(rawKey.toLowerCase(), value);
+		}
+
+		const authorizationUrl =
+			challengeValues.get("authorization_uri") ||
+			challengeValues.get("authorization_url") ||
+			challengeValues.get("authorization_endpoint") ||
+			challengeValues.get("authorize_url") ||
+			challengeValues.get("realm");
+		const tokenUrl =
+			challengeValues.get("token_url") || challengeValues.get("token_uri") || challengeValues.get("token_endpoint");
+
+		if (authorizationUrl && tokenUrl) {
+			return {
+				authorizationUrl,
+				tokenUrl,
+				clientId: challengeValues.get("client_id") || clientIdFromAuthUrl(authorizationUrl),
+				scopes: challengeValues.get("scope") || challengeValues.get("scopes") || scopeFromAuthUrl(authorizationUrl),
+			};
+		}
+	}
+
+	// Try to extract from WWW-Authenticate header format
+	// Example: Bearer realm="https://auth.example.com/oauth/authorize" token_url="https://auth.example.com/oauth/token"
+	const wwwAuthMatch = errorMsg.match(/realm="([^"]+)".*token_url="([^"]+)"/);
+	if (wwwAuthMatch) {
+		return {
+			authorizationUrl: wwwAuthMatch[1],
+			tokenUrl: wwwAuthMatch[2],
+			clientId: clientIdFromAuthUrl(wwwAuthMatch[1]),
+			scopes: scopeFromAuthUrl(wwwAuthMatch[1]),
+		};
+	}
+
+	return null;
+}
+
+/**
+ * Analyze an error to determine authentication requirements.
+ * Returns structured info about what auth is needed.
+ */
+export function analyzeAuthError(error: Error): AuthDetectionResult {
+	if (!detectAuthError(error)) {
+		return { requiresAuth: false };
+	}
+
+	// Try to extract OAuth endpoints
+	const oauth = extractOAuthEndpoints(error);
+
+	if (oauth) {
+		return {
+			requiresAuth: true,
+			authType: "oauth",
+			oauth,
+			message: "Server requires OAuth authentication. Launching authorization flow...",
+		};
+	}
+
+	// Check if it might be API key based
+	const errorMsg = error.message.toLowerCase();
+	if (
+		errorMsg.includes("api key") ||
+		errorMsg.includes("api_key") ||
+		errorMsg.includes("token") ||
+		errorMsg.includes("bearer")
+	) {
+		return {
+			requiresAuth: true,
+			authType: "apikey",
+			message: "Server requires API key authentication.",
+		};
+	}
+
+	// Unknown auth type
+	return {
+		requiresAuth: true,
+		authType: "unknown",
+		message: "Server requires authentication but type could not be determined.",
+	};
+}
+
+/**
+ * Try to discover OAuth endpoints by querying the server's well-known endpoints.
+ * This is a fallback when error responses don't include OAuth metadata.
+ */
+export async function discoverOAuthEndpoints(serverUrl: string): Promise<OAuthEndpoints | null> {
+	const wellKnownPaths = [
+		"/.well-known/oauth-authorization-server",
+		"/.well-known/openid-configuration",
+		"/oauth/metadata",
+		"/.mcp/auth",
+		"/authorize", // Some MCP servers expose OAuth config here
+	];
+
+	for (const path of wellKnownPaths) {
+		try {
+			const url = new URL(path, serverUrl);
+			const response = await fetch(url.toString(), {
+				method: "GET",
+				headers: { Accept: "application/json" },
+			});
+
+			if (response.ok) {
+				const metadata = await response.json();
+
+				// Check for standard OAuth discovery format
+				if (metadata.authorization_endpoint && metadata.token_endpoint) {
+					return {
+						authorizationUrl: metadata.authorization_endpoint,
+						tokenUrl: metadata.token_endpoint,
+						clientId:
+							metadata.client_id || metadata.clientId || metadata.default_client_id || metadata.public_client_id,
+						scopes: metadata.scopes_supported?.join(" ") || metadata.scopes || metadata.scope,
+					};
+				}
+
+				// Check for MCP-specific format
+				if (metadata.oauth || metadata.authorization || metadata.auth) {
+					const oauthData = metadata.oauth || metadata.authorization || metadata.auth;
+					if (oauthData.authorization_url && oauthData.token_url) {
+						return {
+							authorizationUrl: oauthData.authorization_url || oauthData.authorizationUrl,
+							tokenUrl: oauthData.token_url || oauthData.tokenUrl,
+							clientId:
+								oauthData.client_id ||
+								oauthData.clientId ||
+								oauthData.default_client_id ||
+								oauthData.public_client_id,
+							scopes: oauthData.scopes || oauthData.scope,
+						};
+					}
+				}
+			}
+		} catch {
+			// Ignore errors, try next path
+		}
+	}
+
+	return null;
+}

package/src/mcp/oauth-flow.ts

@@ -0,0 +1,229 @@
+/**
+ * Generic OAuth flow for MCP servers.
+ *
+ * Allows users to authenticate with any OAuth-compatible MCP server
+ * by providing authorization URL, token URL, and client credentials.
+ */
+
+import type { OAuthController, OAuthCredentials } from "@oh-my-pi/pi-ai";
+import { OAuthCallbackFlow } from "@oh-my-pi/pi-ai/utils/oauth/callback-server";
+
+const DEFAULT_PORT = 3000;
+const CALLBACK_PATH = "/callback";
+
+export interface MCPOAuthConfig {
+	/** Authorization endpoint URL */
+	authorizationUrl: string;
+	/** Token endpoint URL */
+	tokenUrl: string;
+	/** Client ID (optional when already embedded in authorization URL) */
+	clientId?: string;
+	/** Client secret (optional for PKCE flows) */
+	clientSecret?: string;
+	/** OAuth scopes (space-separated) */
+	scopes?: string;
+}
+
+/**
+ * Generic OAuth flow for MCP servers.
+ * Supports standard OAuth 2.0 authorization code flow with PKCE.
+ */
+export class MCPOAuthFlow extends OAuthCallbackFlow {
+	#resolvedClientId?: string;
+	#registeredClientSecret?: string;
+	#codeVerifier?: string;
+
+	constructor(
+		private config: MCPOAuthConfig,
+		ctrl: OAuthController,
+	) {
+		super(ctrl, DEFAULT_PORT, CALLBACK_PATH);
+		this.#resolvedClientId = this.#resolveClientId(config);
+	}
+
+	async generateAuthUrl(state: string, redirectUri: string): Promise<{ url: string; instructions?: string }> {
+		if (!this.#resolvedClientId) {
+			await this.#tryRegisterClient(redirectUri);
+		}
+
+		const authUrl = new URL(this.config.authorizationUrl);
+		const params = authUrl.searchParams;
+
+		if (!params.get("response_type")) {
+			params.set("response_type", "code");
+		}
+		if (this.#resolvedClientId && !params.get("client_id")) {
+			params.set("client_id", this.#resolvedClientId);
+		}
+		if (this.config.scopes && !params.get("scope")) {
+			params.set("scope", this.config.scopes);
+		}
+		params.set("redirect_uri", redirectUri);
+		params.set("state", state);
+
+		// Add PKCE challenge (some providers require it)
+		const codeVerifier = this.#generateCodeVerifier();
+		const codeChallenge = await this.#generateCodeChallenge(codeVerifier);
+		params.set("code_challenge", codeChallenge);
+		params.set("code_challenge_method", "S256");
+
+		// Store code verifier for token exchange
+		this.#codeVerifier = codeVerifier;
+
+		return { url: authUrl.toString() };
+	}
+
+	async exchangeToken(code: string, _state: string, redirectUri: string): Promise<OAuthCredentials> {
+		const params = new URLSearchParams({
+			grant_type: "authorization_code",
+			code,
+			redirect_uri: redirectUri,
+		});
+		if (this.#resolvedClientId) {
+			params.set("client_id", this.#resolvedClientId);
+		}
+
+		// Add code verifier for PKCE
+		if (this.#codeVerifier) {
+			params.set("code_verifier", this.#codeVerifier);
+		}
+		this.#codeVerifier = undefined;
+
+		// Add client secret if provided
+		const clientSecret = this.config.clientSecret ?? this.#registeredClientSecret;
+		if (clientSecret) {
+			params.set("client_secret", clientSecret);
+		}
+
+		const response = await fetch(this.config.tokenUrl, {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/x-www-form-urlencoded",
+			},
+			body: params.toString(),
+		});
+
+		if (!response.ok) {
+			const errorText = await response.text();
+			throw new Error(`Token exchange failed: ${response.status} ${errorText}`);
+		}
+
+		const data = (await response.json()) as {
+			access_token: string;
+			refresh_token?: string;
+			expires_in?: number;
+			token_type?: string;
+		};
+
+		// Calculate expiry timestamp
+		const expiresIn = data.expires_in ?? 3600; // Default to 1 hour
+		const expires = Date.now() + expiresIn * 1000;
+
+		return {
+			access: data.access_token,
+			refresh: data.refresh_token ?? "",
+			expires,
+		};
+	}
+
+	/**
+	 * Generate PKCE code verifier (random string).
+	 */
+	#generateCodeVerifier(): string {
+		const bytes = new Uint8Array(32);
+		crypto.getRandomValues(bytes);
+		return this.#base64UrlEncode(bytes);
+	}
+
+	/**
+	 * Generate PKCE code challenge from verifier.
+	 */
+	async #generateCodeChallenge(verifier: string): Promise<string> {
+		const encoder = new TextEncoder();
+		const data = encoder.encode(verifier);
+		const hash = await crypto.subtle.digest("SHA-256", data);
+		return this.#base64UrlEncode(new Uint8Array(hash));
+	}
+
+	/**
+	 * Base64 URL encode (without padding).
+	 */
+	#base64UrlEncode(bytes: Uint8Array): string {
+		const base64 = btoa(String.fromCharCode(...bytes));
+		return base64.replace(/\+/g, "-").replace(/\//g, "_").replace(/=/g, "");
+	}
+
+	#resolveClientId(config: MCPOAuthConfig): string | undefined {
+		const fromConfig = config.clientId?.trim();
+		if (fromConfig) return fromConfig;
+
+		try {
+			return new URL(config.authorizationUrl).searchParams.get("client_id") ?? undefined;
+		} catch {
+			return undefined;
+		}
+	}
+
+	/**
+	 * Try OAuth dynamic client registration when provider requires a client_id.
+	 */
+	async #tryRegisterClient(redirectUri: string): Promise<void> {
+		const registrationEndpoint = await this.#resolveRegistrationEndpoint();
+		if (!registrationEndpoint) return;
+
+		try {
+			const response = await fetch(registrationEndpoint, {
+				method: "POST",
+				headers: {
+					"Content-Type": "application/json",
+					Accept: "application/json",
+				},
+				body: JSON.stringify({
+					client_name: "oh-my-pi MCP",
+					redirect_uris: [redirectUri],
+					grant_types: ["authorization_code", "refresh_token"],
+					response_types: ["code"],
+					token_endpoint_auth_method: "none",
+					application_type: "native",
+				}),
+			});
+
+			if (!response.ok) return;
+
+			const data = (await response.json()) as {
+				client_id?: string;
+				client_secret?: string;
+			};
+
+			if (data.client_id && data.client_id.trim() !== "") {
+				this.#resolvedClientId = data.client_id;
+			}
+			if (data.client_secret && data.client_secret.trim() !== "") {
+				this.#registeredClientSecret = data.client_secret;
+			}
+		} catch {
+			// Ignore registration failures and continue without client registration.
+		}
+	}
+
+	async #resolveRegistrationEndpoint(): Promise<string | null> {
+		try {
+			const authorizationEndpoint = new URL(this.config.authorizationUrl);
+			const metadataUrl = new URL("/.well-known/oauth-authorization-server", authorizationEndpoint.origin);
+			const response = await fetch(metadataUrl.toString(), {
+				method: "GET",
+				headers: { Accept: "application/json" },
+			});
+
+			if (!response.ok) return null;
+			const metadata = (await response.json()) as { registration_endpoint?: string };
+			if (metadata.registration_endpoint && metadata.registration_endpoint.trim() !== "") {
+				return metadata.registration_endpoint;
+			}
+		} catch {
+			// Ignore metadata discovery failures.
+		}
+
+		return null;
+	}
+}

package/src/mcp/tool-bridge.ts

@@ -219,8 +219,8 @@ export class DeferredMCPTool implements CustomTool<TSchema, MCPToolDetails> {
 	readonly mcpToolName: string;
 	/** Server name */
 	readonly mcpServerName: string;
-	private readonly fallbackProvider: string | undefined;
-	private readonly fallbackProviderName: string | undefined;
+	readonly #fallbackProvider: string | undefined;
+	readonly #fallbackProviderName: string | undefined;
 
 	/** Create DeferredMCPTool instances for all tools from an MCP server */
 	static fromTools(
@@ -244,8 +244,8 @@ export class DeferredMCPTool implements CustomTool<TSchema, MCPToolDetails> {
 		this.parameters = convertSchema(tool.inputSchema);
 		this.mcpToolName = tool.name;
 		this.mcpServerName = serverName;
-		this.fallbackProvider = source?.provider;
-		this.fallbackProviderName = source?.providerName;
+		this.#fallbackProvider = source?.provider;
+		this.#fallbackProviderName = source?.providerName;
 	}
 
 	renderCall(args: unknown, theme: Theme) {
@@ -273,8 +273,8 @@ export class DeferredMCPTool implements CustomTool<TSchema, MCPToolDetails> {
 				mcpToolName: this.tool.name,
 				isError: result.isError,
 				rawContent: result.content,
-				provider: connection._source?.provider ?? this.fallbackProvider,
-				providerName: connection._source?.providerName ?? this.fallbackProviderName,
+				provider: connection._source?.provider ?? this.#fallbackProvider,
+				providerName: connection._source?.providerName ?? this.#fallbackProviderName,
 			};
 
 			if (result.isError) {
@@ -296,8 +296,8 @@ export class DeferredMCPTool implements CustomTool<TSchema, MCPToolDetails> {
 					serverName: this.serverName,
 					mcpToolName: this.tool.name,
 					isError: true,
-					provider: this.fallbackProvider,
-					providerName: this.fallbackProviderName,
+					provider: this.#fallbackProvider,
+					providerName: this.#fallbackProviderName,
 				},
 			};
 		}