diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index 41102a61ad..4fed085da8 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -77,6 +77,8 @@ export class MCPConnectionFactory { protected readonly oboTrustChecker?: OboTrustChecker; protected readonly upstreamTokenProvider?: UpstreamTokenProvider; protected readonly oboIdentityContext?: AuthIdentityContext; + /** Why the OBO re-exchange failed, when that is more actionable than the server's 401. */ + private oboRefreshError?: Error; private connectionReady = false; /** * Snapshot of the tenant context at factory construction time. Captured eagerly @@ -583,6 +585,9 @@ export class MCPConnectionFactory { } catch { logger.warn(`${this.logPrefix} Failed to clean up rejected MCP connection`); } + if (this.oboRefreshError) { + throw this.oboRefreshError; + } throw error; } } @@ -1160,7 +1165,15 @@ export class MCPConnectionFactory { connection.emit('oauthHandled'); } catch (error) { logger.error(`${this.logPrefix} OBO token re-exchange failed`, error); - this.abandonOboConnection(connection, this.toOboRefreshError(error)); + /** + * `connectClient` rejects its handling promise with this error and then + * rethrows the server's original 401, so a diagnosis like an unrefreshable + * sign-in session would be lost. `createConnection` reads it back and + * surfaces it in place of the generic 401 — the same substitution the + * initial OBO resolution already makes. + */ + this.oboRefreshError = this.toOboRefreshError(error); + this.abandonOboConnection(connection, this.oboRefreshError); } }; diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index f3f703faa1..bdaf218a00 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -1000,9 +1000,19 @@ Please follow these instructions when using tools from the respective MCP server const resolvedHeaders: Record = 'headers' in currentOptions ? { ...(currentOptions.headers || {}) } : {}; - /** Resolve the current OBO token for this tool call; the resolver may serve cached tokens. */ const oboConfig = rawConfig.obo; - if (oboConfig && oboTokenResolver && user) { + const usesObo = Boolean(oboConfig && oboTokenResolver && user); + + /** + * Resolves the downstream token for this call and installs it as the request + * bearer. `forceRefresh` bypasses the resolver's cache, which is what a + * rejected credential needs: a revoked or scope-invalidated token is still + * inside its cached lifetime, so a cached read returns the same dead bearer. + */ + const applyOboAuthorization = async (forceRefresh: boolean): Promise => { + if (!oboConfig || !oboTokenResolver || !user) { + return; + } if (!upstreamTokenProvider) { throw new McpError( ErrorCode.InternalError, @@ -1035,6 +1045,7 @@ Please follow these instructions when using tools from the respective MCP server oboTokenResolver, upstreamTokenProvider, oboIdentityContext, + forceRefresh, ); } catch (error) { if (error instanceof OboTokenResolutionError) { @@ -1053,7 +1064,10 @@ Please follow these instructions when using tools from the respective MCP server ); } resolvedHeaders['Authorization'] = `Bearer ${oboTokens.access_token}`; - } + }; + + /** Resolve the current OBO token for this tool call; the resolver may serve cached tokens. */ + await applyOboAuthorization(false); if ( userId && user && @@ -1150,41 +1164,58 @@ Please follow these instructions when using tools from the respective MCP server try { result = await requestTool(); } catch (error) { - const requestOAuthHandler = attachSharedOAuthHandler; - if (!requestOAuthHandler || !userId) { - throw error; - } - - if (!connection.isOAuthAuthenticationError(error)) { - throw error; - } - - try { - await waitForRecoveryWithoutLease(() => - this.recoverOAuthConnection( - connection!, - error, - serverName, - userId, - requestOAuthHandler, - oauthStart, - oauthEnd, - flowManager, - options?.signal, - !recoveryTakeoverConsumed, - ), + /** + * An OBO server rejecting the bearer mid-session is recoverable here and + * nowhere else: the downstream token is minted from the upstream session + * this request still holds, and `attachSharedOAuthHandler` is never set for + * an OBO-only config, so the OAuth recovery below would rethrow untouched. + * Without this the rejected token is re-served from cache on every later + * call until it expires. + */ + if (usesObo && connection.isOAuthAuthenticationError(error)) { + logger.info( + `${logPrefix}[${toolName}] OBO token rejected by server; re-exchanging and retrying once`, ); - } catch (recoveryError) { - if (recoveryError instanceof OAuthRecoveryTakeoverRequired) { - throw recoveryError; + await applyOboAuthorization(true); + connection.setRequestHeaders(resolvedHeaders); + result = await requestTool(); + } else { + const requestOAuthHandler = attachSharedOAuthHandler; + if (!requestOAuthHandler || !userId) { + throw error; } - if (options?.signal?.aborted) { - throw recoveryError; + + if (!connection.isOAuthAuthenticationError(error)) { + throw error; } - logger.warn(`${logPrefix}[${toolName}] Runtime OAuth recovery failed`, recoveryError); - throw error; + + try { + await waitForRecoveryWithoutLease(() => + this.recoverOAuthConnection( + connection!, + error, + serverName, + userId, + requestOAuthHandler, + oauthStart, + oauthEnd, + flowManager, + options?.signal, + !recoveryTakeoverConsumed, + ), + ); + } catch (recoveryError) { + if (recoveryError instanceof OAuthRecoveryTakeoverRequired) { + throw recoveryError; + } + if (options?.signal?.aborted) { + throw recoveryError; + } + logger.warn(`${logPrefix}[${toolName}] Runtime OAuth recovery failed`, recoveryError); + throw error; + } + result = await requestTool(); } - result = await requestTool(); } const hasPersistentUserConnections = !!userId && (this.userConnections.get(userId)?.size ?? 0) > 0; diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 48d71f8120..d08e179de3 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -5006,6 +5006,26 @@ describe('MCPConnectionFactory', () => { ); }); + it('surfaces why the re-exchange failed instead of the generic 401', async () => { + const authError = new Error('HTTP 401 Unauthorized'); + resolveOboToken + .mockResolvedValueOnce(connectionTokens) + .mockRejectedValueOnce(new Error('Your sign-in session expired. Please sign in again.')); + /** + * `connectClient` rethrows the server's original 401 after the handler fails, + * so without propagation the caller would only ever see `authError`. + */ + mockConnectionInstance.connect.mockImplementationOnce(async () => { + await getAuthHandler()(); + throw authError; + }); + mockConnectionInstance.isConnected.mockResolvedValue(false); + + await expect(createOboConnection()).rejects.toThrow( + 'Your sign-in session expired. Please sign in again.', + ); + }); + it('re-exchanges at most once per connection attempt', async () => { resolveOboToken .mockResolvedValueOnce(connectionTokens) diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index df42a7bdb1..6bda5e75d9 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -2804,6 +2804,7 @@ describe('MCPManager', () => { mockOboTokenResolver, mockUpstreamTokenProvider, undefined, + false, ); expect(appConnections.get).not.toHaveBeenCalled(); expect(getUserConnectionSpy).toHaveBeenCalled(); @@ -2934,6 +2935,113 @@ describe('MCPManager', () => { }); expect(mockResolveOboToken).not.toHaveBeenCalled(); }); + + it('re-exchanges past the token cache when the server rejects the bearer mid-call', async () => { + const authError = new Error('HTTP 401 Unauthorized'); + const rejectingConnection = { + isConnected: jest.fn().mockResolvedValue(true), + setRequestHeaders: jest.fn(), + isOAuthAuthenticationError: jest.fn().mockReturnValue(true), + timeout: 30000, + client: { + request: jest + .fn() + .mockRejectedValueOnce(authError) + .mockResolvedValueOnce({ + content: [{ type: 'text', text: 'Tool result' }], + isError: false, + }), + }, + } as unknown as MCPConnection; + + mockResolveOboToken + .mockResolvedValueOnce({ + access_token: 'rejected-obo-token', + token_type: 'Bearer', + obtained_at: Date.now(), + expires_at: Date.now() + 3600_000, + }) + .mockResolvedValueOnce({ + access_token: 'reminted-obo-token', + token_type: 'Bearer', + obtained_at: Date.now(), + expires_at: Date.now() + 3600_000, + }); + + mockAppConnections({ get: jest.fn().mockResolvedValue(rejectingConnection) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + jest.spyOn(manager, 'getUserConnection').mockResolvedValue(rejectingConnection); + + await manager.callTool({ + user: mockUser as IUser, + serverName, + toolName: 'test_tool', + provider: 'openai', + flowManager: mockFlowManager as unknown as Parameters< + typeof manager.callTool + >[0]['flowManager'], + oboTokenResolver: mockOboTokenResolver, + upstreamTokenProvider: mockUpstreamTokenProvider, + }); + + /** The rejected token is still inside its cached lifetime, so the cache is bypassed. */ + expect(mockResolveOboToken).toHaveBeenNthCalledWith( + 2, + mockUser, + serverConfig.obo, + mockOboTokenResolver, + mockUpstreamTokenProvider, + undefined, + true, + ); + expect(rejectingConnection.setRequestHeaders as jest.Mock).toHaveBeenLastCalledWith( + expect.objectContaining({ Authorization: 'Bearer reminted-obo-token' }), + ); + expect((rejectingConnection.client.request as jest.Mock).mock.calls).toHaveLength(2); + }); + + it('does not retry a non-auth failure', async () => { + const toolError = new Error('tool blew up'); + const failingConnection = { + isConnected: jest.fn().mockResolvedValue(true), + setRequestHeaders: jest.fn(), + isOAuthAuthenticationError: jest.fn().mockReturnValue(false), + timeout: 30000, + client: { request: jest.fn().mockRejectedValue(toolError) }, + } as unknown as MCPConnection; + + mockResolveOboToken.mockResolvedValue({ + access_token: 'fresh-obo-token', + token_type: 'Bearer', + obtained_at: Date.now(), + expires_at: Date.now() + 3600_000, + }); + + mockAppConnections({ get: jest.fn().mockResolvedValue(failingConnection) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + jest.spyOn(manager, 'getUserConnection').mockResolvedValue(failingConnection); + + await expect( + manager.callTool({ + user: mockUser as IUser, + serverName, + toolName: 'test_tool', + provider: 'openai', + flowManager: mockFlowManager as unknown as Parameters< + typeof manager.callTool + >[0]['flowManager'], + oboTokenResolver: mockOboTokenResolver, + upstreamTokenProvider: mockUpstreamTokenProvider, + }), + ).rejects.toThrow('tool blew up'); + + expect(mockResolveOboToken).toHaveBeenCalledTimes(1); + expect((failingConnection.client.request as jest.Mock).mock.calls).toHaveLength(1); + }); }); describe('getConnection', () => {