CERBERUS/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2026-08-27 04:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 13

🧷 fix: Safely Recover Runtime MCP OAuth Rejections (#14684)

Browse source 
* fix runtime MCP OAuth recovery

* style: sort LC-008 imports

* fix: single-flight runtime OAuth handlers

* fix: retain transport OAuth failures for recovery

* fix(mcp): preserve OAuth recovery connections

* test(mcp): type request-scoped config fixture

* fix(mcp): harden shared OAuth recovery

* fix(mcp): bound OAuth recovery escalation

* style(mcp): sort OAuth integration imports

* fix(mcp): harden OAuth recovery boundaries

* fix(mcp): abort shared recovery waiters

* fix(mcp): bound request OAuth recovery phases

* fix(mcp): close OAuth recovery ownership gaps

* fix(mcp): retry borrowers closed by OAuth recovery

* fix(mcp): drain borrowers before OAuth reconnect

* fix(mcp): preserve eviction across OAuth recovery

* fix(mcp): unify OAuth recovery leases

* fix(mcp): serialize cache reuse with recovery

* fix(mcp): make recovery checkout atomic

* test(mcp): use numeric config timestamp

* fix(mcp): reacquire recovery checkouts

* fix(mcp): retain shared recovery disposal

* fix(mcp): restart checkout after recovery takeover

* fix(mcp): close recovery lifecycle gaps

* refactor(mcp): deepen OAuth recovery lifecycle

* fix(mcp): harden OAuth lifecycle disposal

* style(mcp): sort OAuth lifecycle imports

* fix: lease MCP OAuth lifecycle edges

* fix(mcp): isolate shared OAuth flows from aborts

---------

Co-authored-by: Dennis Schenk <dennis@gridonic.ch>
This commit is contained in:
Danny Avila 2026-08-10 10:38:34 -04:00 • committed by GitHub
parent 26bcbb713c
commit 7fc62023eb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 3439 additions and 870 deletions
Download patch file Download diff file Expand all files Collapse all files

136
api/server/routes/__tests__/mcp.spec.js
View file

@ -187,6 +187,11 @@ const mockOAuthCompletion = (tokens) => {
);
};
const createLeasedMcpManager = (connection, overrides = {}) => ({
...overrides,
withUserConnectionLease: jest.fn((_options, useConnection) => useConnection(connection)),
});
describe('MCP Routes', () => {
let app;
let mongoServer;
@ -853,11 +858,9 @@ describe('MCP Routes', () => {
MCPTokenStorage.storeTokens.mockResolvedValue();
mockRegistryInstance.getServerConfig.mockResolvedValue({});
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
};
const mockMcpManager = createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
});
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
@ -915,11 +918,9 @@ describe('MCP Routes', () => {
MCPTokenStorage.storeTokens.mockResolvedValue();
mockRegistryInstance.getServerConfig.mockResolvedValue({});
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
};
const mockMcpManager = createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
});
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
@ -973,11 +974,7 @@ describe('MCP Routes', () => {
const fetchOrderedToolsSnapshot = jest
.fn()
.mockResolvedValue({ tools: fetchedTools, complete: true });
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
fetchOrderedToolsSnapshot,
}),
};
const mockMcpManager = createLeasedMcpManager({ fetchOrderedToolsSnapshot });
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
@ -992,8 +989,9 @@ describe('MCP Routes', () => {
expect(response.status).toBe(302);
expect(mockResolveAllMcpConfigs).toHaveBeenCalledWith('test-user-id');
expect(fetchOrderedToolsSnapshot).toHaveBeenCalledTimes(1);
expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith(
expect(mockMcpManager.withUserConnectionLease).toHaveBeenCalledWith(
expect.objectContaining({ serverConfig: mergedServerConfig }),
expect.any(Function),
);
expect(updateMCPServerTools).toHaveBeenCalledWith({
userId: 'test-user-id',
@ -1044,13 +1042,9 @@ describe('MCP Routes', () => {
[`mcp_test-server`]: { LITELLM_KEY: 'sk-real-user-key' },
});
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
fetchToolsSnapshot: jest
.fn()
.mockResolvedValue({ tools: fetchedTools, complete: true }),
}),
};
const mockMcpManager = createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: fetchedTools, complete: true }),
});
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
@ -1068,8 +1062,9 @@ describe('MCP Routes', () => {
servers: ['test-server'],
findPluginAuthsByKeys: require('~/models').findPluginAuthsByKeys,
});
expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith(
expect(mockMcpManager.withUserConnectionLease).toHaveBeenCalledWith(
expect.objectContaining({ customUserVars: { LITELLM_KEY: 'sk-real-user-key' } }),
expect.any(Function),
);
});
@ -1109,13 +1104,9 @@ describe('MCP Routes', () => {
mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'test-server': mergedServerConfig });
require('@librechat/api').getUserMCPAuthMap.mockClear();
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
fetchToolsSnapshot: jest
.fn()
.mockResolvedValue({ tools: fetchedTools, complete: true }),
}),
};
const mockMcpManager = createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: fetchedTools, complete: true }),
});
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
@ -1129,8 +1120,9 @@ describe('MCP Routes', () => {
expect(response.status).toBe(302);
expect(require('@librechat/api').getUserMCPAuthMap).not.toHaveBeenCalled();
expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith(
expect(mockMcpManager.withUserConnectionLease).toHaveBeenCalledWith(
expect.objectContaining({ customUserVars: undefined }),
expect.any(Function),
);
});
@ -1241,7 +1233,9 @@ describe('MCP Routes', () => {
}),
};
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
withUserConnectionLease: jest.fn((_options, useConnection) =>
useConnection(mockUserConnection),
),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
@ -1285,7 +1279,7 @@ describe('MCP Routes', () => {
);
const storeInvocation = MCPTokenStorage.storeTokens.mock.invocationCallOrder[0];
const flowCompletionInvocation = mockFlowManager.completeFlow.mock.invocationCallOrder[0];
const connectInvocation = mockMcpManager.getUserConnection.mock.invocationCallOrder[0];
const connectInvocation = mockMcpManager.withUserConnectionLease.mock.invocationCallOrder[0];
expect(storeInvocation).toBeLessThan(flowCompletionInvocation);
expect(storeInvocation).toBeLessThan(connectInvocation);
expect(mockFlowManager.completeFlow).toHaveBeenCalledWith(
@ -1331,11 +1325,11 @@ describe('MCP Routes', () => {
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
});
require('~/config').getMCPManager.mockReturnValue({
getUserConnection: jest.fn().mockResolvedValue({
require('~/config').getMCPManager.mockReturnValue(
createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
});
);
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
@ -1406,11 +1400,11 @@ describe('MCP Routes', () => {
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
});
require('~/config').getMCPManager.mockReturnValue({
getUserConnection: jest.fn().mockResolvedValue({
require('~/config').getMCPManager.mockReturnValue(
createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
});
);
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
@ -1464,11 +1458,11 @@ describe('MCP Routes', () => {
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
});
require('~/config').getMCPManager.mockReturnValue({
getUserConnection: jest.fn().mockResolvedValue({
require('~/config').getMCPManager.mockReturnValue(
createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
});
);
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
@ -1518,11 +1512,11 @@ describe('MCP Routes', () => {
require('~/config').getOAuthReconnectionManager.mockReturnValue({
clearReconnection: jest.fn(),
});
require('~/config').getMCPManager.mockReturnValue({
getUserConnection: jest.fn().mockResolvedValue({
require('~/config').getMCPManager.mockReturnValue(
createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
});
);
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
@ -1639,7 +1633,7 @@ describe('MCP Routes', () => {
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
const mockMcpManager = {
getUserConnection: jest.fn().mockRejectedValue(new Error('Reconnection failed')),
withUserConnectionLease: jest.fn().mockRejectedValue(new Error('Reconnection failed')),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
@ -1692,7 +1686,7 @@ describe('MCP Routes', () => {
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
const mockMcpManager = {
getUserConnection: jest.fn(),
withUserConnectionLease: jest.fn(),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
@ -1711,7 +1705,7 @@ describe('MCP Routes', () => {
expect(response.status).toBe(302);
expect(response.headers.location).toBe(`${basePath}/oauth/error?error=callback_failed`);
expect(mockFlowManager.completeFlow).not.toHaveBeenCalled();
expect(mockMcpManager.getUserConnection).not.toHaveBeenCalled();
expect(mockMcpManager.withUserConnectionLease).not.toHaveBeenCalled();
});
it('should use original flow state credentials when storing tokens', async () => {
@ -1755,9 +1749,7 @@ describe('MCP Routes', () => {
const mockUserConnection = {
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
};
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
};
const mockMcpManager = createLeasedMcpManager(mockUserConnection);
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getOAuthReconnectionManager = jest.fn().mockReturnValue({
clearReconnection: jest.fn(),
@ -2927,11 +2919,9 @@ describe('MCP Routes', () => {
};
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
}),
};
const mockMcpManager = createLeasedMcpManager({
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
});
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
const flowId = 'test-user-id:test-server';
@ -2980,15 +2970,15 @@ describe('MCP Routes', () => {
MCPTokenStorage.storeTokens.mockResolvedValue();
mockRegistryInstance.getServerConfig.mockResolvedValue({});
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue({
const mockMcpManager = createLeasedMcpManager(
{
fetchToolsSnapshot: jest.fn().mockResolvedValue({
tools: [{ name: 'test-tool', description: 'Test tool' }],
complete: true,
}),
}),
getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'),
};
},
{ getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation') },
);
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
const flowId = 'test-user-id:test-server';
@ -3035,15 +3025,19 @@ describe('MCP Routes', () => {
mockOAuthCompletion(mockTokens);
MCPTokenStorage.storeTokens.mockResolvedValue();
mockRegistryInstance.getServerConfig.mockResolvedValue({});
require('~/config').getMCPManager.mockReturnValue({
getUserConnection: jest.fn().mockResolvedValue({
fetchToolsSnapshot: jest.fn().mockResolvedValue({
tools: [{ name: 'partial-tool', description: 'Only the first page' }],
complete: false,
}),
}),
getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'),
});
require('~/config').getMCPManager.mockReturnValue(
createLeasedMcpManager(
{
fetchToolsSnapshot: jest.fn().mockResolvedValue({
tools: [{ name: 'partial-tool', description: 'Only the first page' }],
complete: false,
}),
},
{
getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'),
},
),
);
const flowId = 'test-user-id:test-server';
const csrfToken = generateTestCsrfToken(flowId);

52
api/server/routes/mcp.js
View file

@ -522,33 +522,39 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
}
const customUserVars = getServerCustomUserVars(userMCPAuthMap, serverName);
const userConnection = await mcpManager.getUserConnection({
user,
serverName,
flowManager,
serverConfig,
customUserVars,
tokenMethods: {
findToken: db.findToken,
updateToken: db.updateToken,
createToken: db.createToken,
deleteTokens: db.deleteTokens,
const { snapshot, publicationGeneration } = await mcpManager.withUserConnectionLease(
{
user,
serverName,
flowManager,
serverConfig,
customUserVars,
tokenMethods: {
findToken: db.findToken,
updateToken: db.updateToken,
createToken: db.createToken,
deleteTokens: db.deleteTokens,
},
},
});
async (userConnection) => {
logger.info(
`[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`,
);
logger.info(
`[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`,
const oauthReconnectionManager = getOAuthReconnectionManager();
oauthReconnectionManager.clearReconnection(flowState.userId, serverName);
const snapshot =
typeof userConnection.fetchOrderedToolsSnapshot === 'function'
? await userConnection.fetchOrderedToolsSnapshot()
: await userConnection.fetchToolsSnapshot();
return {
snapshot,
publicationGeneration: mcpManager.getToolPublicationGeneration?.(userConnection),
};
},
);
const oauthReconnectionManager = getOAuthReconnectionManager();
oauthReconnectionManager.clearReconnection(flowState.userId, serverName);
const snapshot =
typeof userConnection.fetchOrderedToolsSnapshot === 'function'
? await userConnection.fetchOrderedToolsSnapshot()
: await userConnection.fetchToolsSnapshot();
if (snapshot.complete) {
const publicationGeneration = mcpManager.getToolPublicationGeneration?.(userConnection);
await updateMCPServerTools({
userId: flowState.userId,
serverName,

133
api/server/services/MCP.js
View file

@ -608,24 +608,6 @@ function createOAuthEnd({ res, stepId, toolCall, streamId = null, jobCreatedAt }
};
}
/**
* @param {object} params
* @param {string} params.userId - The ID of the user.
* @param {string} params.serverName - The name of the server.
* @param {string} params.toolName - The name of the tool.
* @param {string} [params.tenantId] - The tenant ID for the current request.
* @param {FlowStateManager<any>} params.flowManager - The flow manager instance.
*/
function createAbortHandler({ userId, serverName, toolName, tenantId, flowManager }) {
return function () {
logger.info(`[MCP][User: ${userId}][${serverName}][${toolName}] Tool call aborted`);
const flowId = getOAuthFlowId(userId, serverName, tenantId);
// Clean up both mcp_oauth and mcp_get_tokens flows
flowManager.failFlow(flowId, 'mcp_oauth', new Error('Tool call aborted'));
flowManager.failFlow(flowId, 'mcp_get_tokens', new Error('Tool call aborted'));
};
}
/**
* @param {Object} params
* @param {() => Promise<void>} params.runStepEmitter
@ -696,66 +678,43 @@ async function reconnectServer({
serverName,
});
// Set up abort handler to clean up OAuth flows if request is aborted
const tenantId = user?.tenantId ?? getTenantId();
const oauthFlowId = getOAuthFlowId(user.id, serverName, tenantId);
const abortHandler = () => {
logger.info(
`[MCP][User: ${user.id}][${serverName}] Tool loading aborted, cleaning up OAuth flows`,
);
// Clean up both mcp_oauth and mcp_get_tokens flows
flowManager.failFlow(oauthFlowId, 'mcp_oauth', new Error('Tool loading aborted'));
flowManager.failFlow(oauthFlowId, 'mcp_get_tokens', new Error('Tool loading aborted'));
};
if (signal) {
signal.addEventListener('abort', abortHandler, { once: true });
}
try {
const runStepEmitter = createRunStepEmitter({
res,
index,
runId,
stepId,
toolCall,
streamId,
jobCreatedAt,
});
const runStepDeltaEmitter = createRunStepDeltaEmitter({
res,
stepId,
toolCall,
streamId,
jobCreatedAt,
});
const callback = createOAuthCallback({ runStepEmitter, runStepDeltaEmitter });
const oauthStart = createOAuthStart({
res,
flowId,
callback,
flowManager,
});
return await reinitMCPServer({
user,
signal,
serverName,
configServers,
oauthStart,
flowManager,
userMCPAuthMap,
requestBody,
requestScopedConnections,
forceNew: true,
returnOnOAuth: false,
connectionTimeout: Time.THIRTY_SECONDS,
});
} finally {
// Clean up abort handler to prevent memory leaks
if (signal) {
signal.removeEventListener('abort', abortHandler);
}
}
const runStepEmitter = createRunStepEmitter({
res,
index,
runId,
stepId,
toolCall,
streamId,
jobCreatedAt,
});
const runStepDeltaEmitter = createRunStepDeltaEmitter({
res,
stepId,
toolCall,
streamId,
jobCreatedAt,
});
const callback = createOAuthCallback({ runStepEmitter, runStepDeltaEmitter });
const oauthStart = createOAuthStart({
res,
flowId,
callback,
flowManager,
});
return await reinitMCPServer({
user,
signal,
serverName,
configServers,
oauthStart,
flowManager,
userMCPAuthMap,
requestBody,
requestScopedConnections,
forceNew: true,
returnOnOAuth: false,
connectionTimeout: Time.THIRTY_SECONDS,
});
}
/**
@ -1090,11 +1049,6 @@ function createToolInstance({
const effectiveUser = config?.configurable?.user ?? capturedUser;
const permissionUser = effectiveUser;
const userId = effectiveUser?.id || config?.configurable?.user_id || capturedUser?.id;
/** @type {ReturnType<typeof createAbortHandler>} */
let abortHandler = null;
/** @type {AbortSignal} */
let derivedSignal = null;
try {
const provider = (config?.metadata?.provider || capturedProvider)?.toLowerCase();
const canUseMCP = mcpPermissionContext
@ -1105,7 +1059,7 @@ function createToolInstance({
}
const flowsCache = getLogStores(CacheKeys.FLOWS);
const flowManager = getFlowStateManager(flowsCache);
derivedSignal = config?.signal ? AbortSignal.any([config.signal]) : undefined;
const derivedSignal = config?.signal ? AbortSignal.any([config.signal]) : undefined;
const mcpManager = getMCPManager(userId);
const { args: _args, stepId, ...toolCall } = config.toolCall ?? {};
@ -1130,12 +1084,6 @@ function createToolInstance({
jobCreatedAt,
});
if (derivedSignal) {
const tenantId = config?.configurable?.user?.tenantId ?? getTenantId();
abortHandler = createAbortHandler({ userId, serverName, toolName, tenantId, flowManager });
derivedSignal.addEventListener('abort', abortHandler, { once: true });
}
const customUserVars =
config?.configurable?.userMCPAuthMap?.[`${Constants.mcp_prefix}${serverName}`];
@ -1205,11 +1153,6 @@ function createToolInstance({
throw new Error(
`[MCP][${serverName}][${toolName}] tool call failed${error?.message ? `: ${error?.message}` : '.'}`,
);
} finally {
// Clean up abort handler to prevent memory leaks
if (abortHandler && derivedSignal) {
derivedSignal.removeEventListener('abort', abortHandler);
}
}
};

96
api/server/services/MCP.spec.js
View file

@ -1411,7 +1411,7 @@ describe('User parameter passing tests', () => {
}
});
it('should fail tenant-scoped OAuth flows when tool loading is aborted', async () => {
it('does not fail shared OAuth flows when tool loading is aborted', async () => {
const mockUser = { id: 'tenant-user', name: 'Tenant User' };
const mockRes = { write: jest.fn(), flush: jest.fn() };
const abortController = new AbortController();
@ -1419,9 +1419,7 @@ describe('User parameter passing tests', () => {
createFlowWithHandler: jest.fn(),
failFlow: jest.fn(),
};
mockGetTenantId.mockReturnValue('tenant/a');
mockGetFlowStateManager.mockReturnValue(mockFlowManager);
MCPOAuthHandler.generateFlowId.mockReturnValue('tenant-flow-id');
let resolveReinit;
mockReinitMCPServer.mockImplementation(
@ -1445,21 +1443,7 @@ describe('User parameter passing tests', () => {
resolveReinit({ tools: [], availableTools: {} });
await createToolsPromise;
expect(MCPOAuthHandler.generateFlowId).toHaveBeenCalledWith(
mockUser.id,
'tenant-abort-server',
'tenant/a',
);
expect(mockFlowManager.failFlow).toHaveBeenCalledWith(
'tenant-flow-id',
'mcp_oauth',
expect.any(Error),
);
expect(mockFlowManager.failFlow).toHaveBeenCalledWith(
'tenant-flow-id',
'mcp_get_tokens',
expect.any(Error),
);
expect(mockFlowManager.failFlow).not.toHaveBeenCalled();
});
it('should throw error if user is not provided', async () => {
@ -1487,6 +1471,82 @@ describe('User parameter passing tests', () => {
});
describe('createMCPTool', () => {
it('keeps shared OAuth recovery alive when one tool caller aborts', async () => {
const mockUser = { id: 'shared-recovery-user', role: 'USER' };
const mockRes = { write: jest.fn(), flush: jest.fn() };
const ownerAbort = new AbortController();
const waiterAbort = new AbortController();
const flowManager = {
getFlowState: jest.fn().mockResolvedValue(null),
createFlowWithHandler: jest.fn(),
failFlow: jest.fn(),
};
let completeRecovery;
const sharedRecovery = new Promise((resolve) => {
completeRecovery = resolve;
});
const callTool = jest.fn(({ options }) => {
const signal = options?.signal;
return new Promise((resolve, reject) => {
const onAbort = () => {
signal?.removeEventListener('abort', onAbort);
reject(new Error('tool caller aborted'));
};
signal?.addEventListener('abort', onAbort, { once: true });
sharedRecovery.then(() => {
signal?.removeEventListener('abort', onAbort);
resolve(['ok', null]);
});
});
});
const { getRoleByName } = require('~/models');
getRoleByName.mockResolvedValue({
permissions: {
[PermissionTypes.MCP_SERVERS]: {
[Permissions.USE]: true,
},
},
});
mockGetFlowStateManager.mockReturnValue(flowManager);
mockGetMCPManager.mockReturnValue({ callTool });
const mcpTool = await createMCPTool({
res: mockRes,
user: mockUser,
config: { url: 'https://runtime-oauth.example.com/mcp' },
toolKey: `test-tool${D}test-server`,
provider: 'openai',
userMCPAuthMap: {},
availableTools: {
[`test-tool${D}test-server`]: {
function: {
description: 'Cached tool',
parameters: { type: 'object', properties: {} },
},
},
},
});
const createConfig = (signal) => ({
signal,
configurable: { user: mockUser },
metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' },
toolCall: {},
});
const ownerCall = mcpTool.invoke({}, createConfig(ownerAbort.signal));
const waiterCall = mcpTool.invoke({}, createConfig(waiterAbort.signal));
await new Promise((resolve) => setImmediate(resolve));
ownerAbort.abort();
await expect(ownerCall).rejects.toThrow('Aborted');
expect(flowManager.failFlow).not.toHaveBeenCalled();
completeRecovery();
await expect(waiterCall).resolves.toBe('ok');
expect(callTool).toHaveBeenCalledTimes(2);
});
it.each(['OAuth flow initiated - return early', 'Pending OAuth flow reused - return early'])(
'preserves runtime-detected OAuth for the internal signal: %s',
async (oauthSignal) => {