fix(agent-core-v2): bound MCP OAuth auth-server requests and the shutdown drain

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<string, McpOAuthClientProvider>();
   private readonly listeners = new Set<McpOAuthEventListener>();
   private readonly refreshes = new Map<string, Promise<void>>();
@@ -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<void> {
@@ -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<void> {
+    return new Promise<void>((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<typeof fetch>[0], init?: Parameters<typeof fetch>[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(

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<void>((resolve) => httpServer.listen(0, '127.0.0.1', resolve));
+  cleanups.push(
+    () =>
+      new Promise<void>((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<void>;
@@ -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<boolean>((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);