diff --git a/packages/agent-core-v2/src/mcpCore/oauth/service.ts b/packages/agent-core-v2/src/mcpCore/oauth/service.ts index 905acc841..c4493ec82 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/service.ts +++ b/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>(); private readonly refreshTimers = new Map(); private readonly activeAuthorizations = new Map>(); + private shuttingDown = false; + private shutdownPromise: Promise | 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 { + 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 { + shutdown(): Promise { + 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 { + 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); diff --git a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts index 1b34e9bae..1c1f535bc 100644 --- a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts +++ b/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 }))); 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 3304f7fc9..4a7f8703d 100644 --- a/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts +++ b/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; + readonly releaseWrite: () => void; +}> { + const memory = createMemoryMcpOAuthStore(); + let signalWriteStarted: () => void = () => undefined; + const writeStarted = new Promise((resolve) => { + signalWriteStarted = resolve; + }); + let releaseWrite: () => void = () => undefined; + const writeReleased = new Promise((resolve) => { + releaseWrite = resolve; + }); + const store: McpOAuthStore = { + ...memory, + async write(key: string, value: unknown): Promise { + 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 { 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()); diff --git a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts index 3fe50b984..62ca5f458 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts +++ b/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 }); });