From 09ccb3423a94b2fe4a36f67c491ee3101447d4cd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 3 Sep 2026 23:26:30 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=81=20fix:=20Re-Exchange=20Rejected=20?= =?UTF-8?q?OBO=20Tokens=20on=20Tool=20Calls=20Too?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses both codex findings on c77be03. A tool-call 401 never reached the connection-level handler: `oauthRequired` is emitted only by `connectClient`, and `callTool`'s own recovery attaches solely when `isOAuthServer(currentOptions) || connection.usesOAuth()` — both false for an OBO-only config. It rethrew untouched, and the resolver kept serving the rejected bearer from cache on every later call until it expired. This is the dominant path for OBO traffic, so the cache-invalidation half of the previous commit was not actually reachable in practice. The per-call token resolution moves into `applyOboAuthorization(forceRefresh)`, which the auth-error path re-runs with the cache bypassed before retrying once. Recovery belongs here rather than at the connection level: the request is still live, so `upstreamTokenProvider` can mint a replacement. Also propagate the re-exchange failure. `connectClient` rejects its handling promise with the mapped error and then rethrows the server's original 401, so a diagnosis like an unrefreshable sign-in session was lost. `createConnection` now surfaces it in place of the 401 — the same substitution the initial OBO resolution already makes — and only for a genuine exchange failure, not for the control-flow deferrals. --- packages/api/src/mcp/MCPConnectionFactory.ts | 15 ++- packages/api/src/mcp/MCPManager.ts | 99 ++++++++++------ .../__tests__/MCPConnectionFactory.test.ts | 20 ++++ .../api/src/mcp/__tests__/MCPManager.test.ts | 108 ++++++++++++++++++ 4 files changed, 207 insertions(+), 35 deletions(-) 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', () => {