MCP OAuth Authorization

┌─────────────────────────────────────────────────────────────────────────────────┐
│                           MCP OAuth Authorization Flow                           │
├─────────────────────────────────────────────────────────────────────────────────┤
│                                                                                  │
│  ┌──────────────┐    1. MCP Request     ┌──────────────────┐                    │
│  │              │ ───────────────────►  │                  │                    │
│  │  Kilo Code   │                       │   MCP Server     │                    │
│  │  Extension   │  ◄─────────────────── │  (Resource       │                    │
│  │              │    2. 401 + metadata  │   Server)        │                    │
│  └──────┬───────┘                       └──────────────────┘                    │
│         │                                                                        │
│         │ 3. Fetch resource metadata                                            │
│         │ 4. Fetch auth server metadata                                         │
│         ▼                                                                        │
│  ┌──────────────┐                       ┌──────────────────┐                    │
│  │   OAuth      │    5. Auth Request    │                  │                    │
│  │   Service    │ ───────────────────►  │  Authorization   │                    │
│  │              │                       │  Server          │                    │
│  │  - Discovery │  ◄─────────────────── │                  │                    │
│  │  - PKCE      │    8. Token Response  │  - User Auth     │                    │
│  │  - Tokens    │                       │  - Consent       │                    │
│  └──────┬───────┘                       └──────────────────┘                    │
│         │                                        ▲                               │
│         │ 6. Open browser                        │ 7. User authenticates        │
│         ▼                                        │                               │
│  ┌──────────────┐                       ┌────────┴─────────┐                    │
│  │   Browser    │ ─────────────────────►│      User        │                    │
│  │              │                       │                  │                    │
│  └──────────────┘                       └──────────────────┘                    │
│                                                                                  │
└─────────────────────────────────────────────────────────────────────────────────┘

// src/services/mcp/oauth/McpOAuthService.ts

interface McpOAuthService {
	/**
	 * Initiates OAuth flow for an MCP server that returned 401
	 * @param serverUrl The MCP server URL
	 * @param wwwAuthenticateHeader The WWW-Authenticate header from 401 response
	 * @returns Promise resolving to access token
	 */
	initiateOAuthFlow(serverUrl: string, wwwAuthenticateHeader: string): Promise<OAuthTokens>

	/**
	 * Gets stored tokens for a server, if available and valid
	 */
	getStoredTokens(serverUrl: string): Promise<OAuthTokens | null>

	/**
	 * Clears stored tokens for a server (for logout/re-auth)
	 */
	clearTokens(serverUrl: string): Promise<void>

	/**
	 * Refreshes tokens if refresh token is available
	 */
	refreshTokens(serverUrl: string): Promise<OAuthTokens | null>
}

interface OAuthTokens {
	accessToken: string
	tokenType: string
	expiresAt?: number
	refreshToken?: string
	scope?: string
}

// src/services/mcp/oauth/McpAuthorizationDiscovery.ts

interface McpAuthorizationDiscovery {
	/**
	 * Discovers authorization server from WWW-Authenticate header or well-known URIs
	 */
	discoverAuthorizationServer(serverUrl: string, wwwAuthenticateHeader?: string): Promise<AuthorizationServerMetadata>

	/**
	 * Fetches Protected Resource Metadata (RFC 9728)
	 */
	fetchResourceMetadata(metadataUrl: string): Promise<ProtectedResourceMetadata>

	/**
	 * Fetches Authorization Server Metadata (RFC 8414 / OIDC Discovery)
	 */
	fetchAuthServerMetadata(issuerUrl: string): Promise<AuthorizationServerMetadata>
}

interface ProtectedResourceMetadata {
	resource: string
	authorization_servers: string[]
	scopes_supported?: string[]
	// ... other RFC 9728 fields
}

interface AuthorizationServerMetadata {
	issuer: string
	authorization_endpoint: string
	token_endpoint: string
	scopes_supported?: string[]
	response_types_supported: string[]
	code_challenge_methods_supported?: string[]
	client_id_metadata_document_supported?: boolean
	registration_endpoint?: string
	// ... other RFC 8414 fields
}

// src/services/mcp/oauth/McpOAuthTokenStorage.ts

interface McpOAuthTokenStorage {
	/**
	 * Stores tokens securely using VS Code SecretStorage
	 */
	storeTokens(serverUrl: string, tokens: OAuthTokens): Promise<void>

	/**
	 * Retrieves stored tokens
	 */
	getTokens(serverUrl: string): Promise<OAuthTokens | null>

	/**
	 * Removes stored tokens
	 */
	removeTokens(serverUrl: string): Promise<void>

	/**
	 * Lists all servers with stored tokens
	 */
	listServers(): Promise<string[]>
}

{
	"client_id": "https://kilocode.ai/.well-known/oauth-client/vscode-extension.json",
	"client_name": "Kilo Code",
	"client_uri": "https://kilocode.ai",
	"logo_uri": "https://kilocode.ai/logo.png",
	"redirect_uris": ["http://127.0.0.1:0/callback", "vscode://kilocode.kilo-code/oauth/callback"],
	"grant_types": ["authorization_code"],
	"response_types": ["code"],
	"token_endpoint_auth_method": "none"
}

// Modifications to McpHub.ts

class McpHub {
	private oauthService: McpOAuthService

	private async connectToServer(name: string, config: ServerConfig, source: "global" | "project"): Promise<void> {
		// ... existing connection logic ...

		// For HTTP-based transports, handle OAuth
		if (config.type === "sse" || config.type === "streamable-http") {
			try {
				await this.connectWithOAuth(name, config, source)
			} catch (error) {
				if (this.isOAuthRequired(error)) {
					// Initiate OAuth flow
					const tokens = await this.oauthService.initiateOAuthFlow(config.url, error.wwwAuthenticateHeader)
					// Retry connection with token
					await this.connectWithToken(name, config, source, tokens)
				} else {
					throw error
				}
			}
		}
	}

	private isOAuthRequired(error: unknown): boolean {
		// Check if error is 401 with WWW-Authenticate header
		return error instanceof HttpError && error.status === 401 && error.headers?.["www-authenticate"]
	}
}

// Extended server config for OAuth-enabled servers
const OAuthServerConfigSchema = BaseConfigSchema.extend({
	type: z.enum(["sse", "streamable-http"]),
	url: z.string().url(),
	headers: z.record(z.string()).optional(),

	// OAuth configuration
	oauth: z
		.object({
			// Override client_id if pre-registered
			clientId: z.string().optional(),
			clientSecret: z.string().optional(),

			// Override scopes to request
			scopes: z.array(z.string()).optional(),

			// Disable OAuth for this server (use static headers instead)
			disabled: z.boolean().optional(),
		})
		.optional(),
})

// src/services/mcp/oauth/McpOAuthBrowserFlow.ts

interface McpOAuthBrowserFlow {
	/**
	 * Opens browser for authorization and waits for callback
	 */
	authorize(params: AuthorizationParams): Promise<AuthorizationResult>
}

interface AuthorizationParams {
	authorizationEndpoint: string
	clientId: string
	redirectUri: string
	scope: string
	state: string
	codeChallenge: string
	codeChallengeMethod: "S256"
	resource: string
}

interface AuthorizationResult {
	code: string
	state: string
}

// Key format: mcp-oauth-{serverUrlHash}
const storageKey = `mcp-oauth-${hashServerUrl(serverUrl)}`

// Stored value (encrypted by VS Code)
interface StoredTokenData {
	accessToken: string
	refreshToken?: string
	expiresAt?: number
	scope?: string
	serverUrl: string
	issuedAt: number
}

// OAuth-specific errors
class McpOAuthError extends Error {
	constructor(
		message: string,
		public code: OAuthErrorCode,
		public serverUrl: string,
		public details?: Record<string, unknown>,
	) {
		super(message)
	}
}

enum OAuthErrorCode {
	DISCOVERY_FAILED = "discovery_failed",
	AUTHORIZATION_FAILED = "authorization_failed",
	TOKEN_EXCHANGE_FAILED = "token_exchange_failed",
	TOKEN_REFRESH_FAILED = "token_refresh_failed",
	PKCE_NOT_SUPPORTED = "pkce_not_supported",
	USER_CANCELLED = "user_cancelled",
	TIMEOUT = "timeout",
}

┌─────────────────────────────────────────────────────────────┐
│ MCP Server: github-mcp                                      │
├─────────────────────────────────────────────────────────────┤
│ Status: Connected                                           │
│ Type: streamable-http                                       │
│ URL: https://mcp.github.com                                 │
│                                                             │
│ Authentication                                              │
│ - Method: OAuth 2.0                                         │
│ - Status: Authenticated                                     │
│ - Expires: 2024-01-15 10:30 AM                              │
│ - [Sign Out] [Re-authenticate]                              │
└─────────────────────────────────────────────────────────────┘

function generatePKCE(): { verifier: string; challenge: string } {
	// Generate 32-byte random verifier
	const verifier = base64UrlEncode(crypto.randomBytes(32))

	// Create S256 challenge
	const challenge = base64UrlEncode(crypto.createHash("sha256").update(verifier).digest())

	return { verifier, challenge }
}

const state = base64UrlEncode(crypto.randomBytes(32))
// Store state locally and verify on callback

const authUrl = new URL(authorizationEndpoint)
authUrl.searchParams.set("resource", mcpServerUrl)