fix(mcp): drain OAuth refreshes during shutdown

packages/agent-core-v2/src/mcpCore/oauth/service.ts

@@ -1,6 +1,5 @@
 import { auth, type OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js';
 
-import { Disposable } from '#/_base/di/lifecycle';
 import type { ILogger as Logger } from '#/_base/log/log';
 import { ErrorCodes, Error2, isError2 } from '#/errors';
 
@@ -112,7 +111,7 @@ const defaultScheduler: McpOAuthScheduler = {
   },
 };
 
-export class McpOAuthService extends Disposable {
+export class McpOAuthService {
   private readonly store: McpOAuthStore;
   private readonly clientLabel: string | undefined;
   private readonly resolveClientName: (() => string | undefined) | undefined;
@@ -123,19 +122,19 @@ export class McpOAuthService extends Disposable {
   private readonly refreshes = new Map<string, Promise<void>>();
   private readonly refreshTimers = new Map<string, McpOAuthScheduledTask>();
   private readonly activeAuthorizations = new Map<string, Promise<SharedAuthorizationFlow>>();
+  private shuttingDown = false;
+  private shutdownPromise: Promise<void> | undefined;
 
   constructor(options: McpOAuthServiceOptions) {
-    super();
     this.store = options.store;
     this.clientLabel = options.clientLabel;
     this.resolveClientName = options.resolveClientName;
     this.log = options.log ?? defaultLog;
     this.scheduler = options.scheduler ?? defaultScheduler;
-    this._register({
-      dispose: () => {
-        void this.shutdown();
-      },
-    });
+  }
+
+  dispose(): Promise<void> {
+    return this.shutdown();
   }
 
   /** Returns the cached provider for `serverName` + `serverUrl`, constructing it on first use. */
@@ -196,6 +195,9 @@ export class McpOAuthService extends Disposable {
     const storeKey = mcpOAuthStoreKey(serverName, serverUrl);
     const existing = this.refreshes.get(storeKey);
     if (existing !== undefined) return existing;
+    if (this.shuttingDown) {
+      throw new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'MCP OAuth service is shutting down');
+    }
     const task = this.refreshNow(serverName, serverUrl).finally(() => {
       this.refreshes.delete(storeKey);
     });
@@ -237,21 +239,33 @@ export class McpOAuthService extends Disposable {
 
   /**
    * Release everything the service owns: pending proactive-refresh timers,
-   * in-flight interactive flows (closing their callback listeners), event
-   * listeners, and cached providers. Idempotent.
+   * in-flight refreshes and interactive flows (closing their callback
+   * listeners), event listeners, and cached providers. Idempotent.
    */
-  async shutdown(): Promise<void> {
+  shutdown(): Promise<void> {
+    if (this.shutdownPromise !== undefined) return this.shutdownPromise;
+    this.shuttingDown = true;
     this.stopProactiveRefresh();
-    const inFlight = [...this.activeAuthorizations.values()];
+    const authorizations = [...this.activeAuthorizations.values()];
+    const refreshes = [...this.refreshes.values()];
     this.activeAuthorizations.clear();
-    await Promise.all(
-      inFlight.map(async (started) => {
-        const flow = await started.catch(() => undefined);
-        await flow?.cancelUnderlying();
-      }),
-    );
-    this.listeners.clear();
-    this.providers.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),
+        ]);
+      } finally {
+        this.listeners.clear();
+        this.providers.clear();
+      }
+    })();
+    return this.shutdownPromise;
   }
 
   /**
@@ -270,6 +284,9 @@ export class McpOAuthService extends Disposable {
     serverUrl: string | URL,
     options: BeginAuthorizationOptions = {},
   ): Promise<BeginAuthorizationResult> {
+    if (this.shuttingDown) {
+      throw new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'MCP OAuth service is shutting down');
+    }
     const storeKey = mcpOAuthStoreKey(serverName, serverUrl);
     const inFlight = this.activeAuthorizations.get(storeKey);
     if (inFlight !== undefined) {
@@ -503,6 +520,7 @@ export class McpOAuthService extends Disposable {
   }
 
   private scheduleRefresh(serverName: string, serverUrl: string | URL, expiresAt: number): void {
+    if (this.shuttingDown) return;
     const canonicalUrl = canonicalMcpOAuthResource(serverUrl);
     const storeKey = mcpOAuthStoreKey(serverName, canonicalUrl);
     this.cancelScheduledRefresh(serverName, canonicalUrl);

packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts

@@ -163,7 +163,7 @@ describe('McpManagementService', () => {
 
   afterEach(async () => {
     disposables.dispose();
-    oauth.dispose();
+    await oauth.dispose();
     vi.unstubAllEnvs();
     await Promise.all(httpServers.map((server) => server.close()));
     await Promise.all(tempDirs.map((dir) => rm(dir, { recursive: true, force: true })));

packages/agent-core-v2/test/mcpCore/oauth/service.test.ts

@@ -141,6 +141,50 @@ function authServerState(authServerUrl: string) {
   };
 }
 
+async function blockedRefreshFixture(): Promise<{
+  readonly fixture: Fixture;
+  readonly authServer: FakeAuthServer;
+  readonly writeStarted: Promise<void>;
+  readonly releaseWrite: () => void;
+}> {
+  const memory = createMemoryMcpOAuthStore();
+  let signalWriteStarted: () => void = () => undefined;
+  const writeStarted = new Promise<void>((resolve) => {
+    signalWriteStarted = resolve;
+  });
+  let releaseWrite: () => void = () => undefined;
+  const writeReleased = new Promise<void>((resolve) => {
+    releaseWrite = resolve;
+  });
+  const store: McpOAuthStore = {
+    ...memory,
+    async write(key: string, value: unknown): Promise<void> {
+      const accessToken =
+        typeof value === 'object' && value !== null
+          ? (value as { readonly access_token?: unknown }).access_token
+          : undefined;
+      if (accessToken === 'fresh-token') {
+        signalWriteStarted();
+        await writeReleased;
+      }
+      await memory.write(key, value);
+    },
+  };
+  const fixture = makeFixture(store);
+  cleanups.push(() => fixture.service.dispose());
+  const authServer = await startFakeAuthServer({ refreshExpiresIn: 60 });
+  const provider = await readyProvider(fixture);
+  const state = authServerState(authServer.url);
+  await provider.saveDiscoveryState(state.discovery);
+  await provider.saveClientInformation(state.client);
+  await provider.saveTokens({
+    access_token: 'stale-access-token',
+    refresh_token: 'stale-refresh-token',
+    token_type: 'Bearer',
+  });
+  return { fixture, authServer, writeStarted, releaseWrite };
+}
+
 async function deliverCallback(flow: BeginAuthorizationResult): Promise<void> {
   const redirectUri = flow.authorizationUrl.searchParams.get('redirect_uri');
   const state = flow.authorizationUrl.searchParams.get('state');
@@ -650,9 +694,7 @@ describe('McpOAuthService proactive refresh scheduling', () => {
 
   it('waits another midpoint after refreshing into another 60-second grant', async () => {
     const fixture = makeFixture();
-    cleanups.push(() => {
-      fixture.service.dispose();
-    });
+    cleanups.push(() => fixture.service.dispose());
     const authServer = await startFakeAuthServer({ refreshExpiresIn: 60 });
 
     const provider = await readyProvider(fixture);
@@ -723,6 +765,37 @@ describe('McpOAuthService proactive refresh scheduling', () => {
 });
 
 describe('McpOAuthService shutdown', () => {
+  it('keeps shutdown pending while a token refresh is in flight', async () => {
+    const { fixture, writeStarted, releaseWrite } = await blockedRefreshFixture();
+    const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL);
+    await writeStarted;
+
+    const shutdown = fixture.service.shutdown();
+    let shutdownSettled = false;
+    void shutdown.then(() => {
+      shutdownSettled = true;
+    });
+    await Promise.resolve();
+    const pendingWhileRefreshInFlight = !shutdownSettled;
+
+    releaseWrite();
+    await Promise.all([refresh, shutdown]);
+    expect(pendingWhileRefreshInFlight).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);
+    await writeStarted;
+
+    const shutdown = fixture.service.shutdown();
+    releaseWrite();
+    await Promise.all([refresh, shutdown]);
+    await fixture.scheduler.advanceBy(30_000);
+
+    expect(authServer.counts.refresh).toBe(1);
+  }, 15000);
+
   it('cancels active flows on shutdown', async () => {
     const fixture = makeFixture();
     cleanups.push(() => fixture.service.dispose());

packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts

@@ -93,7 +93,7 @@ describe('WorkspaceMcpService', () => {
   afterEach(async () => {
     vi.restoreAllMocks();
     await manager?.shutdown();
-    oauthService.dispose();
+    await oauthService.dispose();
     disposables.dispose();
     await rm(cwd, { recursive: true, force: true });
   });