mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-10-22 18:13:27 +00:00
🔁 fix: Re-Exchange Rejected OBO Tokens on Tool Calls Too
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.
This commit is contained in:
parent
c77be0340f
commit
09ccb3423a
4 changed files with 207 additions and 35 deletions
|
|
@ -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);
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -1000,9 +1000,19 @@ Please follow these instructions when using tools from the respective MCP server
|
|||
const resolvedHeaders: Record<string, string> =
|
||||
'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<void> => {
|
||||
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;
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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', () => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue