diff --git a/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts b/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts index f190c4dd7..cb4141f86 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts @@ -8,10 +8,24 @@ export interface CallbackResult { export interface CallbackServer { readonly redirectUri: string; + /** + * Resolves with the OAuth callback payload, or rejects when: + * - `signal` aborts → AbortError + * - `timeoutMs` elapses → Error('OAuth callback timed out') + * - the user's authorization server returns an error → Error('OAuth error: ') + * - `close()` is called → OAuthCallbackClosedError + */ waitForCode(opts: { signal?: AbortSignal; timeoutMs?: number }): Promise; close(): Promise; } +export class OAuthCallbackClosedError extends Error { + constructor() { + super('OAuth callback listener closed'); + this.name = 'OAuthCallbackClosedError'; + } +} + const SUCCESS_HTML = 'Authorized' + '' + @@ -29,12 +43,29 @@ const ERROR_HTML = export async function startCallbackServer(): Promise { let resolveCode: ((value: CallbackResult) => void) | undefined; let rejectCode: ((reason: Error) => void) | undefined; - let settled = false; + let cleanupWait: (() => void) | undefined; + let outcome: + | { readonly status: 'pending' } + | { readonly status: 'resolved'; readonly value: CallbackResult } + | { readonly status: 'rejected'; readonly reason: Error } = { status: 'pending' }; - const settle = (fn: () => void) => { - if (settled) return; - settled = true; - fn(); + const settle = ( + next: + | { readonly status: 'resolved'; readonly value: CallbackResult } + | { readonly status: 'rejected'; readonly reason: Error }, + ) => { + if (outcome.status !== 'pending') return; + outcome = next; + cleanupWait?.(); + cleanupWait = undefined; + if (next.status === 'resolved') { + resolveCode?.(next.value); + } else { + rejectCode?.(next.reason); + } + resolveCode = undefined; + rejectCode = undefined; + void closeServer(); }; const server: Server = createServer((req, res) => { @@ -61,26 +92,26 @@ export async function startCallbackServer(): Promise { if (errorParam !== null) { const description = url.searchParams.get('error_description') ?? ''; res.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }).end(ERROR_HTML); - settle(() => { - rejectCode?.( - new Error(`OAuth error: ${errorParam}${description ? ` — ${description}` : ''}`), - ); + settle({ + status: 'rejected', + reason: new Error( + `OAuth error: ${errorParam}${description ? ` — ${description}` : ''}`, + ), }); return; } const code = url.searchParams.get('code'); if (code === null || code.length === 0) { res.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }).end(ERROR_HTML); - settle(() => { - rejectCode?.(new Error('OAuth callback missing authorization code')); + settle({ + status: 'rejected', + reason: new Error('OAuth callback missing authorization code'), }); return; } const state = url.searchParams.get('state') ?? undefined; res.writeHead(200, { 'content-type': 'text/html; charset=utf-8' }).end(SUCCESS_HTML); - settle(() => { - resolveCode?.({ code, state }); - }); + settle({ status: 'resolved', value: { code, state } }); } await new Promise((resolve, reject) => { @@ -93,44 +124,49 @@ export async function startCallbackServer(): Promise { const port = (server.address() as AddressInfo).port; const redirectUri = `http://127.0.0.1:${port}/callback`; - let closed = false; - const close = async () => { - if (closed) return; - closed = true; - await new Promise((resolve) => { + let closeServerPromise: Promise | undefined; + const closeServer = (): Promise => { + closeServerPromise ??= new Promise((resolve) => { server.close(() => { resolve(); }); }); + return closeServerPromise; + }; + const close = async () => { + settle({ status: 'rejected', reason: new OAuthCallbackClosedError() }); + await closeServer(); }; const waitForCode: CallbackServer['waitForCode'] = ({ signal, timeoutMs } = {}) => { return new Promise((resolve, reject) => { + if (outcome.status === 'resolved') { + resolve(outcome.value); + return; + } + if (outcome.status === 'rejected') { + reject(outcome.reason); + return; + } + let timer: NodeJS.Timeout | undefined; const onAbort = () => { - settle(() => - rejectCode?.( + settle({ + status: 'rejected', + reason: signal?.reason instanceof Error ? signal.reason : new Error('OAuth flow aborted'), - ), - ); + }); }; const cleanup = () => { if (timer !== undefined) clearTimeout(timer); signal?.removeEventListener('abort', onAbort); }; - resolveCode = (value) => { - cleanup(); - void close(); - resolve(value); - }; - rejectCode = (reason) => { - cleanup(); - void close(); - reject(reason); - }; + cleanupWait = cleanup; + resolveCode = resolve; + rejectCode = reject; if (timeoutMs !== undefined) { timer = setTimeout(() => { - settle(() => rejectCode?.(new Error('OAuth callback timed out'))); + settle({ status: 'rejected', reason: new Error('OAuth callback timed out') }); }, timeoutMs); } if (signal !== undefined) { diff --git a/packages/agent-core-v2/test/mcpCore/oauth/callback-server.test.ts b/packages/agent-core-v2/test/mcpCore/oauth/callback-server.test.ts new file mode 100644 index 000000000..5ad70c11f --- /dev/null +++ b/packages/agent-core-v2/test/mcpCore/oauth/callback-server.test.ts @@ -0,0 +1,134 @@ +import { createServer as createHttpServer, type Server as HttpServer } from 'node:http'; +import type { AddressInfo as HttpAddress } from 'node:net'; + +import { afterEach, describe, expect, it } from 'vitest'; + +import { + type CallbackServer, + OAuthCallbackClosedError, + startCallbackServer, +} from '#/mcpCore/oauth/callback-server'; +import { McpOAuthService } from '#/mcpCore/oauth/service'; + +import { createMemoryMcpOAuthStore } from '../stubs'; + +const cleanups: Array<() => Promise | void> = []; +afterEach(async () => { + while (cleanups.length > 0) { + await cleanups.pop()?.(); + } +}); + +function trackServer(server: CallbackServer): void { + cleanups.push(() => server.close()); +} + +async function startRegistrationServer(): Promise<{ readonly url: string }> { + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.method !== 'POST' || req.url !== '/register') { + res.writeHead(404).end(); + return; + } + let body = ''; + req.on('data', (chunk: Buffer) => { + body += chunk.toString('utf-8'); + }); + req.on('end', () => { + const metadata = JSON.parse(body) as Record; + res.writeHead(201, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ ...metadata, client_id: 'test-client' })); + }); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const port = (httpServer.address() as HttpAddress).port; + return { url: `http://127.0.0.1:${port}` }; +} + +describe('startCallbackServer', () => { + it('resolves a late waitForCode with a callback that arrived before it', async () => { + const server = await startCallbackServer(); + trackServer(server); + + const response = await fetch(`${server.redirectUri}?code=early-code&state=early-state`); + expect(response.status).toBe(200); + + await expect(server.waitForCode({ timeoutMs: 10_000 })).resolves.toEqual({ + code: 'early-code', + state: 'early-state', + }); + }); + + it('delivers the callback payload to a pending wait', async () => { + const server = await startCallbackServer(); + trackServer(server); + const pending = server.waitForCode({ timeoutMs: 10_000 }); + + await fetch(`${server.redirectUri}?code=code-1&state=state-1`); + + await expect(pending).resolves.toEqual({ code: 'code-1', state: 'state-1' }); + }); + + it('rejects a pending wait with a closed error when explicitly closed', async () => { + const server = await startCallbackServer(); + trackServer(server); + const pending = server.waitForCode({ timeoutMs: 10_000 }); + const rejection = expect(pending).rejects.toBeInstanceOf(OAuthCallbackClosedError); + + await server.close(); + + await rejection; + }); + + it('rejects a late waitForCode after close', async () => { + const server = await startCallbackServer(); + trackServer(server); + + await server.close(); + + await expect(server.waitForCode({ timeoutMs: 10_000 })).rejects.toBeInstanceOf( + OAuthCallbackClosedError, + ); + }); +}); + +describe('McpOAuthService cancellation', () => { + it('rejects an in-flight completion when the authorization flow is cancelled', async () => { + const service = new McpOAuthService({ store: createMemoryMcpOAuthStore() }); + cleanups.push(() => service.dispose()); + const registrationServer = await startRegistrationServer(); + const provider = service.getProvider('example', 'https://mcp.example.test/rpc'); + await provider.ready; + await provider.saveDiscoveryState({ + authorizationServerUrl: registrationServer.url, + authorizationServerMetadata: { + issuer: registrationServer.url, + authorization_endpoint: `${registrationServer.url}/authorize`, + token_endpoint: `${registrationServer.url}/token`, + registration_endpoint: `${registrationServer.url}/register`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code'], + token_endpoint_auth_methods_supported: ['none'], + }, + }); + + const flow = await service.beginAuthorization('example', 'https://mcp.example.test/rpc'); + const completion = flow.complete({ timeoutMs: 10_000 }); + const rejection = expect(completion).rejects.toThrow('OAuth callback listener closed'); + + await flow.cancel(); + + await rejection; + }, 15000); +});