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>
2026-07-30 08:28:48 +02:00 · 2026-03-27 19:15:37 +01:00
parent f9757d5ce2
commit ac48aa58e0
13 changed files with 635 additions and 8 deletions
@@ -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<AppServiceEvents> {
    return this.persistAndBroadcast(workspace);
  }

+  async startSessionMcpAuth(sessionId: string): Promise<WorkspaceState> {
+    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<AppServiceEvents> {
    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<boolean> {
@@ -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) =>
@@ -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<McpOAuthFlowResult> {
+  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<string> {
+  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<AuthServerMetadata> {
+  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<string> {
+  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<string>;
+  close: () => void;
+}
+
+function startCallbackServer(): Promise<CallbackServerHandle> {
+  return new Promise((resolve, reject) => {
+    let settled = false;
+    let callbackResolve: (code: string) => void;
+    let callbackReject: (err: Error) => void;
+
+    const callbackPromise = new Promise<string>((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('<html><body><h2>Authentication successful</h2><p>You can close this tab.</p></body></html>');
+        callbackResolve(code);
+      } else {
+        const msg = errorDescription ?? error ?? 'Unknown error';
+        res.end(`<html><body><h2>Authentication failed</h2><p>${escapeHtml(msg)}</p></body></html>`);
+        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<McpOAuthToken> {
+  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, '&amp;')
+    .replace(/</g, '&lt;')
+    .replace(/>/g, '&gt;')
+    .replace(/"/g, '&quot;');
+}
@@ -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<string, McpOAuthToken>();
+
+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(/\/+$/, '');
+  }
+}
@@ -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<string, McpServerDefinition>(
    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<string, string> | undefined,
+  tokenLookup: ((serverUrl: string) => string | undefined) | undefined,
+): Record<string, string> | undefined {
+  const bearerToken = tokenLookup?.(serverUrl);
+  if (!bearerToken) {
+    return configHeaders ? { ...configHeaders } : undefined;
+  }
+
+  return {
+    ...(configHeaders ?? {}),
+    Authorization: `Bearer ${bearerToken}`,
+  };
+}
@@ -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),
@@ -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,
@@ -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({
            <div className="mb-3">
              <McpAuthBanner
                mcpAuth={pendingMcpAuth}
+                onAuthenticate={handleAuthenticateMcp}
                onDismiss={handleDismissMcpAuth}
              />
            </div>
@@ -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({
            <p className="mt-2 text-[12px] text-red-400">{mcpAuth.errorMessage}</p>
          )}

-          <p className="mt-3 text-[12px] leading-relaxed text-zinc-400">
-            {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.'}
-          </p>
+          <div className="mt-3 flex items-center gap-3">
+            <button
+              className="inline-flex items-center gap-1.5 rounded-lg bg-amber-500/20 px-3 py-1.5 text-[12px] font-medium text-amber-200 transition hover:bg-amber-500/30 disabled:opacity-50"
+              disabled={isAuthenticating}
+              onClick={handleAuthenticate}
+              type="button"
+            >
+              {isAuthenticating ? (
+                <>
+                  <Loader2 className="size-3.5 animate-spin" />
+                  Authenticating…
+                </>
+              ) : hasFailed ? (
+                'Retry authentication'
+              ) : (
+                'Authenticate in browser'
+              )}
+            </button>
+            <span className="text-[11px] text-zinc-500">
+              {isAuthenticating
+                ? 'Waiting for consent in the browser…'
+                : 'Opens your browser for OAuth consent. Token is stored for this session only.'}
+            </span>
+          </div>
        </div>
      </div>
    </div>
@@ -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',
@@ -120,6 +120,10 @@ export interface DismissSessionMcpAuthInput {
  sessionId: string;
 }

+export interface StartSessionMcpAuthInput {
+  sessionId: string;
+}
+
 export interface ElectronApi {
  describeSidecarCapabilities(): Promise<SidecarCapabilities>;
  refreshSidecarCapabilities(): Promise<SidecarCapabilities>;
@@ -150,6 +154,7 @@ export interface ElectronApi {
  setSessionInteractionMode(input: SetSessionInteractionModeInput): Promise<WorkspaceState>;
  dismissSessionPlanReview(input: DismissSessionPlanReviewInput): Promise<WorkspaceState>;
  dismissSessionMcpAuth(input: DismissSessionMcpAuthInput): Promise<WorkspaceState>;
+  startSessionMcpAuth(input: StartSessionMcpAuthInput): Promise<WorkspaceState>;
  updateSessionModelConfig(input: UpdateSessionModelConfigInput): Promise<WorkspaceState>;
  querySessions(input: QuerySessionsInput): Promise<SessionQueryResult[]>;
  selectProject(projectId?: string): Promise<WorkspaceState>;
@@ -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);
+  });
+});
@@ -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' });
+  });
 });