From ac48aa58e05a35ce73912c2d631aefbe084fe984 Mon Sep 17 00:00:00 2001 From: David Kaya Date: Fri, 27 Mar 2026 19:15:37 +0100 Subject: [PATCH] feat: implement MCP OAuth 2.1 flow, token storage, and injection - Create mcpTokenStore: in-memory token storage keyed by normalized server URL with automatic expiry checking - Create mcpOAuthService: full OAuth 2.1 + PKCE flow implementation - Protected Resource Metadata discovery (RFC 9728) - Authorization Server Metadata fetch (RFC 8414) - Dynamic Client Registration (RFC 7591) when no static client ID - PKCE S256 code challenge generation - Local HTTP callback server for auth code receipt - Browser-based consent via Electron shell.openExternal - Authorization code to token exchange - Add startSessionMcpAuth IPC channel and handler to trigger OAuth flow - Inject stored OAuth tokens as Authorization headers in buildRunTurnToolingConfig - Update McpAuthBanner with 'Authenticate in browser' button and loading state - Add tests for token store (7 tests) and token injection (3 tests) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/main/AryxAppService.ts | 43 ++- src/main/ipc/registerIpcHandlers.ts | 4 + src/main/services/mcpOAuthService.ts | 333 ++++++++++++++++++ src/main/services/mcpTokenStore.ts | 49 +++ src/main/sessionToolingConfig.ts | 19 +- src/preload/index.ts | 1 + src/renderer/App.tsx | 3 + src/renderer/components/ChatPane.tsx | 7 + .../components/chat/McpAuthBanner.tsx | 37 +- src/shared/contracts/channels.ts | 1 + src/shared/contracts/ipc.ts | 5 + tests/main/mcpTokenStore.test.ts | 86 +++++ tests/main/sessionToolingConfig.test.ts | 55 +++ 13 files changed, 635 insertions(+), 8 deletions(-) create mode 100644 src/main/services/mcpOAuthService.ts create mode 100644 src/main/services/mcpTokenStore.ts create mode 100644 tests/main/mcpTokenStore.test.ts diff --git a/src/main/AryxAppService.ts b/src/main/AryxAppService.ts index 8aa67b6..a39cdc4 100644 --- a/src/main/AryxAppService.ts +++ b/src/main/AryxAppService.ts @@ -109,6 +109,8 @@ import { buildRunTurnToolingConfig as buildSessionToolingConfig, validateSessionToolingSelectionIds, } from '@main/sessionToolingConfig'; +import { getStoredToken } from '@main/services/mcpTokenStore'; +import { performMcpOAuthFlow } from '@main/services/mcpOAuthService'; const { dialog, shell } = electron; @@ -889,6 +891,42 @@ export class AryxAppService extends EventEmitter { return this.persistAndBroadcast(workspace); } + async startSessionMcpAuth(sessionId: string): Promise { + const workspace = await this.loadWorkspace(); + const session = this.requireSession(workspace, sessionId); + + if (!session.pendingMcpAuth) { + return workspace; + } + + session.pendingMcpAuth.status = 'authenticating'; + session.updatedAt = nowIso(); + await this.persistAndBroadcast(workspace); + + const result = await performMcpOAuthFlow({ + serverUrl: session.pendingMcpAuth.serverUrl, + staticClientConfig: session.pendingMcpAuth.staticClientConfig, + }); + + const workspaceAfter = await this.loadWorkspace(); + const sessionAfter = this.requireSession(workspaceAfter, sessionId); + + if (!sessionAfter.pendingMcpAuth) { + return workspaceAfter; + } + + if (result.success) { + sessionAfter.pendingMcpAuth.status = 'authenticated'; + sessionAfter.pendingMcpAuth.completedAt = nowIso(); + } else { + sessionAfter.pendingMcpAuth.status = 'failed'; + sessionAfter.pendingMcpAuth.errorMessage = result.error ?? 'Authentication failed'; + } + + sessionAfter.updatedAt = nowIso(); + return this.persistAndBroadcast(workspaceAfter); + } + async updateSessionTooling( sessionId: string, enabledMcpServerIds: string[], @@ -1600,7 +1638,10 @@ export class AryxAppService extends EventEmitter { const tooling = resolveProjectToolingSettings(workspace.settings, project.discoveredTooling); const selection = resolveSessionToolingSelection(session); validateSessionToolingSelectionIds(tooling, selection); - return buildSessionToolingConfig(tooling, selection); + return buildSessionToolingConfig(tooling, selection, (serverUrl) => { + const token = getStoredToken(serverUrl); + return token?.accessToken; + }); } private async syncUserDiscoveredTooling(workspace: WorkspaceState): Promise { diff --git a/src/main/ipc/registerIpcHandlers.ts b/src/main/ipc/registerIpcHandlers.ts index 7b33d40..5586472 100644 --- a/src/main/ipc/registerIpcHandlers.ts +++ b/src/main/ipc/registerIpcHandlers.ts @@ -9,6 +9,7 @@ import type { ResolveWorkspaceDiscoveredToolingInput, DismissSessionPlanReviewInput, DismissSessionMcpAuthInput, + StartSessionMcpAuthInput, DuplicateSessionInput, RenameSessionInput, RescanProjectConfigsInput, @@ -126,6 +127,9 @@ export function registerIpcHandlers(window: BrowserWindow, service: AryxAppServi ipcMain.handle(ipcChannels.dismissSessionMcpAuth, (_event, input: DismissSessionMcpAuthInput) => service.dismissSessionMcpAuth(input.sessionId), ); + ipcMain.handle(ipcChannels.startSessionMcpAuth, (_event, input: StartSessionMcpAuthInput) => + service.startSessionMcpAuth(input.sessionId), + ); ipcMain.handle( ipcChannels.updateSessionModelConfig, (_event, input: UpdateSessionModelConfigInput) => diff --git a/src/main/services/mcpOAuthService.ts b/src/main/services/mcpOAuthService.ts new file mode 100644 index 0000000..0495e25 --- /dev/null +++ b/src/main/services/mcpOAuthService.ts @@ -0,0 +1,333 @@ +import { randomBytes, createHash } from 'node:crypto'; +import { createServer, type Server, type IncomingMessage, type ServerResponse } from 'node:http'; + +import { shell } from 'electron'; + +import type { McpOauthStaticClientConfig } from '@shared/domain/mcpAuth'; + +import { storeToken, type McpOAuthToken } from './mcpTokenStore'; + +/* ── Public API ──────────────────────────────────────────────── */ + +export interface McpOAuthFlowOptions { + serverUrl: string; + staticClientConfig?: McpOauthStaticClientConfig; + onStatusChange?: (status: 'discovering' | 'awaiting-consent' | 'exchanging') => void; +} + +export interface McpOAuthFlowResult { + success: boolean; + token?: McpOAuthToken; + error?: string; +} + +/** + * Performs the full MCP OAuth 2.1 + PKCE flow: + * 1. Discover protected resource metadata (RFC 9728) + * 2. Fetch authorization server metadata (RFC 8414) + * 3. Resolve client ID (static config or dynamic registration per RFC 7591) + * 4. PKCE code verifier + challenge + * 5. Open browser for user consent + * 6. Local callback server receives auth code + * 7. Exchange code for token + */ +export async function performMcpOAuthFlow(options: McpOAuthFlowOptions): Promise { + const { serverUrl, staticClientConfig, onStatusChange } = options; + + try { + onStatusChange?.('discovering'); + + const authServerUrl = await discoverAuthorizationServer(serverUrl); + const metadata = await fetchAuthServerMetadata(authServerUrl); + + const clientId = staticClientConfig?.clientId + ?? await dynamicClientRegistration(metadata, serverUrl); + + const { verifier, challenge } = generatePkceChallenge(); + const { port, redirectUri, waitForCallback, close } = await startCallbackServer(); + + try { + const scopes = metadata.scopes_supported?.join(' ') ?? ''; + const authUrl = buildAuthorizationUrl(metadata.authorization_endpoint, { + clientId, + redirectUri, + codeChallenge: challenge, + scope: scopes, + }); + + onStatusChange?.('awaiting-consent'); + await shell.openExternal(authUrl); + + const code = await waitForCallback(); + + onStatusChange?.('exchanging'); + const token = await exchangeCodeForToken(metadata.token_endpoint, { + code, + clientId, + redirectUri, + codeVerifier: verifier, + }); + + storeToken(serverUrl, token); + return { success: true, token }; + } finally { + close(); + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + return { success: false, error: message }; + } +} + +/* ── Discovery ───────────────────────────────────────────────── */ + +interface ProtectedResourceMetadata { + resource: string; + authorization_servers?: string[]; +} + +interface AuthServerMetadata { + issuer: string; + authorization_endpoint: string; + token_endpoint: string; + registration_endpoint?: string; + scopes_supported?: string[]; +} + +async function discoverAuthorizationServer(serverUrl: string): Promise { + const base = serverUrl.replace(/\/+$/, ''); + const prmUrl = `${base}/.well-known/oauth-protected-resource`; + + const response = await fetch(prmUrl); + if (!response.ok) { + throw new Error(`Protected Resource Metadata discovery failed: ${response.status} ${response.statusText}`); + } + + const metadata: ProtectedResourceMetadata = await response.json(); + const authServer = metadata.authorization_servers?.[0]; + if (!authServer) { + throw new Error('No authorization server found in Protected Resource Metadata'); + } + + return authServer; +} + +async function fetchAuthServerMetadata(authServerUrl: string): Promise { + const base = authServerUrl.replace(/\/+$/, ''); + const metadataUrl = `${base}/.well-known/oauth-authorization-server`; + + const response = await fetch(metadataUrl); + if (!response.ok) { + throw new Error(`Authorization Server Metadata fetch failed: ${response.status} ${response.statusText}`); + } + + const metadata: AuthServerMetadata = await response.json(); + if (!metadata.authorization_endpoint || !metadata.token_endpoint) { + throw new Error('Authorization server metadata is missing required endpoints'); + } + + return metadata; +} + +/* ── Dynamic Client Registration (RFC 7591) ──────────────────── */ + +async function dynamicClientRegistration(metadata: AuthServerMetadata, serverUrl: string): Promise { + if (!metadata.registration_endpoint) { + throw new Error( + 'No static client ID provided and the authorization server does not support dynamic client registration', + ); + } + + const response = await fetch(metadata.registration_endpoint, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + client_name: 'Aryx', + redirect_uris: ['http://127.0.0.1/callback'], + grant_types: ['authorization_code'], + response_types: ['code'], + token_endpoint_auth_method: 'none', + scope: metadata.scopes_supported?.join(' ') ?? '', + }), + }); + + if (!response.ok) { + throw new Error(`Dynamic client registration failed: ${response.status} ${response.statusText}`); + } + + const registration = await response.json(); + if (!registration.client_id) { + throw new Error('Dynamic client registration response is missing client_id'); + } + + return registration.client_id; +} + +/* ── PKCE ────────────────────────────────────────────────────── */ + +function generatePkceChallenge(): { verifier: string; challenge: string } { + const verifier = randomBytes(32).toString('base64url'); + const challenge = createHash('sha256').update(verifier).digest('base64url'); + return { verifier, challenge }; +} + +/* ── Authorization URL ───────────────────────────────────────── */ + +function buildAuthorizationUrl( + authorizationEndpoint: string, + params: { + clientId: string; + redirectUri: string; + codeChallenge: string; + scope: string; + }, +): string { + const url = new URL(authorizationEndpoint); + url.searchParams.set('response_type', 'code'); + url.searchParams.set('client_id', params.clientId); + url.searchParams.set('redirect_uri', params.redirectUri); + url.searchParams.set('code_challenge', params.codeChallenge); + url.searchParams.set('code_challenge_method', 'S256'); + if (params.scope) { + url.searchParams.set('scope', params.scope); + } + url.searchParams.set('state', randomBytes(16).toString('hex')); + return url.toString(); +} + +/* ── Local callback server ───────────────────────────────────── */ + +interface CallbackServerHandle { + port: number; + redirectUri: string; + waitForCallback: () => Promise; + close: () => void; +} + +function startCallbackServer(): Promise { + return new Promise((resolve, reject) => { + let settled = false; + let callbackResolve: (code: string) => void; + let callbackReject: (err: Error) => void; + + const callbackPromise = new Promise((res, rej) => { + callbackResolve = res; + callbackReject = rej; + }); + + const server: Server = createServer((req: IncomingMessage, res: ServerResponse) => { + if (!req.url?.startsWith('/callback')) { + res.writeHead(404); + res.end(); + return; + } + + const url = new URL(req.url, `http://127.0.0.1`); + const code = url.searchParams.get('code'); + const error = url.searchParams.get('error'); + const errorDescription = url.searchParams.get('error_description'); + + res.writeHead(200, { 'Content-Type': 'text/html' }); + if (code) { + res.end('

Authentication successful

You can close this tab.

'); + callbackResolve(code); + } else { + const msg = errorDescription ?? error ?? 'Unknown error'; + res.end(`

Authentication failed

${escapeHtml(msg)}

`); + callbackReject(new Error(`OAuth callback error: ${msg}`)); + } + }); + + server.on('error', (err) => { + if (!settled) { + settled = true; + reject(err); + } + }); + + server.listen(0, '127.0.0.1', () => { + settled = true; + const addr = server.address(); + if (!addr || typeof addr === 'string') { + reject(new Error('Failed to bind callback server')); + return; + } + + resolve({ + port: addr.port, + redirectUri: `http://127.0.0.1:${addr.port}/callback`, + waitForCallback: () => callbackPromise, + close: () => server.close(), + }); + }); + + setTimeout(() => { + if (!settled) { + settled = true; + server.close(); + reject(new Error('Callback server start timed out')); + } + }, 5_000); + }); +} + +/* ── Token exchange ──────────────────────────────────────────── */ + +async function exchangeCodeForToken( + tokenEndpoint: string, + params: { + code: string; + clientId: string; + redirectUri: string; + codeVerifier: string; + }, +): Promise { + const body = new URLSearchParams({ + grant_type: 'authorization_code', + code: params.code, + client_id: params.clientId, + redirect_uri: params.redirectUri, + code_verifier: params.codeVerifier, + }); + + const response = await fetch(tokenEndpoint, { + method: 'POST', + headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, + body: body.toString(), + }); + + if (!response.ok) { + throw new Error(`Token exchange failed: ${response.status} ${response.statusText}`); + } + + const data = await response.json(); + if (!data.access_token) { + throw new Error('Token response is missing access_token'); + } + + const token: McpOAuthToken = { + accessToken: data.access_token, + tokenType: data.token_type ?? 'Bearer', + scope: data.scope, + }; + + if (data.expires_in && typeof data.expires_in === 'number') { + token.expiresAt = Date.now() + data.expires_in * 1_000; + } + + if (data.refresh_token) { + token.refreshToken = data.refresh_token; + } + + return token; +} + +/* ── Utilities ───────────────────────────────────────────────── */ + +function escapeHtml(text: string): string { + return text + .replace(/&/g, '&') + .replace(//g, '>') + .replace(/"/g, '"'); +} diff --git a/src/main/services/mcpTokenStore.ts b/src/main/services/mcpTokenStore.ts new file mode 100644 index 0000000..e6f106b --- /dev/null +++ b/src/main/services/mcpTokenStore.ts @@ -0,0 +1,49 @@ +/** + * In-memory OAuth token store keyed by MCP server URL. + * Tokens are lost on app restart by design (phase 1). + */ + +export interface McpOAuthToken { + accessToken: string; + tokenType: string; + expiresAt?: number; + refreshToken?: string; + scope?: string; +} + +const tokens = new Map(); + +export function getStoredToken(serverUrl: string): McpOAuthToken | undefined { + const token = tokens.get(normalizeUrl(serverUrl)); + if (!token) { + return undefined; + } + + if (token.expiresAt && Date.now() >= token.expiresAt) { + tokens.delete(normalizeUrl(serverUrl)); + return undefined; + } + + return token; +} + +export function storeToken(serverUrl: string, token: McpOAuthToken): void { + tokens.set(normalizeUrl(serverUrl), token); +} + +export function clearToken(serverUrl: string): void { + tokens.delete(normalizeUrl(serverUrl)); +} + +export function clearAllTokens(): void { + tokens.clear(); +} + +function normalizeUrl(url: string): string { + try { + const parsed = new URL(url); + return parsed.origin + parsed.pathname.replace(/\/+$/, ''); + } catch { + return url.toLowerCase().replace(/\/+$/, ''); + } +} diff --git a/src/main/sessionToolingConfig.ts b/src/main/sessionToolingConfig.ts index bb4fc8f..d0d23fb 100644 --- a/src/main/sessionToolingConfig.ts +++ b/src/main/sessionToolingConfig.ts @@ -30,6 +30,7 @@ export function validateSessionToolingSelectionIds( export function buildRunTurnToolingConfig( tooling: WorkspaceToolingSettings, selection: SessionToolingSelection, + tokenLookup?: (serverUrl: string) => string | undefined, ): RunTurnToolingConfig | undefined { const mcpServersById = new Map( tooling.mcpServers.map((server) => [server.id, server]), @@ -68,7 +69,7 @@ export function buildRunTurnToolingConfig( tools: [...server.tools], timeoutMs: server.timeoutMs, url: server.url, - headers: server.headers ? { ...server.headers } : undefined, + headers: mergeAuthorizationHeader(server.url, server.headers, tokenLookup), }, ]; }); @@ -100,3 +101,19 @@ export function buildRunTurnToolingConfig( lspProfiles, }; } + +function mergeAuthorizationHeader( + serverUrl: string, + configHeaders: Record | undefined, + tokenLookup: ((serverUrl: string) => string | undefined) | undefined, +): Record | undefined { + const bearerToken = tokenLookup?.(serverUrl); + if (!bearerToken) { + return configHeaders ? { ...configHeaders } : undefined; + } + + return { + ...(configHeaders ?? {}), + Authorization: `Bearer ${bearerToken}`, + }; +} diff --git a/src/preload/index.ts b/src/preload/index.ts index 9b872b9..3f66352 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -40,6 +40,7 @@ const api: ElectronApi = { setSessionInteractionMode: (input) => ipcRenderer.invoke(ipcChannels.setSessionInteractionMode, input), dismissSessionPlanReview: (input) => ipcRenderer.invoke(ipcChannels.dismissSessionPlanReview, input), dismissSessionMcpAuth: (input) => ipcRenderer.invoke(ipcChannels.dismissSessionMcpAuth, input), + startSessionMcpAuth: (input) => ipcRenderer.invoke(ipcChannels.startSessionMcpAuth, input), updateSessionModelConfig: (input) => ipcRenderer.invoke(ipcChannels.updateSessionModelConfig, input), querySessions: (input) => ipcRenderer.invoke(ipcChannels.querySessions, input), diff --git a/src/renderer/App.tsx b/src/renderer/App.tsx index 41048b6..4bab4e8 100644 --- a/src/renderer/App.tsx +++ b/src/renderer/App.tsx @@ -262,6 +262,9 @@ export default function App() { onDismissMcpAuth={() => { void api.dismissSessionMcpAuth({ sessionId: selectedSession.id }); }} + onAuthenticateMcp={() => { + void api.startSessionMcpAuth({ sessionId: selectedSession.id }); + }} onUpdateSessionModelConfig={(config) => api.updateSessionModelConfig({ sessionId: selectedSession.id, diff --git a/src/renderer/components/ChatPane.tsx b/src/renderer/components/ChatPane.tsx index 1977425..54682e5 100644 --- a/src/renderer/components/ChatPane.tsx +++ b/src/renderer/components/ChatPane.tsx @@ -44,6 +44,7 @@ interface ChatPaneProps { onSetInteractionMode?: (mode: InteractionMode) => void; onDismissPlanReview?: () => void; onDismissMcpAuth?: () => void; + onAuthenticateMcp?: () => void; onUpdateSessionModelConfig?: (config: { model: string; reasoningEffort?: ReasoningEffort; @@ -66,6 +67,7 @@ export function ChatPane({ onSetInteractionMode, onDismissPlanReview, onDismissMcpAuth, + onAuthenticateMcp, onUpdateSessionModelConfig, onUpdateSessionTooling, onUpdateSessionApprovalSettings, @@ -145,6 +147,10 @@ export function ChatPane({ onDismissMcpAuth?.(); } + function handleAuthenticateMcp() { + onAuthenticateMcp?.(); + } + async function handleSessionModelConfigChange(config: { model: string; reasoningEffort?: ReasoningEffort; @@ -414,6 +420,7 @@ export function ChatPane({
diff --git a/src/renderer/components/chat/McpAuthBanner.tsx b/src/renderer/components/chat/McpAuthBanner.tsx index a4b0694..ef268aa 100644 --- a/src/renderer/components/chat/McpAuthBanner.tsx +++ b/src/renderer/components/chat/McpAuthBanner.tsx @@ -1,15 +1,21 @@ import { useCallback } from 'react'; -import { KeyRound, X } from 'lucide-react'; +import { KeyRound, Loader2, X } from 'lucide-react'; import type { PendingMcpAuthRecord } from '@shared/domain/mcpAuth'; export function McpAuthBanner({ mcpAuth, + onAuthenticate, onDismiss, }: { mcpAuth: PendingMcpAuthRecord; + onAuthenticate: () => void; onDismiss: () => void; }) { + const handleAuthenticate = useCallback(() => { + onAuthenticate(); + }, [onAuthenticate]); + const handleDismiss = useCallback(() => { onDismiss(); }, [onDismiss]); @@ -51,11 +57,30 @@ export function McpAuthBanner({

{mcpAuth.errorMessage}

)} -

- {isAuthenticating - ? 'Waiting for authentication to complete in the browser…' - : 'Authentication support for HTTP MCP servers is not yet available. Configure a static access token in the MCP server headers instead.'} -

+
+ + + {isAuthenticating + ? 'Waiting for consent in the browser…' + : 'Opens your browser for OAuth consent. Token is stored for this session only.'} + +
diff --git a/src/shared/contracts/channels.ts b/src/shared/contracts/channels.ts index 3e0d076..4f27b91 100644 --- a/src/shared/contracts/channels.ts +++ b/src/shared/contracts/channels.ts @@ -30,6 +30,7 @@ export const ipcChannels = { setSessionInteractionMode: 'sessions:set-interaction-mode', dismissSessionPlanReview: 'sessions:dismiss-plan-review', dismissSessionMcpAuth: 'sessions:dismiss-mcp-auth', + startSessionMcpAuth: 'sessions:start-mcp-auth', querySessions: 'sessions:query', updateSessionModelConfig: 'sessions:update-model-config', selectProject: 'selection:project', diff --git a/src/shared/contracts/ipc.ts b/src/shared/contracts/ipc.ts index 028b086..8d8b8c6 100644 --- a/src/shared/contracts/ipc.ts +++ b/src/shared/contracts/ipc.ts @@ -120,6 +120,10 @@ export interface DismissSessionMcpAuthInput { sessionId: string; } +export interface StartSessionMcpAuthInput { + sessionId: string; +} + export interface ElectronApi { describeSidecarCapabilities(): Promise; refreshSidecarCapabilities(): Promise; @@ -150,6 +154,7 @@ export interface ElectronApi { setSessionInteractionMode(input: SetSessionInteractionModeInput): Promise; dismissSessionPlanReview(input: DismissSessionPlanReviewInput): Promise; dismissSessionMcpAuth(input: DismissSessionMcpAuthInput): Promise; + startSessionMcpAuth(input: StartSessionMcpAuthInput): Promise; updateSessionModelConfig(input: UpdateSessionModelConfigInput): Promise; querySessions(input: QuerySessionsInput): Promise; selectProject(projectId?: string): Promise; diff --git a/tests/main/mcpTokenStore.test.ts b/tests/main/mcpTokenStore.test.ts new file mode 100644 index 0000000..6102fa6 --- /dev/null +++ b/tests/main/mcpTokenStore.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, test, beforeEach } from 'bun:test'; + +import { + getStoredToken, + storeToken, + clearToken, + clearAllTokens, + type McpOAuthToken, +} from '@main/services/mcpTokenStore'; + +describe('MCP OAuth token store', () => { + beforeEach(() => { + clearAllTokens(); + }); + + test('stores and retrieves tokens by server URL', () => { + const token: McpOAuthToken = { + accessToken: 'abc123', + tokenType: 'Bearer', + }; + + storeToken('https://mcp.example.com/api', token); + + expect(getStoredToken('https://mcp.example.com/api')).toEqual(token); + }); + + test('normalizes trailing slashes and case in server URLs', () => { + const token: McpOAuthToken = { + accessToken: 'xyz', + tokenType: 'Bearer', + }; + + storeToken('https://MCP.Example.com/api/', token); + + expect(getStoredToken('https://mcp.example.com/api')).toEqual(token); + expect(getStoredToken('https://mcp.example.com/api/')).toEqual(token); + }); + + test('returns undefined for unknown server URLs', () => { + expect(getStoredToken('https://unknown.example.com')).toBeUndefined(); + }); + + test('clears a token for a specific server', () => { + storeToken('https://a.example.com', { accessToken: 'a', tokenType: 'Bearer' }); + storeToken('https://b.example.com', { accessToken: 'b', tokenType: 'Bearer' }); + + clearToken('https://a.example.com'); + + expect(getStoredToken('https://a.example.com')).toBeUndefined(); + expect(getStoredToken('https://b.example.com')).toBeDefined(); + }); + + test('clears all stored tokens', () => { + storeToken('https://a.example.com', { accessToken: 'a', tokenType: 'Bearer' }); + storeToken('https://b.example.com', { accessToken: 'b', tokenType: 'Bearer' }); + + clearAllTokens(); + + expect(getStoredToken('https://a.example.com')).toBeUndefined(); + expect(getStoredToken('https://b.example.com')).toBeUndefined(); + }); + + test('returns undefined for expired tokens', () => { + const token: McpOAuthToken = { + accessToken: 'expired', + tokenType: 'Bearer', + expiresAt: Date.now() - 1_000, + }; + + storeToken('https://mcp.example.com', token); + + expect(getStoredToken('https://mcp.example.com')).toBeUndefined(); + }); + + test('returns valid tokens that have not expired', () => { + const token: McpOAuthToken = { + accessToken: 'valid', + tokenType: 'Bearer', + expiresAt: Date.now() + 60_000, + }; + + storeToken('https://mcp.example.com', token); + + expect(getStoredToken('https://mcp.example.com')).toEqual(token); + }); +}); diff --git a/tests/main/sessionToolingConfig.test.ts b/tests/main/sessionToolingConfig.test.ts index 3b69b64..aee53e4 100644 --- a/tests/main/sessionToolingConfig.test.ts +++ b/tests/main/sessionToolingConfig.test.ts @@ -117,4 +117,59 @@ describe('session tooling config helpers', () => { }), ).toBeUndefined(); }); + + test('injects OAuth token as Authorization header for remote MCP servers', () => { + const tokenLookup = (url: string) => + url === 'https://example.com/mcp' ? 'oauth-access-token' : undefined; + + const config = buildRunTurnToolingConfig( + TOOLING, + { enabledMcpServerIds: ['mcp-remote'], enabledLspProfileIds: [] }, + tokenLookup, + ); + + expect(config?.mcpServers[0]).toMatchObject({ + id: 'mcp-remote', + headers: { Authorization: 'Bearer oauth-access-token' }, + }); + }); + + test('preserves existing headers when injecting OAuth token', () => { + const toolingWithHeaders: WorkspaceToolingSettings = { + ...TOOLING, + mcpServers: [ + { + id: 'mcp-custom', + name: 'Custom MCP', + transport: 'http', + url: 'https://custom.example.com/mcp', + headers: { 'X-Custom': 'value' }, + tools: [], + createdAt: TIMESTAMP, + updatedAt: TIMESTAMP, + }, + ], + }; + + const config = buildRunTurnToolingConfig( + toolingWithHeaders, + { enabledMcpServerIds: ['mcp-custom'], enabledLspProfileIds: [] }, + () => 'my-token', + ); + + expect(config?.mcpServers[0].headers).toEqual({ + 'X-Custom': 'value', + Authorization: 'Bearer my-token', + }); + }); + + test('does not inject Authorization header when no token is available', () => { + const config = buildRunTurnToolingConfig( + TOOLING, + { enabledMcpServerIds: ['mcp-remote'], enabledLspProfileIds: [] }, + () => undefined, + ); + + expect(config?.mcpServers[0].headers).toEqual({ Authorization: 'Bearer token' }); + }); });