diff --git a/packages/agent-core-v2/src/mcpCore/oauth/service.ts b/packages/agent-core-v2/src/mcpCore/oauth/service.ts index 6f3a277ff..35e086997 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/service.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/service.ts @@ -26,6 +26,10 @@ export interface McpOAuthServiceOptions { readonly resolveClientName?: () => string | undefined; readonly log?: Logger; readonly scheduler?: McpOAuthScheduler; + /** Per-request bound for OAuth-flow HTTP (discovery, registration, grants). */ + readonly authRequestTimeoutMs?: number; + /** Upper bound for awaiting in-flight flows and refreshes during shutdown. */ + readonly shutdownDrainTimeoutMs?: number; } export interface McpOAuthScheduledTask { @@ -101,6 +105,8 @@ export interface McpOAuthTokenState { const REFRESH_AHEAD_MS = 120_000; const MAX_TIMER_DELAY_MS = 0x7fffffff; +const DEFAULT_AUTH_REQUEST_TIMEOUT_MS = 30_000; +const DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS = 30_000; const defaultScheduler: McpOAuthScheduler = { now: () => Date.now(), @@ -117,6 +123,8 @@ export class McpOAuthService { private readonly resolveClientName: (() => string | undefined) | undefined; private readonly log: Logger; private readonly scheduler: McpOAuthScheduler; + private readonly authRequestTimeoutMs: number; + private readonly shutdownDrainTimeoutMs: number; private readonly providers = new Map(); private readonly listeners = new Set(); private readonly refreshes = new Map>(); @@ -132,6 +140,8 @@ export class McpOAuthService { this.resolveClientName = options.resolveClientName; this.log = options.log ?? defaultLog; this.scheduler = options.scheduler ?? defaultScheduler; + this.authRequestTimeoutMs = options.authRequestTimeoutMs ?? DEFAULT_AUTH_REQUEST_TIMEOUT_MS; + this.shutdownDrainTimeoutMs = options.shutdownDrainTimeoutMs ?? DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS; } dispose(): Promise { @@ -264,15 +274,18 @@ export class McpOAuthService { this.activeAuthorizations.clear(); this.shutdownPromise = (async () => { try { - await Promise.all([ - Promise.all( - authorizations.map(async (started) => { - const flow = await started.catch(() => undefined); - await flow?.cancelUnderlying(); - }), - ), - Promise.allSettled(refreshes), - Promise.allSettled(backgroundTasks), + await Promise.race([ + Promise.all([ + Promise.all( + authorizations.map(async (started) => { + const flow = await started.catch(() => undefined); + await flow?.cancelUnderlying(); + }), + ), + Promise.allSettled(refreshes), + Promise.allSettled(backgroundTasks), + ]), + this.drainDeadline(), ]); } finally { this.listeners.clear(); @@ -282,6 +295,28 @@ export class McpOAuthService { return this.shutdownPromise; } + private drainDeadline(): Promise { + return new Promise((resolve) => { + this.scheduler.schedule(this.shutdownDrainTimeoutMs, () => { + this.log.warn('mcp oauth shutdown drain timed out; continuing teardown'); + resolve(); + }); + }); + } + + private authFetch(provider: McpOAuthClientProvider): typeof fetch { + const fetchFn = provider.createOAuthFetch(); + const timeoutMs = this.authRequestTimeoutMs; + return (async (input: Parameters[0], init?: Parameters[1]) => { + const timeout = AbortSignal.timeout(timeoutMs); + const signal = + init?.signal === undefined || init.signal === null + ? timeout + : AbortSignal.any([init.signal, timeout]); + return fetchFn(input, { ...init, signal }); + }) as typeof fetch; + } + /** * Drive the SDK `auth()` orchestrator far enough to surface an * authorization URL. The caller is responsible for displaying the URL @@ -384,7 +419,7 @@ export class McpOAuthService { try { const result = await auth(provider as OAuthClientProvider, { serverUrl, - fetchFn: provider.createOAuthFetch(), + fetchFn: this.authFetch(provider), }); if (result !== 'REDIRECT') { await callbackServer.close(); @@ -544,7 +579,7 @@ export class McpOAuthService { try { const result = await auth(provider as OAuthClientProvider, { serverUrl, - fetchFn: provider.createOAuthFetch(), + fetchFn: this.authFetch(provider), }); if (result !== 'AUTHORIZED') { throw new Error2( diff --git a/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts b/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts index 049cd4c5f..34da0d9a5 100644 --- a/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts +++ b/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts @@ -30,10 +30,13 @@ interface Fixture { readonly scheduler: ManualMcpOAuthScheduler; } -function makeFixture(store: McpOAuthStore = createMemoryMcpOAuthStore()): Fixture { +function makeFixture( + store: McpOAuthStore = createMemoryMcpOAuthStore(), + options: { readonly authRequestTimeoutMs?: number; readonly shutdownDrainTimeoutMs?: number } = {}, +): Fixture { const events: McpOAuthEvent[] = []; const scheduler = new ManualMcpOAuthScheduler(1_000_000); - const service = new McpOAuthService({ store, scheduler }); + const service = new McpOAuthService({ store, scheduler, ...options }); service.onEvent((event) => events.push(event)); return { service, store, events, scheduler }; } @@ -118,6 +121,28 @@ async function startFakeAuthServer( return { url: `http://127.0.0.1:${port}`, counts }; } +async function startHangingServer(): Promise<{ readonly url: string; readonly counts: { requests: number } }> { + const counts = { requests: 0 }; + const httpServer: HttpServer = createHttpServer(() => { + counts.requests += 1; + }); + 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}`, counts }; +} + function authServerState(authServerUrl: string) { return { discovery: { @@ -142,7 +167,9 @@ function authServerState(authServerUrl: string) { }; } -async function blockedRefreshFixture(): Promise<{ +async function blockedRefreshFixture(options?: { + readonly shutdownDrainTimeoutMs?: number; +}): Promise<{ readonly fixture: Fixture; readonly authServer: FakeAuthServer; readonly writeStarted: Promise; @@ -171,7 +198,7 @@ async function blockedRefreshFixture(): Promise<{ await memory.write(key, value); }, }; - const fixture = makeFixture(store); + const fixture = makeFixture(store, options ?? {}); cleanups.push(() => fixture.service.dispose()); const authServer = await startFakeAuthServer({ refreshExpiresIn: 60 }); const provider = await readyProvider(fixture); @@ -351,6 +378,22 @@ describe('McpOAuthService single-flight refresh', () => { expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(2); }, 15000); + it('bounds a hung authorization-server request by the configured request timeout', async () => { + const hanging = await startHangingServer(); + const fixture = makeFixture(createMemoryMcpOAuthStore(), { authRequestTimeoutMs: 50 }); + cleanups.push(() => fixture.service.dispose()); + const provider = fixture.service.getProvider(SERVER_NAME, hanging.url); + await provider.ready; + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + + await expect(fixture.service.refresh(SERVER_NAME, hanging.url)).rejects.toThrow(); + expect(hanging.counts.requests).toBeGreaterThan(0); + }); + it('rejects when no refresh token is stored', async () => { const fixture = makeFixture(); cleanups.push(() => fixture.service.dispose()); @@ -870,6 +913,25 @@ describe('McpOAuthService shutdown', () => { expect(pendingWhileRefreshInFlight).toBe(true); }, 15000); + it('caps the shutdown drain when an in-flight refresh outlives the drain timeout', async () => { + const { fixture, writeStarted, releaseWrite } = await blockedRefreshFixture({ + shutdownDrainTimeoutMs: 50, + }); + const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL); + await writeStarted; + + const shutdown = fixture.service.shutdown(); + await fixture.scheduler.advanceBy(60); + const settledBeforeRelease = await Promise.race([ + shutdown.then(() => true), + new Promise((resolve) => setTimeout(() => resolve(false), 1_000)), + ]); + releaseWrite(); + await Promise.all([refresh, shutdown]); + + expect(settledBeforeRelease).toBe(true); + }, 15000); + it('prevents a completing refresh from scheduling work after shutdown', async () => { const { fixture, authServer, writeStarted, releaseWrite } = await blockedRefreshFixture(); const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL);