diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index 1b722dc1a7..41102a61ad 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -451,8 +451,13 @@ export class MCPConnectionFactory { } } - /** Resolves OBO tokens when the server config specifies obo, returns null otherwise */ - protected async getOboTokens(): Promise { + /** + * Resolves OBO tokens when the server config specifies obo, returns null otherwise. + * + * @param forceRefresh Bypasses the OBO token cache; used when the downstream server + * has rejected the credential we currently hold. + */ + protected async getOboTokens(forceRefresh = false): Promise { const oboConfig = this.serverConfig.obo; if (!oboConfig || !this.oboTokenResolver || !this.user) { return null; @@ -488,6 +493,7 @@ export class MCPConnectionFactory { this.oboTokenResolver, this.upstreamTokenProvider, this.oboIdentityContext, + forceRefresh, ); } @@ -543,6 +549,8 @@ export class MCPConnectionFactory { let cleanupOAuthHandlers: (() => void) | null = null; if (this.useOAuth && !this.usesObo) { cleanupOAuthHandlers = this.handleOAuthEvents(connection); + } else if (this.usesObo) { + cleanupOAuthHandlers = this.handleOboEvents(connection); } else { const nonOAuthHandler = () => { logger.info( @@ -1109,6 +1117,81 @@ export class MCPConnectionFactory { return expiresAt > Date.now() ? expiresAt : undefined; } + /** + * Sets up the authentication handler for an OBO connection. + * + * An OBO server rejecting our bearer is recoverable in a way the non-OAuth + * fallback cannot express: the downstream credential is minted from the user's + * live upstream session, so a fresh exchange produces a working token without + * any interactive step. What it needs is that live session, which only exists + * while the request that created this connection is still running. Past that + * point `upstreamTokenProvider` closes over a finished request, so a cached + * connection stops reconnecting instead and its next borrower rebuilds it + * against a live provider. + */ + protected handleOboEvents(connection: MCPConnection): () => void { + let refreshAttempted = false; + + const oboHandler = async (): Promise => { + if (this.connectionReady) { + logger.info( + `${this.logPrefix} Cached OBO connection was rejected; deferring to a live request for re-exchange`, + ); + this.abandonOboConnection(connection, new Error('OBO re-exchange requires a live request')); + return; + } + + if (refreshAttempted) { + logger.warn(`${this.logPrefix} Refreshed OBO token was rejected as well`); + this.abandonOboConnection(connection, new Error('Refreshed OBO token was rejected')); + return; + } + refreshAttempted = true; + + logger.info(`${this.logPrefix} OBO token rejected by server; re-running token exchange`); + try { + const tokens = await this.getOboTokens(true); + if (!tokens?.access_token) { + throw new Error(`OBO token exchange returned no token for "${this.serverName}".`); + } + connection.setOAuthTokens(tokens); + connection.setAuthorizationHeader(tokens.access_token); + logger.info(`${this.logPrefix} OBO token re-exchanged; retrying connection`); + connection.emit('oauthHandled'); + } catch (error) { + logger.error(`${this.logPrefix} OBO token re-exchange failed`, error); + this.abandonOboConnection(connection, this.toOboRefreshError(error)); + } + }; + + connection.on('oauthRequired', oboHandler); + + return () => { + connection.removeListener('oauthRequired', oboHandler); + }; + } + + private toOboRefreshError(error: unknown): Error { + if (error instanceof OboTokenResolutionError) { + return this.createOboConnectionError(error); + } + if (error instanceof Error) { + return error; + } + return new Error(`OBO token re-exchange failed for "${this.serverName}".`); + } + + /** + * Ends recovery for an OBO connection holding a credential nothing can replace. + * Reconnect attempts would replay the rejected bearer through every backoff step + * and spend the circuit breaker's cycle budget, so the connection is retired and + * left for its next borrower to dispose and rebuild. + */ + private abandonOboConnection(connection: MCPConnection, error: Error): void { + connection.stopReconnecting(); + connection.emit('oauthFailed', error); + } + /** Sets up OAuth event handlers for the connection */ protected handleOAuthEvents( connection: MCPConnection, diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 094a6d396c..48d71f8120 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -136,6 +136,8 @@ describe('MCPConnectionFactory', () => { connect: jest.fn(), isConnected: jest.fn(), setOAuthTokens: jest.fn(), + setAuthorizationHeader: jest.fn(), + stopReconnecting: jest.fn(), on: jest.fn().mockReturnValue(mockConnectionInstance), once: jest.fn().mockReturnValue(mockConnectionInstance), off: jest.fn().mockReturnValue(mockConnectionInstance), @@ -4859,6 +4861,7 @@ describe('MCPConnectionFactory', () => { oboTokenResolver, upstreamTokenProvider, undefined, + false, ); }); @@ -4890,4 +4893,153 @@ describe('MCPConnectionFactory', () => { expect(resolveOboToken).not.toHaveBeenCalled(); }); }); + + describe('OBO authentication failures', () => { + const oboServerConfig = { + type: 'sse' as const, + url: 'https://obo.example.com', + requiresOAuth: false, + obo: { scopes: 'api://obo-server/Mcp.Tools.ReadWrite' }, + } as unknown as t.MCPOptions; + + const connectionTokens = { + access_token: 'rejected-obo-token', + token_type: 'Bearer' as const, + obtained_at: Date.now(), + expires_at: Date.now() + 3600_000, + }; + const refreshedTokens = { + access_token: 'refreshed-obo-token', + token_type: 'Bearer' as const, + obtained_at: Date.now(), + expires_at: Date.now() + 3600_000, + }; + + let resolveOboToken: jest.Mock; + + /** The `oauthRequired` listener the factory attached for this connection. */ + const getAuthHandler = (): (() => Promise) => { + const onCall = (mockConnectionInstance.on as jest.Mock).mock.calls.find( + ([event]: [string]) => event === 'oauthRequired', + ); + return onCall![1] as () => Promise; + }; + + const createOboConnection = () => + MCPConnectionFactory.create( + { serverName: 'obo-srv', serverConfig: oboServerConfig }, + { + useOAuth: true, + user: mockUser, + flowManager: mockFlowManager, + tokenMethods: { + findToken: jest.fn(), + createToken: jest.fn(), + updateToken: jest.fn(), + deleteTokens: jest.fn(), + }, + oboTokenResolver: jest.fn(), + upstreamTokenProvider: jest.fn(), + }, + ); + + beforeEach(() => { + mockProcessMCPEnv.mockReturnValue(oboServerConfig); + mockConnectionInstance.connect.mockResolvedValue(undefined); + mockConnectionInstance.isConnected.mockResolvedValue(true); + ({ resolveOboToken } = jest.requireMock('~/mcp/oauth') as { + resolveOboToken: jest.Mock; + }); + resolveOboToken.mockReset(); + resolveOboToken.mockResolvedValue(connectionTokens); + }); + + it('re-exchanges a rejected token and lets the connect retry proceed', async () => { + resolveOboToken + .mockResolvedValueOnce(connectionTokens) + .mockResolvedValueOnce(refreshedTokens); + /** The real handler runs inside `connectClient`, so fire it from within `connect()`. */ + mockConnectionInstance.connect.mockImplementationOnce(async () => { + await getAuthHandler()(); + }); + + await createOboConnection(); + + expect(resolveOboToken).toHaveBeenCalledTimes(2); + /** The rejected credential is still inside its cached lifetime, so the cache is bypassed. */ + expect(resolveOboToken).toHaveBeenLastCalledWith( + mockUser, + oboServerConfig.obo, + expect.any(Function), + expect.any(Function), + undefined, + true, + ); + expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(refreshedTokens); + /** Runtime headers outrank the transport's, so the stale bearer has to be replaced too. */ + expect(mockConnectionInstance.setAuthorizationHeader).toHaveBeenCalledWith( + 'refreshed-obo-token', + ); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled'); + expect(mockConnectionInstance.emit).not.toHaveBeenCalledWith( + 'oauthFailed', + expect.anything(), + ); + expect(mockConnectionInstance.stopReconnecting).not.toHaveBeenCalled(); + }); + + it('retires the connection when the re-exchange fails', async () => { + resolveOboToken + .mockResolvedValueOnce(connectionTokens) + .mockRejectedValueOnce(new Error('IdP rejected the exchange')); + mockConnectionInstance.connect.mockImplementationOnce(async () => { + await getAuthHandler()(); + }); + + await createOboConnection(); + + expect(mockConnectionInstance.setOAuthTokens).not.toHaveBeenCalled(); + expect(mockConnectionInstance.stopReconnecting).toHaveBeenCalled(); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith( + 'oauthFailed', + expect.objectContaining({ message: 'IdP rejected the exchange' }), + ); + }); + + it('re-exchanges at most once per connection attempt', async () => { + resolveOboToken + .mockResolvedValueOnce(connectionTokens) + .mockResolvedValueOnce(refreshedTokens); + mockConnectionInstance.connect.mockImplementationOnce(async () => { + const handler = getAuthHandler(); + await handler(); + await handler(); + }); + + await createOboConnection(); + + expect(resolveOboToken).toHaveBeenCalledTimes(2); + expect(mockConnectionInstance.stopReconnecting).toHaveBeenCalled(); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith( + 'oauthFailed', + expect.objectContaining({ message: 'Refreshed OBO token was rejected' }), + ); + }); + + it('defers to a live request instead of re-exchanging on a cached connection', async () => { + await createOboConnection(); + expect(resolveOboToken).toHaveBeenCalledTimes(1); + + await getAuthHandler()(); + + /** The creating request is over, so `upstreamTokenProvider` no longer has a live session. */ + expect(resolveOboToken).toHaveBeenCalledTimes(1); + expect(mockConnectionInstance.setOAuthTokens).not.toHaveBeenCalled(); + expect(mockConnectionInstance.stopReconnecting).toHaveBeenCalled(); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith( + 'oauthFailed', + expect.objectContaining({ message: 'OBO re-exchange requires a live request' }), + ); + }); + }); }); diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index 2ec69c9830..cd61548ec1 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -1158,6 +1158,30 @@ export class MCPConnection extends EventEmitter { return this.requestHeaders; } + /** + * Replaces only the runtime bearer, leaving the rest of the per-request headers + * a tool call installed intact. Runtime headers outrank the ones baked into the + * transport (see `buildFetchInit`), so a credential refreshed mid-connection has + * to land here as well or the stale `authorization` from the last tool call would + * keep winning on the retry. + */ + setAuthorizationHeader(accessToken: string): void { + this.requestHeaders = { + ...(this.requestHeaders ?? {}), + authorization: `Bearer ${accessToken}`, + }; + } + + /** + * Halts reconnect attempts for a credential this connection cannot refresh on its + * own. Retrying would replay the rejected bearer through every backoff step and + * spend the circuit breaker's cycle budget; the connection is instead left for its + * next borrower to dispose and rebuild with a live credential. + */ + stopReconnecting(): void { + this.shouldStopReconnecting = true; + } + constructor(params: MCPConnectionParams) { super(); this.options = params.serverConfig; diff --git a/packages/api/src/mcp/oauth/obo.spec.ts b/packages/api/src/mcp/oauth/obo.spec.ts index 9a837694b7..6f90b2c4aa 100644 --- a/packages/api/src/mcp/oauth/obo.spec.ts +++ b/packages/api/src/mcp/oauth/obo.spec.ts @@ -94,6 +94,30 @@ describe('resolveOboToken', () => { expect(result.access_token).toBe('exchanged-mcp-token'); }); + it('bypasses the resolver cache when forceRefresh is set', async () => { + /** + * A downstream 401 can arrive while the credential is still inside its cached + * lifetime (revoked, or its scopes invalidated), so a cached read would hand + * back the same rejected bearer. + */ + await resolveOboToken( + mockUser as IUser, + oboConfig, + mockResolver, + liveProvider, + undefined, + true, + ); + + expect(mockResolver).toHaveBeenCalledWith( + mockUser, + 'live-access-token', + 'api://mcp-server-id/Mcp.Tools.ReadWrite', + false, + undefined, + ); + }); + it('throws missing_upstream_token when federated fallback token is invalid', async () => { mockExtractOpenIDTokenInfo.mockReturnValue({ accessToken: 'federated-access-token', diff --git a/packages/api/src/mcp/oauth/obo.ts b/packages/api/src/mcp/oauth/obo.ts index a87e444e86..65d26ef9ee 100644 --- a/packages/api/src/mcp/oauth/obo.ts +++ b/packages/api/src/mcp/oauth/obo.ts @@ -191,6 +191,11 @@ function buildUpstreamTokenInfo( * itself carries that verified upstream bearer — still works. Browser requests * whose Express session was cleared reject in the provider instead of reaching * this fallback with a stale strategy-time snapshot. + * + * @param forceRefresh Bypasses the resolver's token cache. Set it when the downstream + * server has rejected the current credential: a revoked or scope-invalidated token is + * still inside its cached lifetime, so a cached read would hand back the same rejected + * bearer instead of minting a replacement. */ export async function resolveOboToken( user: IUser, @@ -198,6 +203,7 @@ export async function resolveOboToken( oboTokenResolver: OboTokenResolver, upstreamTokenProvider: UpstreamTokenProvider, identityContext?: AuthIdentityContext, + forceRefresh = false, ): Promise { let liveTokens: OIDCTokens | null; try { @@ -240,7 +246,7 @@ export async function resolveOboToken( user, tokenInfo.accessToken, oboConfig.scopes, - true, + !forceRefresh, identityContext, );