mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-27 04:07:05 +00:00
🧷 fix: Safely Recover Runtime MCP OAuth Rejections (#14684)
* 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
GitHub
parent
26bcbb713c
commit
7fc62023eb
20 changed files with 3439 additions and 870 deletions
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ import {
|
|||
} from '~/mcp/oauth';
|
||||
import { sanitizeUrlForLogging, isClientRejectionMessage, isOAuthServer } from './utils';
|
||||
import { PENDING_STALE_MS, normalizeExpiresAt } from '~/flow/manager';
|
||||
import { isOAuthAuthenticationError } from './errors';
|
||||
import { preProcessGraphTokens } from '~/utils/graph';
|
||||
import { withTimeout } from '~/utils/promise';
|
||||
import { MCPConnection } from './connection';
|
||||
|
|
@ -44,6 +45,8 @@ type OAuthRequiredEvent = {
|
|||
skipSilentRefresh?: boolean;
|
||||
};
|
||||
|
||||
type OAuthRecoveryPhase = 'silent-refresh' | 'interactive' | 'terminal';
|
||||
|
||||
/**
|
||||
* Factory for creating MCP connections with optional OAuth authentication.
|
||||
* Handles OAuth flows, token management, and connection retry logic.
|
||||
|
|
@ -988,16 +991,46 @@ export class MCPConnectionFactory {
|
|||
connection: MCPConnection,
|
||||
eventName: 'oauthRequired' | 'oauthReauthenticationRequired' = 'oauthRequired',
|
||||
): () => void {
|
||||
const oauthHandler = async (data: OAuthRequiredEvent) => {
|
||||
const isRequestRecovery = eventName === 'oauthReauthenticationRequired';
|
||||
let recoveryPhase: OAuthRecoveryPhase = 'silent-refresh';
|
||||
let eventHandling: Promise<void> | null = null;
|
||||
|
||||
const handleOAuthEvent = async (data: OAuthRequiredEvent) => {
|
||||
logger.info(`${this.logPrefix} oauthRequired event received`);
|
||||
|
||||
if (!data.skipSilentRefresh && this.shouldAttemptSilentTokenRefresh(data)) {
|
||||
const refreshedTokens = await this.attemptSilentTokenRefresh();
|
||||
if (refreshedTokens) {
|
||||
connection.setOAuthTokens(refreshedTokens);
|
||||
connection.emit('oauthHandled');
|
||||
if (this.connectionReady) {
|
||||
const emitted = connection.emit('oauthReauthenticationRequired', {
|
||||
...data,
|
||||
skipSilentRefresh: data.skipSilentRefresh,
|
||||
});
|
||||
if (emitted) {
|
||||
return;
|
||||
}
|
||||
logger.info(`${this.logPrefix} Cached connection requires a live OAuth request handler`);
|
||||
connection.emit('oauthFailed', new Error('OAuth reauthentication required'));
|
||||
return;
|
||||
}
|
||||
|
||||
if (isRequestRecovery && recoveryPhase === 'terminal') {
|
||||
logger.warn(`${this.logPrefix} OAuth recovery phase budget exhausted`);
|
||||
connection.emit('oauthFailed', new Error('OAuth recovery phase budget exhausted'));
|
||||
return;
|
||||
}
|
||||
|
||||
if (!isRequestRecovery || recoveryPhase === 'silent-refresh') {
|
||||
recoveryPhase = 'interactive';
|
||||
if (!data.skipSilentRefresh && this.shouldAttemptSilentTokenRefresh(data)) {
|
||||
const refreshedTokens = await this.attemptSilentTokenRefresh();
|
||||
if (refreshedTokens) {
|
||||
connection.setOAuthTokens(refreshedTokens);
|
||||
connection.emit('oauthHandled', 'silent-refresh' satisfies t.OAuthHandledSource);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (isRequestRecovery) {
|
||||
recoveryPhase = 'terminal';
|
||||
}
|
||||
|
||||
// Silent refresh failed and we're about to fall through to interactive
|
||||
|
|
@ -1007,21 +1040,6 @@ export class MCPConnectionFactory {
|
|||
// window in `handleOAuthRequired`).
|
||||
await this.invalidateCompletedOAuthFlow();
|
||||
|
||||
if (this.connectionReady) {
|
||||
const emitted = connection.emit('oauthReauthenticationRequired', {
|
||||
...data,
|
||||
skipSilentRefresh: true,
|
||||
});
|
||||
if (emitted) {
|
||||
return;
|
||||
}
|
||||
logger.info(
|
||||
`${this.logPrefix} Silent refresh did not recover cached connection; requiring fresh OAuth prompt`,
|
||||
);
|
||||
connection.emit('oauthFailed', new Error('OAuth reauthentication required'));
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.returnOnOAuth) {
|
||||
try {
|
||||
const config = this.serverConfig;
|
||||
|
|
@ -1098,15 +1116,10 @@ export class MCPConnectionFactory {
|
|||
|
||||
// Start monitoring in background — createFlow will find the existing PENDING state
|
||||
// written by initFlow above, so metadata arg is unused (pass {} to make that explicit)
|
||||
this.flowManager!.createFlow(newFlowId, 'mcp_oauth', {}, this.signal).catch(
|
||||
async (error) => {
|
||||
logger.debug(`${this.logPrefix} OAuth flow monitor ended`, error);
|
||||
await this.clearStaleClientIfRejected(
|
||||
flowMetadata.reusedClientCredentialSetId,
|
||||
error,
|
||||
);
|
||||
},
|
||||
);
|
||||
this.flowManager!.createFlow(newFlowId, 'mcp_oauth', {}).catch(async (error) => {
|
||||
logger.debug(`${this.logPrefix} OAuth flow monitor ended`, error);
|
||||
await this.clearStaleClientIfRejected(flowMetadata.reusedClientCredentialSetId, error);
|
||||
});
|
||||
|
||||
if (this.oauthStart) {
|
||||
logger.info(`${this.logPrefix} OAuth flow started, issuing authorization URL`);
|
||||
|
|
@ -1187,7 +1200,7 @@ export class MCPConnectionFactory {
|
|||
|
||||
// Only emit oauthHandled if we actually got tokens (OAuth succeeded)
|
||||
if (result?.tokens) {
|
||||
connection.emit('oauthHandled');
|
||||
connection.emit('oauthHandled', 'interactive' satisfies t.OAuthHandledSource);
|
||||
} else {
|
||||
await this.clearStaleClientIfRejected(result?.reusedClientCredentialSetId, result?.error);
|
||||
logger.warn(`${this.logPrefix} OAuth failed, emitting oauthFailed event`);
|
||||
|
|
@ -1195,6 +1208,23 @@ export class MCPConnectionFactory {
|
|||
}
|
||||
};
|
||||
|
||||
const oauthHandler = (data: OAuthRequiredEvent): Promise<void> => {
|
||||
if (!isRequestRecovery) {
|
||||
return handleOAuthEvent(data);
|
||||
}
|
||||
if (eventHandling) {
|
||||
return eventHandling;
|
||||
}
|
||||
|
||||
const handling = handleOAuthEvent(data).finally(() => {
|
||||
if (eventHandling === handling) {
|
||||
eventHandling = null;
|
||||
}
|
||||
});
|
||||
eventHandling = handling;
|
||||
return handling;
|
||||
};
|
||||
|
||||
connection.on(eventName, oauthHandler);
|
||||
|
||||
return () => {
|
||||
|
|
@ -1279,7 +1309,7 @@ export class MCPConnectionFactory {
|
|||
throw error;
|
||||
}
|
||||
|
||||
if (this.useOAuth && this.isOAuthError(error)) {
|
||||
if (this.useOAuth && isOAuthAuthenticationError(error)) {
|
||||
logger.info(`${this.logPrefix} OAuth required, stopping connection attempts`);
|
||||
throw error;
|
||||
}
|
||||
|
|
@ -1332,48 +1362,6 @@ export class MCPConnectionFactory {
|
|||
return false;
|
||||
}
|
||||
|
||||
// Determines if an error indicates OAuth authentication is required
|
||||
private isOAuthError(error: unknown): boolean {
|
||||
if (!error || typeof error !== 'object') {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check for error code
|
||||
if ('code' in error) {
|
||||
const code = (error as { code?: number }).code;
|
||||
if (code === 401 || code === 403) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// Check message for various auth error indicators
|
||||
if ('message' in error && typeof error.message === 'string') {
|
||||
const message = error.message.toLowerCase();
|
||||
// Check for 401 status
|
||||
if (message.includes('401') || message.includes('non-200 status code (401)')) {
|
||||
return true;
|
||||
}
|
||||
// Check for invalid_token (OAuth servers return this for expired/revoked tokens)
|
||||
if (message.includes('invalid_token')) {
|
||||
return true;
|
||||
}
|
||||
// Check for invalid_grant (OAuth servers return this for expired/revoked grants)
|
||||
if (message.includes('invalid_grant')) {
|
||||
return true;
|
||||
}
|
||||
// Check for authentication required
|
||||
if (message.includes('authentication required') || message.includes('unauthorized')) {
|
||||
return true;
|
||||
}
|
||||
// Check for missing authorization values (e.g., Amazon Ads MCP returns HTTP 400 with this)
|
||||
if (message.includes('no authorization')) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/** Manages OAuth flow initiation and completion */
|
||||
protected async handleOAuthRequired(): Promise<{
|
||||
tokens: MCPOAuthTokens | null;
|
||||
|
|
@ -1437,7 +1425,7 @@ export class MCPConnectionFactory {
|
|||
|
||||
reusedStoredClient = flowMeta?.reusedStoredClient === true;
|
||||
reusedClientCredentialSetId = flowMeta?.reusedClientCredentialSetId;
|
||||
const tokens = await this.flowManager.createFlow(flowId, 'mcp_oauth', {}, this.signal);
|
||||
const tokens = await this.waitForSharedOAuthFlow(flowId);
|
||||
if (typeof this.oauthEnd === 'function') {
|
||||
await this.oauthEnd();
|
||||
}
|
||||
|
|
@ -1542,7 +1530,7 @@ export class MCPConnectionFactory {
|
|||
|
||||
// createFlow will find the existing PENDING state written by initFlow above,
|
||||
// so metadata arg is unused (pass {} to make that explicit)
|
||||
const tokens = await this.flowManager.createFlow(newFlowId, 'mcp_oauth', {}, this.signal);
|
||||
const tokens = await this.waitForSharedOAuthFlow(newFlowId);
|
||||
if (typeof this.oauthEnd === 'function') {
|
||||
await this.oauthEnd();
|
||||
}
|
||||
|
|
@ -1562,4 +1550,40 @@ export class MCPConnectionFactory {
|
|||
return { tokens: null, reusedStoredClient, reusedClientCredentialSetId, error };
|
||||
}
|
||||
}
|
||||
|
||||
private waitForSharedOAuthFlow(flowId: string): Promise<MCPOAuthTokens | null> {
|
||||
const flow = this.flowManager!.createFlow(flowId, 'mcp_oauth', {});
|
||||
const signal = this.signal;
|
||||
if (!signal) {
|
||||
return flow;
|
||||
}
|
||||
|
||||
return new Promise<MCPOAuthTokens | null>((resolve, reject) => {
|
||||
const cleanup = () => signal.removeEventListener('abort', onAbort);
|
||||
const onAbort = () => {
|
||||
cleanup();
|
||||
reject(
|
||||
signal.reason instanceof Error ? signal.reason : new Error('MCP OAuth flow wait aborted'),
|
||||
);
|
||||
};
|
||||
|
||||
flow.then(
|
||||
(tokens) => {
|
||||
cleanup();
|
||||
resolve(tokens);
|
||||
},
|
||||
(error: unknown) => {
|
||||
cleanup();
|
||||
reject(error);
|
||||
},
|
||||
);
|
||||
|
||||
if (signal.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
|
||||
signal.addEventListener('abort', onAbort, { once: true });
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,9 +28,11 @@ import { UserConnectionManager } from './UserConnectionManager';
|
|||
import { ConnectionsRepository } from './ConnectionsRepository';
|
||||
import { MCPConnectionFactory } from './MCPConnectionFactory';
|
||||
import { processMCPEnv, isPluginSourced } from '~/utils/env';
|
||||
import { OAuthLifecycleRelay } from './oauth/pending';
|
||||
import { preProcessGraphTokens } from '~/utils/graph';
|
||||
import { formatToolContent } from './parsers';
|
||||
import { MCPConnection } from './connection';
|
||||
import { mcpConfig } from './mcpConfig';
|
||||
|
||||
function createOboToolCallErrorMessage(
|
||||
logPrefix: string,
|
||||
|
|
@ -48,12 +50,35 @@ function createOboToolCallErrorMessage(
|
|||
return `${logPrefix} ${error.userMessage} Cannot execute tool ${toolName}. ${failureSuffix}`;
|
||||
}
|
||||
|
||||
class OAuthRecoveryTakeoverRequired extends Error {}
|
||||
|
||||
type OAuthReconnectResult =
|
||||
| { connected: true }
|
||||
| {
|
||||
connected: false;
|
||||
error: unknown;
|
||||
oauthHandled: boolean;
|
||||
source?: t.OAuthHandledSource;
|
||||
};
|
||||
|
||||
const OAUTH_RECOVERY_RECONNECT_ATTEMPTS = 3;
|
||||
const OAUTH_RECOVERY_RECONNECT_DELAY_MS = 2000;
|
||||
|
||||
/**
|
||||
* Centralized manager for MCP server connections and tool execution.
|
||||
* Extends UserConnectionManager to handle both app-level and user-specific connections.
|
||||
*/
|
||||
export class MCPManager extends UserConnectionManager {
|
||||
private static instance: MCPManager | null;
|
||||
private readonly oauthRecoveries = new WeakMap<
|
||||
MCPConnection,
|
||||
{
|
||||
promise: Promise<void>;
|
||||
callbacks: OAuthLifecycleRelay;
|
||||
allowsTakeover: boolean;
|
||||
takeoverClaimed?: boolean;
|
||||
}
|
||||
>();
|
||||
|
||||
/** Creates and initializes the singleton MCPManager instance */
|
||||
public static async createInstance(configs: t.MCPServers): Promise<MCPManager> {
|
||||
|
|
@ -75,6 +100,117 @@ export class MCPManager extends UserConnectionManager {
|
|||
this.appConnections = new ConnectionsRepository(undefined);
|
||||
}
|
||||
|
||||
public override async getUserConnection(
|
||||
opts: t.UserMCPConnectionOptions,
|
||||
): Promise<MCPConnection> {
|
||||
const userId = opts.user?.id;
|
||||
if (opts.forceNew || !userId) {
|
||||
return super.getUserConnection(opts);
|
||||
}
|
||||
|
||||
const connectionKey = `${userId}:${opts.serverName}`;
|
||||
const requestConnection = opts.requestScopedConnections?.connections.get(connectionKey) as
|
||||
| MCPConnection
|
||||
| undefined;
|
||||
const connection = requestConnection ?? this.userConnections.get(userId)?.get(opts.serverName);
|
||||
const recovery = connection ? this.oauthRecoveries.get(connection) : undefined;
|
||||
const providedConfigIsNewer =
|
||||
connection != null &&
|
||||
opts.serverConfig?.updatedAt != null &&
|
||||
connection.isStale(opts.serverConfig.updatedAt);
|
||||
if (recovery && !providedConfigIsNewer) {
|
||||
if (recovery.callbacks) {
|
||||
await recovery.callbacks.add({
|
||||
oauthStart: opts.oauthStart,
|
||||
oauthEnd: opts.oauthEnd,
|
||||
flowManager: opts.flowManager,
|
||||
userId,
|
||||
serverName: opts.serverName,
|
||||
});
|
||||
}
|
||||
await this.waitForActiveRecovery(recovery.promise, opts.signal);
|
||||
}
|
||||
|
||||
return super.getUserConnection(opts);
|
||||
}
|
||||
|
||||
/** Runs work against a user connection while preventing recovery from replacing its SDK client. */
|
||||
public async withUserConnectionLease<TResult>(
|
||||
opts: t.UserMCPConnectionOptions,
|
||||
operation: (connection: MCPConnection) => Promise<TResult>,
|
||||
): Promise<TResult> {
|
||||
while (true) {
|
||||
const connection = await this.getUserConnection(opts);
|
||||
this.retainConnection(connection);
|
||||
const recovery = this.oauthRecoveries.get(connection)?.promise;
|
||||
if (recovery) {
|
||||
await this.releaseConnection(connection);
|
||||
await this.waitForActiveRecovery(recovery, opts.signal);
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
return await operation(connection);
|
||||
} finally {
|
||||
await this.releaseConnection(connection);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private waitForActiveRecovery(recovery: Promise<void>, signal?: AbortSignal): Promise<void> {
|
||||
if (!signal) {
|
||||
return recovery;
|
||||
}
|
||||
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
const onRecoveryResolved = () => {
|
||||
signal.removeEventListener('abort', onAbort);
|
||||
resolve();
|
||||
};
|
||||
const onRecoveryRejected = (error: unknown) => {
|
||||
signal.removeEventListener('abort', onAbort);
|
||||
reject(error);
|
||||
};
|
||||
const onAbort = () => {
|
||||
signal.removeEventListener('abort', onAbort);
|
||||
const reason = signal.reason;
|
||||
reject(reason instanceof Error ? reason : new Error('OAuth recovery wait aborted'));
|
||||
};
|
||||
|
||||
if (signal.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
|
||||
signal.addEventListener('abort', onAbort, { once: true });
|
||||
recovery.then(onRecoveryResolved, onRecoveryRejected);
|
||||
});
|
||||
}
|
||||
|
||||
protected override getActiveConnectionRecovery(
|
||||
connection: MCPConnection,
|
||||
): Promise<void> | undefined {
|
||||
return this.oauthRecoveries.get(connection)?.promise;
|
||||
}
|
||||
|
||||
protected override waitForConnectionRecovery(
|
||||
recovery: Promise<void>,
|
||||
signal?: AbortSignal,
|
||||
): Promise<void> {
|
||||
return this.waitForActiveRecovery(recovery, signal);
|
||||
}
|
||||
|
||||
private claimRecoveryTakeover(recovery: {
|
||||
allowsTakeover: boolean;
|
||||
takeoverClaimed?: boolean;
|
||||
}): boolean {
|
||||
if (!recovery.allowsTakeover || recovery.takeoverClaimed) {
|
||||
return false;
|
||||
}
|
||||
recovery.takeoverClaimed = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
/** Retrieves an app-level or user-specific connection based on provided arguments */
|
||||
public async getConnection(
|
||||
args: {
|
||||
|
|
@ -302,53 +438,64 @@ export class MCPManager extends UserConnectionManager {
|
|||
};
|
||||
}
|
||||
|
||||
const userConnections = this.getUserConnections(userId);
|
||||
if (!userConnections || userConnections.size === 0) {
|
||||
return { tools: null };
|
||||
}
|
||||
if (!userConnections.has(serverName)) {
|
||||
return { tools: null };
|
||||
}
|
||||
let awaitedRecovery: Promise<void> | undefined;
|
||||
while (true) {
|
||||
const userConnections = this.getUserConnections(userId);
|
||||
const connection = userConnections?.get(serverName);
|
||||
if (!connection) {
|
||||
return { tools: null };
|
||||
}
|
||||
|
||||
const connection = userConnections.get(serverName)!;
|
||||
if (effectiveConfig == null) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
if (effectiveConfig == null) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
const connectionConfigGeneration = this.getToolConfigGeneration(connection);
|
||||
const effectiveConfigGeneration = getMCPAppToolsPublicationGeneration(effectiveConfig);
|
||||
if (
|
||||
connectionConfigGeneration != null &&
|
||||
effectiveConfigGeneration != null &&
|
||||
connectionConfigGeneration !== effectiveConfigGeneration
|
||||
) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
const publicationGeneration = this.getToolPublicationGeneration(connection);
|
||||
const currentGeneration = await getMCPToolsChangedGeneration({ userId, serverName });
|
||||
if (
|
||||
publicationGeneration != null &&
|
||||
currentGeneration != null &&
|
||||
publicationGeneration !== currentGeneration
|
||||
) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
|
||||
this.retainConnection(connection);
|
||||
const recovery = this.oauthRecoveries.get(connection)?.promise;
|
||||
if (recovery && recovery !== awaitedRecovery) {
|
||||
awaitedRecovery = recovery;
|
||||
await this.releaseConnection(connection);
|
||||
await this.waitForConnectionRecovery(recovery);
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
const tools = await MCPServerInspector.getToolFunctions(serverName, connection);
|
||||
const generationAfterFetch = await getMCPToolsChangedGeneration({ userId, serverName });
|
||||
if (
|
||||
publicationGeneration != null &&
|
||||
generationAfterFetch != null &&
|
||||
publicationGeneration !== generationAfterFetch
|
||||
) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
return { tools, publicationGeneration };
|
||||
} finally {
|
||||
await this.releaseConnection(connection);
|
||||
}
|
||||
}
|
||||
const connectionConfigGeneration = this.getToolConfigGeneration(connection);
|
||||
const effectiveConfigGeneration = getMCPAppToolsPublicationGeneration(effectiveConfig);
|
||||
if (
|
||||
connectionConfigGeneration != null &&
|
||||
effectiveConfigGeneration != null &&
|
||||
connectionConfigGeneration !== effectiveConfigGeneration
|
||||
) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
const publicationGeneration = this.getToolPublicationGeneration(connection);
|
||||
const currentGeneration = await getMCPToolsChangedGeneration({ userId, serverName });
|
||||
if (
|
||||
publicationGeneration != null &&
|
||||
currentGeneration != null &&
|
||||
publicationGeneration !== currentGeneration
|
||||
) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
const tools = await MCPServerInspector.getToolFunctions(serverName, connection);
|
||||
const generationAfterFetch = await getMCPToolsChangedGeneration({ userId, serverName });
|
||||
if (
|
||||
publicationGeneration != null &&
|
||||
generationAfterFetch != null &&
|
||||
publicationGeneration !== generationAfterFetch
|
||||
) {
|
||||
await this.disconnectUserConnection(userId, serverName);
|
||||
return { tools: null };
|
||||
}
|
||||
return {
|
||||
tools,
|
||||
publicationGeneration,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[getServerToolFunctions] Error getting tool functions for server ${serverName}`,
|
||||
|
|
@ -422,6 +569,198 @@ ${formattedInstructions}
|
|||
Please follow these instructions when using tools from the respective MCP servers.`;
|
||||
}
|
||||
|
||||
private async recoverOAuthConnection(
|
||||
connection: MCPConnection,
|
||||
error: unknown,
|
||||
serverName: string,
|
||||
userId: string,
|
||||
attachSharedOAuthHandler: (relay: OAuthLifecycleRelay) => () => void,
|
||||
oauthStart: t.OAuthStartHandler | undefined,
|
||||
oauthEnd: (() => Promise<void>) | undefined,
|
||||
flowManager: FlowStateManager<MCPOAuthTokens | null>,
|
||||
signal?: AbortSignal,
|
||||
allowsTakeover = true,
|
||||
): Promise<void> {
|
||||
const existingRecovery = this.oauthRecoveries.get(connection);
|
||||
if (existingRecovery) {
|
||||
if (existingRecovery.callbacks) {
|
||||
await existingRecovery.callbacks.add({
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
flowManager,
|
||||
userId,
|
||||
serverName,
|
||||
});
|
||||
}
|
||||
try {
|
||||
return await this.waitForActiveRecovery(existingRecovery.promise, signal);
|
||||
} catch (recoveryError) {
|
||||
if (signal?.aborted) {
|
||||
throw recoveryError;
|
||||
}
|
||||
if (!allowsTakeover || !this.claimRecoveryTakeover(existingRecovery)) {
|
||||
throw recoveryError;
|
||||
}
|
||||
if (this.oauthRecoveries.get(connection) === existingRecovery) {
|
||||
this.oauthRecoveries.delete(connection);
|
||||
}
|
||||
throw new OAuthRecoveryTakeoverRequired();
|
||||
}
|
||||
}
|
||||
|
||||
const callbacks = new OAuthLifecycleRelay({
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
logPrefix: `[MCP][User: ${userId}][${serverName}]`,
|
||||
});
|
||||
const recovery = Promise.resolve().then(async () => {
|
||||
const cleanupRequestOAuthHandler = attachSharedOAuthHandler(callbacks);
|
||||
try {
|
||||
await this.waitForOAuthRecovery(connection, () =>
|
||||
connection.emit('oauthReauthenticationRequired', {
|
||||
serverName,
|
||||
error,
|
||||
serverUrl: connection.url,
|
||||
userId,
|
||||
}),
|
||||
);
|
||||
await this.connectAfterOAuthRecovery(connection, async (connectError) => {
|
||||
await this.waitForOAuthRecovery(connection, () =>
|
||||
connection.emit('oauthReauthenticationRequired', {
|
||||
serverName,
|
||||
error: connectError,
|
||||
serverUrl: connection.url,
|
||||
userId,
|
||||
skipSilentRefresh: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
} finally {
|
||||
cleanupRequestOAuthHandler();
|
||||
}
|
||||
});
|
||||
|
||||
const recoveryEntry = { promise: recovery, callbacks, allowsTakeover, takeoverClaimed: false };
|
||||
this.oauthRecoveries.set(connection, recoveryEntry);
|
||||
this.holdDeferredConnectionDisposal(connection);
|
||||
const clearRecovery = () => {
|
||||
if (this.oauthRecoveries.get(connection) === recoveryEntry) {
|
||||
this.oauthRecoveries.delete(connection);
|
||||
}
|
||||
};
|
||||
const releaseRecoveryDisposal = () => this.releaseDeferredConnectionDisposal(connection);
|
||||
void recovery.then(clearRecovery, clearRecovery);
|
||||
void recovery.then(releaseRecoveryDisposal, releaseRecoveryDisposal);
|
||||
await this.waitForActiveRecovery(recovery, signal);
|
||||
}
|
||||
|
||||
private async connectAfterOAuthRecovery(
|
||||
connection: MCPConnection,
|
||||
requestInteractiveRecovery: (error: unknown) => Promise<void>,
|
||||
): Promise<void> {
|
||||
await this.waitForConnectionBorrowersToDrain(connection);
|
||||
const firstAttempt = await this.connectWithTransientRetries(connection);
|
||||
if (firstAttempt.connected) {
|
||||
return;
|
||||
}
|
||||
if (!firstAttempt.oauthHandled) {
|
||||
throw firstAttempt.error;
|
||||
}
|
||||
if (firstAttempt.source === 'silent-refresh') {
|
||||
await requestInteractiveRecovery(firstAttempt.error);
|
||||
}
|
||||
|
||||
const secondAttempt = await this.connectWithTransientRetries(connection);
|
||||
if (!secondAttempt.connected) {
|
||||
throw secondAttempt.error;
|
||||
}
|
||||
}
|
||||
|
||||
private async connectWithTransientRetries(
|
||||
connection: MCPConnection,
|
||||
): Promise<OAuthReconnectResult> {
|
||||
let result: OAuthReconnectResult | undefined;
|
||||
for (let attempt = 1; attempt <= OAUTH_RECOVERY_RECONNECT_ATTEMPTS; attempt++) {
|
||||
result = await this.connectOnceAfterOAuth(connection);
|
||||
if (
|
||||
result.connected ||
|
||||
result.oauthHandled ||
|
||||
connection.isOAuthAuthenticationError(result.error) ||
|
||||
attempt === OAUTH_RECOVERY_RECONNECT_ATTEMPTS
|
||||
) {
|
||||
return result;
|
||||
}
|
||||
await this.waitForOAuthReconnectRetry(attempt);
|
||||
}
|
||||
return result!;
|
||||
}
|
||||
|
||||
private waitForOAuthReconnectRetry(attempt: number): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
setTimeout(resolve, OAUTH_RECOVERY_RECONNECT_DELAY_MS * attempt);
|
||||
});
|
||||
}
|
||||
|
||||
private async connectOnceAfterOAuth(connection: MCPConnection): Promise<OAuthReconnectResult> {
|
||||
let oauthHandled = false;
|
||||
let source: t.OAuthHandledSource | undefined;
|
||||
const handleOAuth = (handledSource?: t.OAuthHandledSource) => {
|
||||
oauthHandled = true;
|
||||
source = handledSource;
|
||||
};
|
||||
connection.on('oauthHandled', handleOAuth);
|
||||
try {
|
||||
await connection.connect();
|
||||
return { connected: true };
|
||||
} catch (error) {
|
||||
return { connected: false, error, oauthHandled, source };
|
||||
} finally {
|
||||
connection.off('oauthHandled', handleOAuth);
|
||||
}
|
||||
}
|
||||
|
||||
private waitForOAuthRecovery(
|
||||
connection: MCPConnection,
|
||||
requestRecovery: () => boolean,
|
||||
): Promise<void> {
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
const cleanup = () => {
|
||||
clearTimeout(timeout);
|
||||
connection.off('oauthHandled', handleSuccess);
|
||||
connection.off('oauthFailed', handleFailure);
|
||||
};
|
||||
const handleSuccess = () => {
|
||||
cleanup();
|
||||
resolve();
|
||||
};
|
||||
const handleFailure = (oauthError: Error) => {
|
||||
cleanup();
|
||||
reject(oauthError);
|
||||
};
|
||||
const timeout = setTimeout(() => {
|
||||
cleanup();
|
||||
reject(new Error(`OAuth recovery timeout after ${mcpConfig.OAUTH_HANDLING_TIMEOUT}ms`));
|
||||
}, mcpConfig.OAUTH_HANDLING_TIMEOUT);
|
||||
|
||||
connection.once('oauthHandled', handleSuccess);
|
||||
connection.once('oauthFailed', handleFailure);
|
||||
|
||||
let recoveryRequested: boolean;
|
||||
try {
|
||||
recoveryRequested = requestRecovery();
|
||||
} catch (error) {
|
||||
cleanup();
|
||||
reject(error);
|
||||
return;
|
||||
}
|
||||
if (recoveryRequested) {
|
||||
return;
|
||||
}
|
||||
cleanup();
|
||||
reject(new Error('OAuth recovery requested without an active request handler'));
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Calls a tool on an MCP server, using either a user-specific connection
|
||||
* (if userId is provided) or an app-level connection. Updates the last activity timestamp
|
||||
|
|
@ -469,184 +808,351 @@ Please follow these instructions when using tools from the respective MCP server
|
|||
oboTokenResolver?: OboTokenResolver;
|
||||
oboTrustChecker?: OboTrustChecker;
|
||||
}): Promise<t.FormattedToolResponse> {
|
||||
/** User-specific connection */
|
||||
let connection: MCPConnection | undefined;
|
||||
let cleanupRequestOAuthHandler: (() => void) | undefined;
|
||||
let disconnectAfterCall = false;
|
||||
const userId = user?.id;
|
||||
const logPrefix = userId ? `[MCP][User: ${userId}][${serverName}]` : `[MCP][${serverName}]`;
|
||||
|
||||
try {
|
||||
connection = await this.getConnection({
|
||||
serverName,
|
||||
user,
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
oboTokenResolver,
|
||||
oboTrustChecker,
|
||||
graphTokenResolver,
|
||||
signal: options?.signal,
|
||||
customUserVars,
|
||||
requestBody,
|
||||
requestScopedConnections,
|
||||
serverConfig: providedConfig,
|
||||
});
|
||||
|
||||
if (!(await connection.isConnected())) {
|
||||
/** May happen if getUserConnection failed silently or app connection dropped */
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError, // Use InternalError for connection issues
|
||||
`${logPrefix} Connection is not active. Cannot execute tool ${toolName}.`,
|
||||
);
|
||||
}
|
||||
|
||||
const registry = MCPServersRegistry.getInstance();
|
||||
const rawConfig = providedConfig ?? (await registry.getServerConfig(serverName, userId));
|
||||
if (!rawConfig) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidRequest,
|
||||
`${logPrefix} Configuration for server "${serverName}" not found.`,
|
||||
);
|
||||
}
|
||||
const isDbSourced = isUserSourced(rawConfig);
|
||||
const ephemeralConnection = !!userId && requiresEphemeralUserConnection(rawConfig);
|
||||
disconnectAfterCall = ephemeralConnection && !requestScopedConnections;
|
||||
|
||||
/**
|
||||
* Pre-process Graph token placeholders (async) before the synchronous processMCPEnv pass.
|
||||
* Plugin-sourced configs are excluded for the same reason processMCPEnv excludes them:
|
||||
* a placeholder a plugin authored must never resolve against the user's Graph token.
|
||||
*/
|
||||
const graphProcessedConfig =
|
||||
isDbSourced || isPluginSourced(rawConfig)
|
||||
? (rawConfig as t.MCPOptions)
|
||||
: await preProcessGraphTokens(rawConfig as t.MCPOptions, {
|
||||
user,
|
||||
graphTokenResolver,
|
||||
scopes: process.env.GRAPH_API_SCOPES,
|
||||
});
|
||||
const currentOptions = processMCPEnv({
|
||||
user,
|
||||
body: requestBody,
|
||||
dbSourced: isDbSourced,
|
||||
options: graphProcessedConfig,
|
||||
customUserVars,
|
||||
});
|
||||
|
||||
const resolvedHeaders: Record<string, string> =
|
||||
'headers' in currentOptions ? { ...(currentOptions.headers || {}) } : {};
|
||||
|
||||
/** Refresh OBO token on each tool call to ensure it's current */
|
||||
const oboConfig = rawConfig.obo;
|
||||
if (oboConfig && oboTokenResolver && user) {
|
||||
const oboTrusted = oboTrustChecker
|
||||
? await oboTrustChecker({
|
||||
source: rawConfig.source,
|
||||
author: rawConfig.author,
|
||||
dbId: rawConfig.dbId,
|
||||
})
|
||||
: true;
|
||||
if (!oboTrusted) {
|
||||
logger.warn(
|
||||
`${logPrefix} OBO config not trusted (author lacks ${PermissionTypes.MCP_SERVERS}.${Permissions.CONFIGURE_OBO}); refusing to mint a downstream token`,
|
||||
);
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
`${logPrefix} OBO is not permitted for server "${serverName}". The user who configured it no longer has permission to use OBO.`,
|
||||
);
|
||||
this.bindRequestScopedConnectionStore(requestScopedConnections);
|
||||
let recoveryTakeoverConsumed = false;
|
||||
while (true) {
|
||||
/** User-specific connection */
|
||||
let connection: MCPConnection | undefined;
|
||||
let connectionRetained = false;
|
||||
let deferredDisposalHeld = false;
|
||||
let attachSharedOAuthHandler: ((relay: OAuthLifecycleRelay) => () => void) | undefined;
|
||||
let disposeAfterCall = false;
|
||||
const retainConnectionLease = () => {
|
||||
if (!connection || connectionRetained) {
|
||||
return;
|
||||
}
|
||||
let oboTokens: MCPOAuthTokens;
|
||||
this.retainConnection(connection);
|
||||
connectionRetained = true;
|
||||
};
|
||||
const releaseConnectionLease = async (preserveDisposalHold = false) => {
|
||||
if (!connection || !connectionRetained) {
|
||||
return;
|
||||
}
|
||||
if (deferredDisposalHeld && !preserveDisposalHold) {
|
||||
await this.releaseDeferredConnectionDisposal(connection);
|
||||
deferredDisposalHeld = false;
|
||||
}
|
||||
connectionRetained = false;
|
||||
await this.releaseConnection(connection);
|
||||
};
|
||||
const waitForRecoveryWithoutLease = async (startRecovery: () => Promise<void>) => {
|
||||
const recovery = startRecovery();
|
||||
// Keep an eviction marker across the temporary lease gap and transfer that
|
||||
// responsibility back to this caller after recovery. An unrelated final
|
||||
// borrower may disconnect the old client, but cannot consume the marker.
|
||||
if (!deferredDisposalHeld) {
|
||||
this.holdDeferredConnectionDisposal(connection!);
|
||||
deferredDisposalHeld = true;
|
||||
}
|
||||
await releaseConnectionLease(true);
|
||||
try {
|
||||
oboTokens = await resolveOboToken(user, oboConfig, oboTokenResolver);
|
||||
} catch (error) {
|
||||
if (error instanceof OboTokenResolutionError) {
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
createOboToolCallErrorMessage(logPrefix, toolName, error),
|
||||
);
|
||||
}
|
||||
throw error;
|
||||
await recovery;
|
||||
} finally {
|
||||
retainConnectionLease();
|
||||
}
|
||||
};
|
||||
|
||||
if (!oboTokens.access_token) {
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
`${logPrefix} OBO token refresh failed. Cannot execute tool ${toolName}. Re-authenticate the user and retry.`,
|
||||
);
|
||||
}
|
||||
resolvedHeaders['Authorization'] = `Bearer ${oboTokens.access_token}`;
|
||||
}
|
||||
if (userId && user && oauthStart && flowManager && isOAuthServer(currentOptions)) {
|
||||
const { allowedDomains, allowedAddresses, useSSRFProtection } =
|
||||
await registry.resolveAllowlists({ userId, role: user?.role });
|
||||
cleanupRequestOAuthHandler = MCPConnectionFactory.attachRequestOAuthHandler(
|
||||
{
|
||||
try {
|
||||
let awaitedCheckoutRecovery: Promise<void> | undefined;
|
||||
while (true) {
|
||||
connection = await this.getConnection({
|
||||
serverName,
|
||||
serverConfig: currentOptions,
|
||||
dbSourced: isDbSourced,
|
||||
skipEnvProcessing: true,
|
||||
useSSRFProtection,
|
||||
allowedDomains,
|
||||
allowedAddresses,
|
||||
},
|
||||
{
|
||||
useOAuth: true,
|
||||
user,
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
signal: options?.signal,
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
oboTokenResolver,
|
||||
oboTrustChecker,
|
||||
graphTokenResolver,
|
||||
signal: options?.signal,
|
||||
customUserVars,
|
||||
requestBody,
|
||||
},
|
||||
connection,
|
||||
);
|
||||
}
|
||||
|
||||
connection.setRequestHeaders(resolvedHeaders);
|
||||
|
||||
const result = await connection.client.request(
|
||||
{
|
||||
method: 'tools/call',
|
||||
params: {
|
||||
name: toolName,
|
||||
arguments: toolArguments,
|
||||
},
|
||||
},
|
||||
CallToolResultSchema,
|
||||
{
|
||||
timeout: connection.timeout,
|
||||
resetTimeoutOnProgress: true,
|
||||
...options,
|
||||
},
|
||||
);
|
||||
const hasPersistentUserConnections =
|
||||
!!userId && (this.userConnections.get(userId)?.size ?? 0) > 0;
|
||||
if (!ephemeralConnection && hasPersistentUserConnections) {
|
||||
await this.updateUserLastActivity(userId);
|
||||
}
|
||||
this.checkIdleConnections();
|
||||
return formatToolContent(result as t.MCPToolCallResponse, provider);
|
||||
} catch (error) {
|
||||
// Log with context and re-throw or handle as needed
|
||||
logger.error(`${logPrefix}[${toolName}] Tool call failed`, error);
|
||||
// Rethrowing allows the caller (createMCPTool) to handle the final user message
|
||||
throw error;
|
||||
} finally {
|
||||
cleanupRequestOAuthHandler?.();
|
||||
// Ephemeral connections are never stored in userConnections, so disconnecting
|
||||
// is the only cleanup needed; removing the map entry here could orphan a
|
||||
// still-connected cached connection from before a config change.
|
||||
if (disconnectAfterCall && connection) {
|
||||
try {
|
||||
await connection.disconnect();
|
||||
} catch (disconnectError) {
|
||||
logger.warn(`${logPrefix}[${toolName}] Failed to disconnect ephemeral connection`, {
|
||||
error: disconnectError,
|
||||
requestScopedConnections,
|
||||
serverConfig: providedConfig,
|
||||
});
|
||||
retainConnectionLease();
|
||||
const checkoutRecovery = this.oauthRecoveries.get(connection);
|
||||
if (!checkoutRecovery || checkoutRecovery.promise === awaitedCheckoutRecovery) {
|
||||
break;
|
||||
}
|
||||
if (checkoutRecovery.callbacks) {
|
||||
await checkoutRecovery.callbacks.add({
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
flowManager,
|
||||
userId: userId!,
|
||||
serverName,
|
||||
});
|
||||
}
|
||||
awaitedCheckoutRecovery = checkoutRecovery.promise;
|
||||
await releaseConnectionLease();
|
||||
try {
|
||||
await this.waitForConnectionRecovery(checkoutRecovery.promise, options?.signal);
|
||||
} catch (recoveryError) {
|
||||
if (
|
||||
options?.signal?.aborted ||
|
||||
recoveryTakeoverConsumed ||
|
||||
!this.claimRecoveryTakeover(checkoutRecovery)
|
||||
) {
|
||||
throw recoveryError;
|
||||
}
|
||||
recoveryTakeoverConsumed = true;
|
||||
if (this.oauthRecoveries.get(connection) === checkoutRecovery) {
|
||||
this.oauthRecoveries.delete(connection);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
const connectionIsActive = await connection.isConnected();
|
||||
const connectionCheckError = connectionIsActive
|
||||
? undefined
|
||||
: connection.getLastConnectionCheckError();
|
||||
|
||||
if (
|
||||
!connectionIsActive &&
|
||||
(!userId || !connection.isOAuthAuthenticationError(connectionCheckError))
|
||||
) {
|
||||
/** May happen if getUserConnection failed silently or app connection dropped */
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
`${logPrefix} Connection is not active. Cannot execute tool ${toolName}.`,
|
||||
);
|
||||
}
|
||||
|
||||
const registry = MCPServersRegistry.getInstance();
|
||||
const rawConfig = providedConfig ?? (await registry.getServerConfig(serverName, userId));
|
||||
if (!rawConfig) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidRequest,
|
||||
`${logPrefix} Configuration for server "${serverName}" not found.`,
|
||||
);
|
||||
}
|
||||
const isDbSourced = isUserSourced(rawConfig);
|
||||
const ephemeralConnection = !!userId && requiresEphemeralUserConnection(rawConfig);
|
||||
disposeAfterCall = ephemeralConnection && !requestScopedConnections;
|
||||
|
||||
/** Plugin-authored placeholders must not resolve against the user's Graph token. */
|
||||
const graphProcessedConfig =
|
||||
isDbSourced || isPluginSourced(rawConfig)
|
||||
? (rawConfig as t.MCPOptions)
|
||||
: await preProcessGraphTokens(rawConfig as t.MCPOptions, {
|
||||
user,
|
||||
graphTokenResolver,
|
||||
scopes: process.env.GRAPH_API_SCOPES,
|
||||
});
|
||||
const currentOptions = processMCPEnv({
|
||||
user,
|
||||
body: requestBody,
|
||||
dbSourced: isDbSourced,
|
||||
options: graphProcessedConfig,
|
||||
customUserVars,
|
||||
});
|
||||
|
||||
const resolvedHeaders: Record<string, string> =
|
||||
'headers' in currentOptions ? { ...(currentOptions.headers || {}) } : {};
|
||||
|
||||
/** Refresh OBO token on each tool call to ensure it's current */
|
||||
const oboConfig = rawConfig.obo;
|
||||
if (oboConfig && oboTokenResolver && user) {
|
||||
const oboTrusted = oboTrustChecker
|
||||
? await oboTrustChecker({
|
||||
source: rawConfig.source,
|
||||
author: rawConfig.author,
|
||||
dbId: rawConfig.dbId,
|
||||
})
|
||||
: true;
|
||||
if (!oboTrusted) {
|
||||
logger.warn(
|
||||
`${logPrefix} OBO config not trusted (author lacks ${PermissionTypes.MCP_SERVERS}.${Permissions.CONFIGURE_OBO}); refusing to mint a downstream token`,
|
||||
);
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
`${logPrefix} OBO is not permitted for server "${serverName}". The user who configured it no longer has permission to use OBO.`,
|
||||
);
|
||||
}
|
||||
let oboTokens: MCPOAuthTokens;
|
||||
try {
|
||||
oboTokens = await resolveOboToken(user, oboConfig, oboTokenResolver);
|
||||
} catch (error) {
|
||||
if (error instanceof OboTokenResolutionError) {
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
createOboToolCallErrorMessage(logPrefix, toolName, error),
|
||||
);
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
|
||||
if (!oboTokens.access_token) {
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
`${logPrefix} OBO token refresh failed. Cannot execute tool ${toolName}. Re-authenticate the user and retry.`,
|
||||
);
|
||||
}
|
||||
resolvedHeaders['Authorization'] = `Bearer ${oboTokens.access_token}`;
|
||||
}
|
||||
if (
|
||||
userId &&
|
||||
user &&
|
||||
oauthStart &&
|
||||
flowManager &&
|
||||
(isOAuthServer(currentOptions) || connection.usesOAuth())
|
||||
) {
|
||||
const { allowedDomains, allowedAddresses, useSSRFProtection } =
|
||||
await registry.resolveAllowlists({ userId, role: user?.role });
|
||||
attachSharedOAuthHandler = (relay) =>
|
||||
MCPConnectionFactory.attachRequestOAuthHandler(
|
||||
{
|
||||
serverName,
|
||||
serverConfig: currentOptions,
|
||||
dbSourced: isDbSourced,
|
||||
skipEnvProcessing: true,
|
||||
useSSRFProtection,
|
||||
allowedDomains,
|
||||
allowedAddresses,
|
||||
},
|
||||
{
|
||||
useOAuth: true,
|
||||
user,
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
oauthStart: relay.start,
|
||||
oauthEnd: relay.end,
|
||||
customUserVars,
|
||||
requestBody,
|
||||
},
|
||||
connection!,
|
||||
);
|
||||
}
|
||||
|
||||
connection.setRequestHeaders(resolvedHeaders);
|
||||
|
||||
if (!connectionIsActive) {
|
||||
const requestOAuthHandler = attachSharedOAuthHandler;
|
||||
if (!requestOAuthHandler || !userId) {
|
||||
throw new McpError(
|
||||
ErrorCode.InternalError,
|
||||
`${logPrefix} Connection is not active. Cannot execute tool ${toolName}.`,
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
await waitForRecoveryWithoutLease(() =>
|
||||
this.recoverOAuthConnection(
|
||||
connection!,
|
||||
connectionCheckError,
|
||||
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}] Connection-check OAuth recovery failed`,
|
||||
recoveryError,
|
||||
);
|
||||
throw connectionCheckError;
|
||||
}
|
||||
}
|
||||
|
||||
const requestTool = () =>
|
||||
connection!.client.request(
|
||||
{
|
||||
method: 'tools/call',
|
||||
params: {
|
||||
name: toolName,
|
||||
arguments: toolArguments,
|
||||
},
|
||||
},
|
||||
CallToolResultSchema,
|
||||
{
|
||||
timeout: connection!.timeout,
|
||||
resetTimeoutOnProgress: true,
|
||||
...options,
|
||||
},
|
||||
);
|
||||
|
||||
let result: Awaited<ReturnType<typeof requestTool>>;
|
||||
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,
|
||||
),
|
||||
);
|
||||
} 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();
|
||||
}
|
||||
const hasPersistentUserConnections =
|
||||
!!userId && (this.userConnections.get(userId)?.size ?? 0) > 0;
|
||||
if (!ephemeralConnection && hasPersistentUserConnections) {
|
||||
await this.updateUserLastActivity(userId);
|
||||
}
|
||||
this.checkIdleConnections();
|
||||
return formatToolContent(result as t.MCPToolCallResponse, provider);
|
||||
} catch (error) {
|
||||
if (error instanceof OAuthRecoveryTakeoverRequired) {
|
||||
recoveryTakeoverConsumed = true;
|
||||
continue;
|
||||
}
|
||||
// Log with context and re-throw or handle as needed
|
||||
logger.error(`${logPrefix}[${toolName}] Tool call failed`, error);
|
||||
// Rethrowing allows the caller (createMCPTool) to handle the final user message
|
||||
throw error;
|
||||
} finally {
|
||||
await releaseConnectionLease();
|
||||
// Ephemeral connections are never stored in userConnections, so disposing
|
||||
// is the only cleanup needed; removing the map entry here could orphan a
|
||||
// still-connected cached connection from before a config change.
|
||||
if (disposeAfterCall && connection) {
|
||||
await this.disposeEvictedConnection(
|
||||
connection,
|
||||
`${logPrefix}[${toolName}] Ephemeral connection`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
import { logger, getTenantId } from '@librechat/data-schemas';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js';
|
||||
import type { MCPOAuthFlowMetadata } from '~/mcp/oauth';
|
||||
import type { FlowState } from '~/flow/types';
|
||||
import type * as t from './types';
|
||||
import {
|
||||
cancelMCPToolsChanged,
|
||||
|
|
@ -18,31 +16,19 @@ import {
|
|||
requiresOAuthMachinery,
|
||||
} from './utils';
|
||||
import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry';
|
||||
import { detectOAuthRequirement, MCPOAuthHandler } from '~/mcp/oauth';
|
||||
import { ConnectionsRepository } from '~/mcp/ConnectionsRepository';
|
||||
import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory';
|
||||
import { processMCPEnv, isPluginSourced } from '~/utils/env';
|
||||
import { OAuthLifecycleRelay } from '~/mcp/oauth/pending';
|
||||
import { preProcessGraphTokens } from '~/utils/graph';
|
||||
import { detectOAuthRequirement } from '~/mcp/oauth';
|
||||
import { isMCPDomainAllowed } from '~/auth/domain';
|
||||
import { PENDING_STALE_MS } from '~/flow/manager';
|
||||
import { MCPConnection } from './connection';
|
||||
import { mcpConfig } from './mcpConfig';
|
||||
|
||||
type PendingOAuthStart = {
|
||||
authURL: string;
|
||||
options?: t.OAuthStartOptions;
|
||||
};
|
||||
|
||||
type PendingOAuthState = {
|
||||
oauthStarts: Set<t.OAuthStartHandler>;
|
||||
emittedAuthUrls: WeakMap<t.OAuthStartHandler, string>;
|
||||
primaryOAuthStart?: t.OAuthStartHandler;
|
||||
lastOAuthStart?: PendingOAuthStart;
|
||||
};
|
||||
|
||||
type PendingConnection = {
|
||||
promise: Promise<MCPConnection>;
|
||||
oauth: PendingOAuthState;
|
||||
oauth: OAuthLifecycleRelay;
|
||||
};
|
||||
|
||||
type ConnectionCreationGuard = { cancelled: boolean };
|
||||
|
|
@ -63,6 +49,10 @@ export abstract class UserConnectionManager {
|
|||
protected userLastActivity: Map<string, number> = new Map();
|
||||
/** In-flight connection promises keyed by `userId:serverName` — coalesces concurrent attempts */
|
||||
protected pendingConnections: Map<string, PendingConnection> = new Map();
|
||||
private readonly connectionBorrowers = new WeakMap<MCPConnection, number>();
|
||||
private readonly connectionBorrowerDrainWaiters = new WeakMap<MCPConnection, Set<() => void>>();
|
||||
private readonly deferredConnectionDisposalHolds = new WeakMap<MCPConnection, number>();
|
||||
private readonly deferredConnectionDisposals = new WeakMap<MCPConnection, string>();
|
||||
/** All durable creations, including forced replacements, visible to mutation teardown. */
|
||||
private readonly activeConnectionCreations: Map<string, Set<ConnectionCreationGuard>> = new Map();
|
||||
/** Serializes explicit durable replacements without coalescing their callers. */
|
||||
|
|
@ -233,24 +223,35 @@ export abstract class UserConnectionManager {
|
|||
? opts.requestScopedConnections
|
||||
: undefined;
|
||||
if (requestScopedConnections) {
|
||||
this.bindRequestScopedConnectionStore(requestScopedConnections);
|
||||
const requestConnectionKey = `${userId}:${serverName}`;
|
||||
const existing = requestScopedConnections.connections.get(requestConnectionKey) as
|
||||
| MCPConnection
|
||||
| undefined;
|
||||
if (existing) {
|
||||
if (!config || (config.updatedAt && existing.isStale(config.updatedAt))) {
|
||||
await existing.disconnect().catch((error) => {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to disconnect stale request-scoped connection`,
|
||||
error,
|
||||
);
|
||||
});
|
||||
requestScopedConnections.connections.delete(requestConnectionKey);
|
||||
} else if (await existing.isConnected()) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing request-scoped connection`);
|
||||
return existing;
|
||||
await this.disposeEvictedConnection(existing, `[MCP][User: ${userId}][${serverName}]`);
|
||||
} else {
|
||||
const activeRecovery = this.getActiveConnectionRecovery(existing);
|
||||
let awaitedRecovery = activeRecovery;
|
||||
if (activeRecovery) {
|
||||
await this.waitForConnectionRecovery(activeRecovery, opts.signal);
|
||||
}
|
||||
let connected = await existing.isConnected();
|
||||
let recovery = this.getActiveConnectionRecovery(existing);
|
||||
while (recovery && recovery !== awaitedRecovery) {
|
||||
awaitedRecovery = recovery;
|
||||
await this.waitForConnectionRecovery(recovery, opts.signal);
|
||||
connected = await existing.isConnected();
|
||||
recovery = this.getActiveConnectionRecovery(existing);
|
||||
}
|
||||
if (connected) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing request-scoped connection`);
|
||||
return existing;
|
||||
}
|
||||
requestScopedConnections.connections.delete(requestConnectionKey);
|
||||
await this.disposeEvictedConnection(existing, `[MCP][User: ${userId}][${serverName}]`);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -264,14 +265,19 @@ export abstract class UserConnectionManager {
|
|||
return pending;
|
||||
}
|
||||
|
||||
const pendingOAuth = this.createPendingOAuthState(opts.oauthStart);
|
||||
const pendingOAuth = new OAuthLifecycleRelay({
|
||||
oauthStart: opts.oauthStart,
|
||||
oauthEnd: opts.oauthEnd,
|
||||
logPrefix: `[MCP][User: ${userId}][${serverName}]`,
|
||||
});
|
||||
const connectionPromise = this.createUserConnectionInternal(
|
||||
{
|
||||
...opts,
|
||||
forceNew: true,
|
||||
ephemeralConnection: true,
|
||||
serverConfig: config,
|
||||
oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth),
|
||||
oauthStart: pendingOAuth.start,
|
||||
oauthEnd: pendingOAuth.end,
|
||||
},
|
||||
userId,
|
||||
forceNew === true,
|
||||
|
|
@ -303,12 +309,22 @@ export abstract class UserConnectionManager {
|
|||
const pending = this.pendingConnections.get(lockKey);
|
||||
if (pending) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Joining in-flight connection attempt`);
|
||||
await this.addPendingOAuthStart(pending.oauth, opts, userId);
|
||||
await pending.oauth.add({
|
||||
oauthStart: opts.oauthStart,
|
||||
oauthEnd: opts.oauthEnd,
|
||||
flowManager: opts.flowManager,
|
||||
userId,
|
||||
serverName,
|
||||
});
|
||||
return pending.promise;
|
||||
}
|
||||
}
|
||||
|
||||
const pendingOAuth = this.createPendingOAuthState(opts.oauthStart);
|
||||
const pendingOAuth = new OAuthLifecycleRelay({
|
||||
oauthStart: opts.oauthStart,
|
||||
oauthEnd: opts.oauthEnd,
|
||||
logPrefix: `[MCP][User: ${userId}][${serverName}]`,
|
||||
});
|
||||
const creationGuard: ConnectionCreationGuard = { cancelled: false };
|
||||
this.registerConnectionCreation(lockKey, creationGuard);
|
||||
const createConnection = () =>
|
||||
|
|
@ -318,7 +334,8 @@ export abstract class UserConnectionManager {
|
|||
forceNew: forceNewConnection,
|
||||
ephemeralConnection,
|
||||
serverConfig: config,
|
||||
oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth),
|
||||
oauthStart: pendingOAuth.start,
|
||||
oauthEnd: pendingOAuth.end,
|
||||
},
|
||||
userId,
|
||||
clearCooldown,
|
||||
|
|
@ -348,170 +365,6 @@ export abstract class UserConnectionManager {
|
|||
}
|
||||
}
|
||||
|
||||
private createPendingOAuthState(oauthStart?: t.OAuthStartHandler): PendingOAuthState {
|
||||
return {
|
||||
oauthStarts: oauthStart ? new Set([oauthStart]) : new Set(),
|
||||
emittedAuthUrls: new WeakMap<t.OAuthStartHandler, string>(),
|
||||
primaryOAuthStart: oauthStart,
|
||||
};
|
||||
}
|
||||
|
||||
private createPendingOAuthStart(
|
||||
serverName: string,
|
||||
userId: string,
|
||||
pendingOAuth: PendingOAuthState,
|
||||
): t.OAuthStartHandler {
|
||||
return async (authURL, options) => {
|
||||
pendingOAuth.lastOAuthStart = { authURL, options };
|
||||
|
||||
let primaryError: unknown;
|
||||
const oauthStarts = Array.from(pendingOAuth.oauthStarts);
|
||||
for (const oauthStart of oauthStarts) {
|
||||
try {
|
||||
await this.emitPendingOAuthStart(pendingOAuth, oauthStart, authURL, options);
|
||||
} catch (error) {
|
||||
if (oauthStart === pendingOAuth.primaryOAuthStart) {
|
||||
primaryError = error;
|
||||
} else {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to notify joined OAuth listener`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (primaryError) {
|
||||
throw primaryError;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private async addPendingOAuthStart(
|
||||
pendingOAuth: PendingOAuthState,
|
||||
opts: t.UserMCPConnectionOptions,
|
||||
userId: string,
|
||||
): Promise<void> {
|
||||
const { oauthStart, serverName } = opts;
|
||||
if (typeof oauthStart !== 'function') {
|
||||
return;
|
||||
}
|
||||
|
||||
pendingOAuth.oauthStarts.add(oauthStart);
|
||||
const lastOAuthStart = pendingOAuth.lastOAuthStart;
|
||||
if (lastOAuthStart) {
|
||||
try {
|
||||
const pendingOAuthStart =
|
||||
lastOAuthStart.options?.expiresAt == null
|
||||
? await this.getFlowPendingOAuthStart(opts, userId)
|
||||
: undefined;
|
||||
const replayOAuthStart =
|
||||
pendingOAuthStart?.authURL === lastOAuthStart.authURL
|
||||
? pendingOAuthStart
|
||||
: lastOAuthStart;
|
||||
await this.emitPendingOAuthStart(
|
||||
pendingOAuth,
|
||||
oauthStart,
|
||||
replayOAuthStart.authURL,
|
||||
replayOAuthStart.options,
|
||||
);
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to re-issue pending OAuth URL`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
await this.reissuePendingOAuthStart(opts, userId, pendingOAuth);
|
||||
}
|
||||
|
||||
private async emitPendingOAuthStart(
|
||||
pendingOAuth: PendingOAuthState,
|
||||
oauthStart: t.OAuthStartHandler,
|
||||
authURL: string,
|
||||
options?: t.OAuthStartOptions,
|
||||
): Promise<void> {
|
||||
if (pendingOAuth.emittedAuthUrls.get(oauthStart) === authURL) {
|
||||
return;
|
||||
}
|
||||
pendingOAuth.emittedAuthUrls.set(oauthStart, authURL);
|
||||
await oauthStart(authURL, options);
|
||||
}
|
||||
|
||||
private getPendingOAuthStart(flow: FlowState | null | undefined): PendingOAuthStart | undefined {
|
||||
if (flow?.status !== 'PENDING') {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const expiresAt = flow.createdAt + PENDING_STALE_MS;
|
||||
if (expiresAt <= Date.now()) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const metadata = flow.metadata as MCPOAuthFlowMetadata | undefined;
|
||||
const authorizationUrl = metadata?.authorizationUrl;
|
||||
if (!authorizationUrl) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return { authURL: authorizationUrl, options: { expiresAt } };
|
||||
}
|
||||
|
||||
private async getFlowPendingOAuthStart(
|
||||
{ flowManager, serverName }: Pick<t.UserMCPConnectionOptions, 'flowManager' | 'serverName'>,
|
||||
userId: string,
|
||||
): Promise<PendingOAuthStart | undefined> {
|
||||
if (!flowManager) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const flowId = MCPOAuthHandler.generateFlowId(userId, serverName, getTenantId());
|
||||
const existingFlow = await flowManager.getFlowState(flowId, 'mcp_oauth');
|
||||
return this.getPendingOAuthStart(existingFlow);
|
||||
}
|
||||
|
||||
private async reissuePendingOAuthStart(
|
||||
{ flowManager, oauthStart, serverName }: t.UserMCPConnectionOptions,
|
||||
userId: string,
|
||||
pendingOAuth?: PendingOAuthState,
|
||||
): Promise<void> {
|
||||
if (!flowManager || typeof oauthStart !== 'function') {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const pendingOAuthStart = await this.getFlowPendingOAuthStart(
|
||||
{ flowManager, serverName },
|
||||
userId,
|
||||
);
|
||||
if (!pendingOAuthStart) {
|
||||
return;
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[MCP][User: ${userId}][${serverName}] Re-issuing stored authorization URL while joining in-flight connection`,
|
||||
);
|
||||
if (pendingOAuth) {
|
||||
pendingOAuth.lastOAuthStart = pendingOAuthStart;
|
||||
await this.emitPendingOAuthStart(
|
||||
pendingOAuth,
|
||||
oauthStart,
|
||||
pendingOAuthStart.authURL,
|
||||
pendingOAuthStart.options,
|
||||
);
|
||||
} else {
|
||||
await oauthStart(pendingOAuthStart.authURL, pendingOAuthStart.options);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to re-issue pending OAuth URL`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private async createUserConnectionInternal(
|
||||
{
|
||||
serverName,
|
||||
|
|
@ -621,22 +474,41 @@ export abstract class UserConnectionManager {
|
|||
if (!config || (config.updatedAt && connection.isStale(config.updatedAt))) {
|
||||
if (config) {
|
||||
logger.info(
|
||||
`[MCP][User: ${userId}][${serverName}] Config was updated, disconnecting stale connection`,
|
||||
`[MCP][User: ${userId}][${serverName}] Config was updated, evicting stale connection`,
|
||||
);
|
||||
}
|
||||
await this.disconnectUserConnection(userId, serverName, creationGuard);
|
||||
connection = undefined;
|
||||
} else if (await connection.isConnected()) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing active connection`);
|
||||
await this.updateUserLastActivity(userId);
|
||||
await this.assertToolPublicationLeaseCurrent(connection, userId, serverName, creationGuard);
|
||||
if (creationGuard?.cancelled) {
|
||||
throw new Error(
|
||||
`[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`,
|
||||
);
|
||||
}
|
||||
return connection;
|
||||
} else {
|
||||
const activeRecovery = this.getActiveConnectionRecovery(connection);
|
||||
let awaitedRecovery = activeRecovery;
|
||||
if (activeRecovery) {
|
||||
await this.waitForConnectionRecovery(activeRecovery, signal);
|
||||
}
|
||||
let connected = await connection.isConnected();
|
||||
let recovery = this.getActiveConnectionRecovery(connection);
|
||||
while (recovery && recovery !== awaitedRecovery) {
|
||||
awaitedRecovery = recovery;
|
||||
await this.waitForConnectionRecovery(recovery, signal);
|
||||
connected = await connection.isConnected();
|
||||
recovery = this.getActiveConnectionRecovery(connection);
|
||||
}
|
||||
if (connected) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing active connection`);
|
||||
await this.updateUserLastActivity(userId);
|
||||
await this.assertToolPublicationLeaseCurrent(
|
||||
connection,
|
||||
userId,
|
||||
serverName,
|
||||
creationGuard,
|
||||
);
|
||||
if (creationGuard?.cancelled) {
|
||||
throw new Error(
|
||||
`[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`,
|
||||
);
|
||||
}
|
||||
return connection;
|
||||
}
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Found existing but disconnected connection object. Cleaning up.`,
|
||||
);
|
||||
|
|
@ -964,6 +836,115 @@ export abstract class UserConnectionManager {
|
|||
logger.debug(`[MCP][User: ${userId}][${serverName}] Removed connection entry.`);
|
||||
}
|
||||
|
||||
protected retainConnection(connection: MCPConnection): void {
|
||||
const borrowers = this.connectionBorrowers.get(connection) ?? 0;
|
||||
this.connectionBorrowers.set(connection, borrowers + 1);
|
||||
}
|
||||
|
||||
protected getActiveConnectionRecovery(_connection: MCPConnection): Promise<void> | undefined {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
protected waitForConnectionRecovery(
|
||||
recovery: Promise<void>,
|
||||
_signal?: AbortSignal,
|
||||
): Promise<void> {
|
||||
return recovery;
|
||||
}
|
||||
|
||||
protected holdDeferredConnectionDisposal(connection: MCPConnection): void {
|
||||
const holds = this.deferredConnectionDisposalHolds.get(connection) ?? 0;
|
||||
this.deferredConnectionDisposalHolds.set(connection, holds + 1);
|
||||
}
|
||||
|
||||
protected async releaseDeferredConnectionDisposal(connection: MCPConnection): Promise<void> {
|
||||
const holds = this.deferredConnectionDisposalHolds.get(connection) ?? 0;
|
||||
if (holds > 1) {
|
||||
this.deferredConnectionDisposalHolds.set(connection, holds - 1);
|
||||
return;
|
||||
}
|
||||
this.deferredConnectionDisposalHolds.delete(connection);
|
||||
await this.finalizeDeferredConnectionDisposal(connection);
|
||||
}
|
||||
|
||||
protected async releaseConnection(connection: MCPConnection): Promise<void> {
|
||||
const borrowers = this.connectionBorrowers.get(connection) ?? 0;
|
||||
if (borrowers > 1) {
|
||||
this.connectionBorrowers.set(connection, borrowers - 1);
|
||||
return;
|
||||
}
|
||||
|
||||
this.connectionBorrowers.delete(connection);
|
||||
await this.finalizeDeferredConnectionDisposal(connection);
|
||||
|
||||
const drainWaiters = this.connectionBorrowerDrainWaiters.get(connection);
|
||||
if (drainWaiters) {
|
||||
this.connectionBorrowerDrainWaiters.delete(connection);
|
||||
for (const resolve of drainWaiters) {
|
||||
resolve();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
protected waitForConnectionBorrowersToDrain(connection: MCPConnection): Promise<void> {
|
||||
if ((this.connectionBorrowers.get(connection) ?? 0) === 0) {
|
||||
return Promise.resolve();
|
||||
}
|
||||
|
||||
return new Promise<void>((resolve) => {
|
||||
const drainWaiters = this.connectionBorrowerDrainWaiters.get(connection) ?? new Set();
|
||||
drainWaiters.add(resolve);
|
||||
this.connectionBorrowerDrainWaiters.set(connection, drainWaiters);
|
||||
});
|
||||
}
|
||||
|
||||
protected bindRequestScopedConnectionStore(
|
||||
requestScopedConnections?: t.RequestScopedMCPConnectionStore,
|
||||
): void {
|
||||
if (!requestScopedConnections || requestScopedConnections.disposeConnection) {
|
||||
return;
|
||||
}
|
||||
|
||||
requestScopedConnections.disposeConnection = async (connectionKey, connection) => {
|
||||
await this.disposeEvictedConnection(
|
||||
connection as MCPConnection,
|
||||
`[MCP][Request-scoped: ${connectionKey}]`,
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
protected async disposeEvictedConnection(
|
||||
connection: MCPConnection,
|
||||
logPrefix: string,
|
||||
): Promise<void> {
|
||||
this.deferredConnectionDisposals.set(connection, logPrefix);
|
||||
await this.finalizeDeferredConnectionDisposal(connection);
|
||||
}
|
||||
|
||||
private async finalizeDeferredConnectionDisposal(connection: MCPConnection): Promise<void> {
|
||||
if (
|
||||
(this.connectionBorrowers.get(connection) ?? 0) > 0 ||
|
||||
(this.deferredConnectionDisposalHolds.get(connection) ?? 0) > 0
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
const logPrefix = this.deferredConnectionDisposals.get(connection);
|
||||
if (!logPrefix) {
|
||||
return;
|
||||
}
|
||||
this.deferredConnectionDisposals.delete(connection);
|
||||
await this.disposeConnection(connection, logPrefix);
|
||||
}
|
||||
|
||||
private async disposeConnection(connection: MCPConnection, logPrefix: string): Promise<void> {
|
||||
try {
|
||||
await connection.dispose();
|
||||
} catch (error) {
|
||||
logger.warn(`${logPrefix} Failed to dispose evicted connection`, error);
|
||||
}
|
||||
}
|
||||
|
||||
/** Disconnects and removes a specific user connection */
|
||||
public async disconnectUserConnection(
|
||||
userId: string,
|
||||
|
|
@ -980,10 +961,11 @@ export abstract class UserConnectionManager {
|
|||
const connection = userMap?.get(serverName);
|
||||
try {
|
||||
if (connection) {
|
||||
logger.info(`[MCP][User: ${userId}][${serverName}] Disconnecting...`);
|
||||
const logPrefix = `[MCP][User: ${userId}][${serverName}]`;
|
||||
logger.info(`${logPrefix} Disconnecting...`);
|
||||
connection.removeAllListeners?.('toolsChanged');
|
||||
this.removeUserConnection(userId, serverName);
|
||||
await connection.dispose();
|
||||
await this.disposeEvictedConnection(connection, logPrefix);
|
||||
}
|
||||
} finally {
|
||||
await cancelMCPToolsChanged({ userId, serverName });
|
||||
|
|
|
|||
|
|
@ -1,16 +1,17 @@
|
|||
/**
|
||||
* Tests for MCPConnection error detection methods.
|
||||
*
|
||||
* These tests use standalone implementations that mirror the private methods in MCPConnection.
|
||||
* This approach was chosen because MCPConnection requires complex dependencies (Client, transport)
|
||||
* that are difficult to mock properly. The standalone implementations are kept in sync with
|
||||
* the actual implementation in connection.ts.
|
||||
* Rate-limit and SSE tests use standalone implementations that mirror private methods in
|
||||
* MCPConnection. OAuth classification exercises the production helper shared by the connection
|
||||
* and factory.
|
||||
*
|
||||
* Alternative approaches considered:
|
||||
* 1. Reflection/type casting - fragile and breaks with refactoring
|
||||
* 2. Protected methods with test subclass - changes public API for testing
|
||||
* 3. Integration tests - tested separately in the full MCP test suite
|
||||
*/
|
||||
import { isOAuthAuthenticationError } from '~/mcp/errors';
|
||||
|
||||
describe('MCPConnection Error Detection', () => {
|
||||
/**
|
||||
* Standalone implementation of isRateLimitError for testing.
|
||||
|
|
@ -45,52 +46,6 @@ describe('MCPConnection Error Detection', () => {
|
|||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Standalone implementation of isOAuthError for testing.
|
||||
* This mirrors the private method in MCPConnection (connection.ts).
|
||||
* Keep in sync with the actual implementation.
|
||||
*/
|
||||
function isOAuthError(error: unknown): boolean {
|
||||
if (!error || typeof error !== 'object') {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check for error code
|
||||
if ('code' in error) {
|
||||
const code = (error as { code?: number }).code;
|
||||
if (code === 401 || code === 403) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// Check message for various auth error indicators
|
||||
if ('message' in error && typeof error.message === 'string') {
|
||||
const message = error.message.toLowerCase();
|
||||
// Check for 401 status
|
||||
if (message.includes('401') || message.includes('non-200 status code (401)')) {
|
||||
return true;
|
||||
}
|
||||
// Check for invalid_grant (OAuth servers return this for expired/revoked grants)
|
||||
if (message.includes('invalid_grant')) {
|
||||
return true;
|
||||
}
|
||||
// Check for invalid_token (OAuth servers return this for expired/revoked tokens)
|
||||
if (message.includes('invalid_token')) {
|
||||
return true;
|
||||
}
|
||||
// Check for authentication required
|
||||
if (message.includes('authentication required') || message.includes('unauthorized')) {
|
||||
return true;
|
||||
}
|
||||
// Check for missing authorization values (e.g., Amazon Ads MCP returns HTTP 400 with this)
|
||||
if (message.includes('no authorization')) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
describe('isRateLimitError', () => {
|
||||
it('should detect rate limit error by code 429', () => {
|
||||
const error = { code: 429, message: 'Too many requests' };
|
||||
|
|
@ -142,30 +97,36 @@ describe('MCPConnection Error Detection', () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe('isOAuthError', () => {
|
||||
it('should detect OAuth error by code 401', () => {
|
||||
const error = { code: 401, message: 'Unauthorized' };
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
describe('isOAuthAuthenticationError', () => {
|
||||
it.each([
|
||||
{ code: 401, message: 'Unauthorized' },
|
||||
{ status: 403, message: 'Forbidden' },
|
||||
{ statusCode: 401, message: 'Authentication required' },
|
||||
{ message: 'Error POSTing to endpoint (HTTP 401): Unauthorized' },
|
||||
{ message: 'Error POSTing to endpoint (HTTP 403): Forbidden' },
|
||||
{ message: 'Non-200 status code (403)' },
|
||||
{ message: '403 Forbidden' },
|
||||
{ message: 'Unauthorized (401)' },
|
||||
{ message: 'Forbidden (403)' },
|
||||
{ message: 'The server rejected the token with insufficient_scope' },
|
||||
])('should detect OAuth authentication error %#', (error) => {
|
||||
expect(isOAuthAuthenticationError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect OAuth error by code 403', () => {
|
||||
const error = { code: 403, message: 'Forbidden' };
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect OAuth error by message containing 401', () => {
|
||||
const error = { message: 'Error POSTing to endpoint (HTTP 401): Unauthorized' };
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should not detect OAuth error for 429 rate limit', () => {
|
||||
const error = { code: 429, message: 'Too many requests' };
|
||||
expect(isOAuthError(error)).toBe(false);
|
||||
it.each([
|
||||
{ code: 429, message: 'Too many requests' },
|
||||
{ message: 'Customer 401 not found' },
|
||||
{ message: 'Order 403 is unavailable' },
|
||||
{ message: 'User is unauthorized to delete this record' },
|
||||
{ message: 'No authorization to delete this record' },
|
||||
{ code: 400, message: 'Bad request: missing required field' },
|
||||
])('should ignore non-authentication error %#', (error) => {
|
||||
expect(isOAuthAuthenticationError(error)).toBe(false);
|
||||
});
|
||||
|
||||
it('should detect OAuth error for invalid_token', () => {
|
||||
const error = { message: 'The access token is invalid_token or expired' };
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
expect(isOAuthAuthenticationError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect OAuth error for invalid_grant', () => {
|
||||
|
|
@ -173,7 +134,7 @@ describe('MCPConnection Error Detection', () => {
|
|||
message:
|
||||
'Streamable HTTP error: Error POSTing to endpoint: {"error":"invalid_grant","error_description":"The provided authorization grant is invalid, expired, or revoked"}',
|
||||
};
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
expect(isOAuthAuthenticationError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect OAuth error for "no authorization" in message (HTTP 400)', () => {
|
||||
|
|
@ -181,17 +142,12 @@ describe('MCPConnection Error Detection', () => {
|
|||
message:
|
||||
'Either no authorization values are specified or it could not be derived from the request',
|
||||
};
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
expect(isOAuthAuthenticationError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should detect OAuth error for "No authorization" with different casing', () => {
|
||||
const error = { message: 'No Authorization header provided' };
|
||||
expect(isOAuthError(error)).toBe(true);
|
||||
});
|
||||
|
||||
it('should not detect OAuth error for unrelated 400 errors', () => {
|
||||
const error = { code: 400, message: 'Bad request: missing required field' };
|
||||
expect(isOAuthError(error)).toBe(false);
|
||||
expect(isOAuthAuthenticationError(error)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -202,10 +158,10 @@ describe('MCPConnection Error Detection', () => {
|
|||
|
||||
// Rate limit error should be detected as rate limit, not OAuth
|
||||
expect(isRateLimitError(rateLimitError)).toBe(true);
|
||||
expect(isOAuthError(rateLimitError)).toBe(false);
|
||||
expect(isOAuthAuthenticationError(rateLimitError)).toBe(false);
|
||||
|
||||
// OAuth error should be detected as OAuth, not rate limit
|
||||
expect(isOAuthError(oauthError)).toBe(true);
|
||||
expect(isOAuthAuthenticationError(oauthError)).toBe(true);
|
||||
expect(isRateLimitError(oauthError)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -549,12 +549,13 @@ describe('MCPConnection SSE 404 handling – session-aware', () => {
|
|||
conn: MCPConnection,
|
||||
transport: ReturnType<typeof makeTransportStub>,
|
||||
code = 404,
|
||||
) {
|
||||
): Error {
|
||||
(
|
||||
conn as unknown as { setupTransportErrorHandlers: (t: unknown) => void }
|
||||
).setupTransportErrorHandlers(transport);
|
||||
const sseError = Object.assign(new Error('Failed to open SSE stream'), { code });
|
||||
transport.onerror?.(sseError);
|
||||
return sseError;
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
|
|
@ -615,6 +616,19 @@ describe('MCPConnection SSE 404 handling – session-aware', () => {
|
|||
expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('session lost'));
|
||||
expect(emitSpy).toHaveBeenCalledWith('connectionChange', 'error');
|
||||
});
|
||||
it('marks an OAuth-challenged connection unusable without blindly reconnecting', async () => {
|
||||
const conn = makeConn();
|
||||
const transport = makeTransportStub();
|
||||
conn.emit('connectionChange', 'connected');
|
||||
const emitSpy = jest.spyOn(conn, 'emit');
|
||||
|
||||
const oauthError = fireSSEError(conn, transport, 401);
|
||||
|
||||
expect(emitSpy).toHaveBeenCalledWith('oauthError', expect.any(Error));
|
||||
expect(emitSpy).not.toHaveBeenCalledWith('connectionChange', 'error');
|
||||
expect(await conn.isConnected()).toBe(false);
|
||||
expect(conn.getLastConnectionCheckError()).toBe(oauthError);
|
||||
});
|
||||
});
|
||||
|
||||
describe('MCPConnection SSE stream disconnect handling', () => {
|
||||
|
|
|
|||
|
|
@ -11,10 +11,12 @@ import {
|
|||
createOAuthMCPServer,
|
||||
type OAuthTestServer,
|
||||
} from './helpers/oauthTestServer';
|
||||
import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry';
|
||||
import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory';
|
||||
import { MCPTokenStorage, MCPOAuthHandler } from '~/mcp/oauth';
|
||||
import { FlowStateManager } from '~/flow/manager';
|
||||
import { MCPConnection } from '~/mcp/connection';
|
||||
import { MCPTokenStorage } from '~/mcp/oauth';
|
||||
import { MCPManager } from '~/mcp/MCPManager';
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
|
|
@ -64,6 +66,16 @@ async function safeDisconnect(conn: MCPConnection | null): Promise<void> {
|
|||
await conn.disconnect().catch(() => undefined);
|
||||
}
|
||||
|
||||
async function waitFor(condition: () => boolean, timeoutMs = 5000): Promise<void> {
|
||||
const deadline = Date.now() + timeoutMs;
|
||||
while (!condition()) {
|
||||
if (Date.now() > deadline) {
|
||||
throw new Error('Timed out waiting for condition');
|
||||
}
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
}
|
||||
}
|
||||
|
||||
function createFlowManager(): FlowStateManager<MCPOAuthTokens | null> {
|
||||
return new FlowStateManager(new MockKeyv<MCPOAuthTokens | null>() as unknown as Keyv, {
|
||||
ttl: 30000,
|
||||
|
|
@ -122,7 +134,7 @@ async function storeTokens(
|
|||
server: OAuthTestServer,
|
||||
tokens: MCPOAuthTokens,
|
||||
scope = 'read',
|
||||
): Promise<void> {
|
||||
): Promise<MCPOAuthTokens> {
|
||||
const clientInfo: OAuthClientInformation = {
|
||||
client_id: CLIENT_ID,
|
||||
redirect_uris: ['http://localhost'],
|
||||
|
|
@ -144,7 +156,7 @@ async function storeTokens(
|
|||
resource: server.resourceUrl,
|
||||
};
|
||||
|
||||
await MCPTokenStorage.storeTokens({
|
||||
return MCPTokenStorage.storeTokens({
|
||||
userId: USER_ID,
|
||||
serverName: SERVER_NAME,
|
||||
tokens,
|
||||
|
|
@ -231,6 +243,258 @@ describe('MCPConnectionFactory OAuth against real SDK Streamable HTTP server', (
|
|||
expect(storedAccessToken?.token).not.toBe(`enc:${initialTokens.access_token}`);
|
||||
});
|
||||
|
||||
it('recovers a tool call rejected after connection and retries with the refreshed token', async () => {
|
||||
server = await createOAuthMCPServer({
|
||||
issueRefreshTokens: true,
|
||||
requireResourceParameter: true,
|
||||
tokenScopes: ['read'],
|
||||
scopesSupported: ['read'],
|
||||
});
|
||||
const initialTokens = await issueTokens(server);
|
||||
await storeTokens(tokenStore, server, initialTokens);
|
||||
const flowManager = createFlowManager();
|
||||
const tokenMethods = {
|
||||
findToken: tokenStore.findToken,
|
||||
createToken: tokenStore.createToken,
|
||||
updateToken: tokenStore.updateToken,
|
||||
deleteTokens: tokenStore.deleteTokens,
|
||||
};
|
||||
const serverConfig = {
|
||||
type: 'streamable-http' as const,
|
||||
url: server.url,
|
||||
initTimeout: 15000,
|
||||
requiresOAuth: true,
|
||||
};
|
||||
|
||||
connection = await MCPConnectionFactory.create(
|
||||
{ serverName: SERVER_NAME, serverConfig },
|
||||
{
|
||||
useOAuth: true,
|
||||
user: { id: USER_ID } as IUser,
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
},
|
||||
);
|
||||
server.issuedTokens.delete(initialTokens.access_token);
|
||||
const isConnectedSpy = jest.spyOn(connection, 'isConnected').mockResolvedValueOnce(true);
|
||||
|
||||
const manager = new MCPManager();
|
||||
jest.spyOn(manager, 'getConnection').mockResolvedValue(connection);
|
||||
const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({
|
||||
resolveAllowlists: jest.fn().mockResolvedValue({
|
||||
allowedDomains: null,
|
||||
allowedAddresses: null,
|
||||
useSSRFProtection: false,
|
||||
}),
|
||||
} as unknown as MCPServersRegistry);
|
||||
const oauthStart = jest.fn(async (_authorizationUrl: string): Promise<void> => undefined);
|
||||
|
||||
try {
|
||||
await expect(
|
||||
manager.callTool({
|
||||
user: { id: USER_ID } as IUser,
|
||||
serverName: SERVER_NAME,
|
||||
serverConfig,
|
||||
toolName: 'echo',
|
||||
toolArguments: { message: 'runtime refresh' },
|
||||
provider: 'openai',
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
oauthStart,
|
||||
}),
|
||||
).resolves.toBeDefined();
|
||||
|
||||
expect(
|
||||
server.tokenRequests.filter((request) => request.grantType === 'refresh_token'),
|
||||
).toHaveLength(1);
|
||||
expect(oauthStart).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
isConnectedSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('lets an in-flight request finish before a concurrent OAuth reconnect', async () => {
|
||||
let markSlowRequestStarted: (() => void) | undefined;
|
||||
let releaseFirstSlowRequest: (() => void) | undefined;
|
||||
const slowRequestStarted = new Promise<void>((resolve) => {
|
||||
markSlowRequestStarted = resolve;
|
||||
});
|
||||
const firstSlowRequestBlocked = new Promise<void>((resolve) => {
|
||||
releaseFirstSlowRequest = resolve;
|
||||
});
|
||||
let slowRequestCount = 0;
|
||||
server = await createOAuthMCPServer({
|
||||
issueRefreshTokens: true,
|
||||
requireResourceParameter: true,
|
||||
tokenScopes: ['read'],
|
||||
scopesSupported: ['read'],
|
||||
echoHandler: async (message) => {
|
||||
if (message === 'slow borrower' && slowRequestCount++ === 0) {
|
||||
markSlowRequestStarted?.();
|
||||
await firstSlowRequestBlocked;
|
||||
}
|
||||
return `echo: ${message}`;
|
||||
},
|
||||
});
|
||||
const initialTokens = await issueTokens(server);
|
||||
await storeTokens(tokenStore, server, initialTokens);
|
||||
const flowManager = createFlowManager();
|
||||
const tokenMethods = {
|
||||
findToken: tokenStore.findToken,
|
||||
createToken: tokenStore.createToken,
|
||||
updateToken: tokenStore.updateToken,
|
||||
deleteTokens: tokenStore.deleteTokens,
|
||||
};
|
||||
const serverConfig = {
|
||||
type: 'streamable-http' as const,
|
||||
url: server.url,
|
||||
initTimeout: 15000,
|
||||
requiresOAuth: true,
|
||||
};
|
||||
|
||||
connection = await MCPConnectionFactory.create(
|
||||
{ serverName: SERVER_NAME, serverConfig },
|
||||
{
|
||||
useOAuth: true,
|
||||
user: { id: USER_ID } as IUser,
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
},
|
||||
);
|
||||
const connectSpy = jest.spyOn(connection, 'connect');
|
||||
const isConnectedSpy = jest.spyOn(connection, 'isConnected').mockResolvedValue(true);
|
||||
|
||||
const manager = new MCPManager();
|
||||
jest.spyOn(manager, 'getConnection').mockResolvedValue(connection);
|
||||
const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({
|
||||
resolveAllowlists: jest.fn().mockResolvedValue({
|
||||
allowedDomains: null,
|
||||
allowedAddresses: null,
|
||||
useSSRFProtection: false,
|
||||
}),
|
||||
} as unknown as MCPServersRegistry);
|
||||
const oauthStart = jest.fn(async (_authorizationUrl: string): Promise<void> => undefined);
|
||||
const callTool = (message: string) =>
|
||||
manager.callTool({
|
||||
user: { id: USER_ID } as IUser,
|
||||
serverName: SERVER_NAME,
|
||||
serverConfig,
|
||||
toolName: 'echo',
|
||||
toolArguments: { message },
|
||||
provider: 'openai',
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
oauthStart,
|
||||
});
|
||||
|
||||
try {
|
||||
const slowCall = callTool('slow borrower');
|
||||
await slowRequestStarted;
|
||||
server.issuedTokens.delete(initialTokens.access_token);
|
||||
|
||||
const recoveringCall = callTool('recovery owner');
|
||||
await waitFor(
|
||||
() =>
|
||||
server.tokenRequests.filter((request) => request.grantType === 'refresh_token').length ===
|
||||
1,
|
||||
);
|
||||
|
||||
expect(connectSpy).not.toHaveBeenCalled();
|
||||
expect(slowRequestCount).toBe(1);
|
||||
releaseFirstSlowRequest?.();
|
||||
|
||||
await expect(Promise.all([slowCall, recoveringCall])).resolves.toHaveLength(2);
|
||||
expect(
|
||||
server.tokenRequests.filter((request) => request.grantType === 'refresh_token'),
|
||||
).toHaveLength(1);
|
||||
expect(slowRequestCount).toBe(1);
|
||||
expect(oauthStart).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
releaseFirstSlowRequest?.();
|
||||
connectSpy.mockRestore();
|
||||
isConnectedSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('escalates to interactive OAuth when the resource rejects refreshed tokens during reconnect', async () => {
|
||||
server = await createOAuthMCPServer({
|
||||
issueRefreshTokens: true,
|
||||
requireResourceParameter: true,
|
||||
tokenScopes: ['read'],
|
||||
scopesSupported: ['read'],
|
||||
rejectRefreshTokens: 10,
|
||||
});
|
||||
const initialTokens = await issueTokens(server);
|
||||
await storeTokens(tokenStore, server, initialTokens);
|
||||
const flowManager = createFlowManager();
|
||||
const tokenMethods = {
|
||||
findToken: tokenStore.findToken,
|
||||
createToken: tokenStore.createToken,
|
||||
updateToken: tokenStore.updateToken,
|
||||
deleteTokens: tokenStore.deleteTokens,
|
||||
};
|
||||
const serverConfig = {
|
||||
type: 'streamable-http' as const,
|
||||
url: server.url,
|
||||
initTimeout: 15000,
|
||||
requiresOAuth: true,
|
||||
};
|
||||
|
||||
connection = await MCPConnectionFactory.create(
|
||||
{ serverName: SERVER_NAME, serverConfig },
|
||||
{
|
||||
useOAuth: true,
|
||||
user: { id: USER_ID } as IUser,
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
},
|
||||
);
|
||||
server.issuedTokens.delete(initialTokens.access_token);
|
||||
const isConnectedSpy = jest.spyOn(connection, 'isConnected').mockResolvedValueOnce(true);
|
||||
|
||||
const manager = new MCPManager();
|
||||
jest.spyOn(manager, 'getConnection').mockResolvedValue(connection);
|
||||
const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({
|
||||
resolveAllowlists: jest.fn().mockResolvedValue({
|
||||
allowedDomains: null,
|
||||
allowedAddresses: null,
|
||||
useSSRFProtection: false,
|
||||
}),
|
||||
} as unknown as MCPServersRegistry);
|
||||
const oauthStart = jest.fn(async (): Promise<void> => {
|
||||
const authorizedTokens = await issueTokens(server);
|
||||
const storedTokens = await storeTokens(tokenStore, server, authorizedTokens);
|
||||
const flowId = MCPOAuthHandler.generateFlowId(USER_ID, SERVER_NAME);
|
||||
await flowManager.completeFlow(flowId, 'mcp_oauth', storedTokens);
|
||||
});
|
||||
|
||||
try {
|
||||
await expect(
|
||||
manager.callTool({
|
||||
user: { id: USER_ID } as IUser,
|
||||
serverName: SERVER_NAME,
|
||||
serverConfig,
|
||||
toolName: 'echo',
|
||||
toolArguments: { message: 'interactive fallback' },
|
||||
provider: 'openai',
|
||||
flowManager,
|
||||
tokenMethods,
|
||||
oauthStart,
|
||||
}),
|
||||
).resolves.toBeDefined();
|
||||
|
||||
expect(
|
||||
server.tokenRequests.filter((request) => request.grantType === 'refresh_token'),
|
||||
).toHaveLength(1);
|
||||
expect(oauthStart).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
isConnectedSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('does not silently refresh an SDK insufficient_scope challenge before starting OAuth', async () => {
|
||||
server = await createOAuthMCPServer({
|
||||
issueRefreshTokens: true,
|
||||
|
|
|
|||
|
|
@ -95,6 +95,9 @@ describe('MCPConnectionFactory', () => {
|
|||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
// Cached runtime handlers now delegate before attempting their own refresh,
|
||||
// so queued one-shot refresh results must not leak into the next test.
|
||||
mockMCPTokenStorage.forceRefreshTokens.mockReset();
|
||||
// Clear process-local silent-refresh in-flight map so a leftover entry
|
||||
// from a prior test (e.g. one that errored before its `finally` ran)
|
||||
// cannot cause a later test to join a stale promise.
|
||||
|
|
@ -207,6 +210,95 @@ describe('MCPConnectionFactory', () => {
|
|||
}
|
||||
});
|
||||
|
||||
it('aborts only the local waiter for a shared OAuth flow', async () => {
|
||||
const abortController = new AbortController();
|
||||
const abortReason = new Error('owner request aborted');
|
||||
const sseConfig = {
|
||||
url: 'https://api.example.com/mcp',
|
||||
type: 'sse' as const,
|
||||
requiresOAuth: true,
|
||||
} as t.SSEOptions;
|
||||
const pendingTokens = new Promise<MCPOAuthTokens | null>(() => undefined);
|
||||
const oauthStart = jest.fn().mockResolvedValue(undefined);
|
||||
|
||||
mockProcessMCPEnv.mockReturnValue(sseConfig);
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
type: 'mcp_oauth',
|
||||
metadata: {
|
||||
authorizationUrl: 'https://auth.example.com/pending',
|
||||
serverUrl: sseConfig.url,
|
||||
},
|
||||
createdAt: Date.now(),
|
||||
});
|
||||
mockFlowManager.createFlow.mockReturnValue(pendingTokens);
|
||||
|
||||
const factory = new InspectableMCPConnectionFactory(
|
||||
{ serverName: 'test-server', serverConfig: sseConfig },
|
||||
{
|
||||
useOAuth: true,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
oauthStart,
|
||||
signal: abortController.signal,
|
||||
},
|
||||
);
|
||||
|
||||
const resultPromise = factory.handleOAuthRequiredForTest();
|
||||
await new Promise((resolve) => setImmediate(resolve));
|
||||
|
||||
expect(mockFlowManager.createFlow).toHaveBeenCalledWith('user123:test-server', 'mcp_oauth', {});
|
||||
|
||||
abortController.abort(abortReason);
|
||||
|
||||
await expect(resultPromise).resolves.toEqual(
|
||||
expect.objectContaining({ tokens: null, error: abortReason }),
|
||||
);
|
||||
expect(mockFlowManager.deleteFlow).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('observes the shared flow after an already-aborted local wait', async () => {
|
||||
const abortController = new AbortController();
|
||||
const abortReason = new Error('request already aborted');
|
||||
const sseConfig = {
|
||||
url: 'https://api.example.com/mcp',
|
||||
type: 'sse' as const,
|
||||
requiresOAuth: true,
|
||||
} as t.SSEOptions;
|
||||
let rejectSharedFlow: ((error: Error) => void) | undefined;
|
||||
const sharedFlow = new Promise<MCPOAuthTokens | null>((_resolve, reject) => {
|
||||
rejectSharedFlow = reject;
|
||||
});
|
||||
|
||||
abortController.abort(abortReason);
|
||||
mockProcessMCPEnv.mockReturnValue(sseConfig);
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
type: 'mcp_oauth',
|
||||
metadata: { authorizationUrl: 'https://auth.example.com/pending' },
|
||||
createdAt: Date.now(),
|
||||
});
|
||||
mockFlowManager.createFlow.mockReturnValue(sharedFlow);
|
||||
|
||||
const factory = new InspectableMCPConnectionFactory(
|
||||
{ serverName: 'test-server', serverConfig: sseConfig },
|
||||
{
|
||||
useOAuth: true,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
signal: abortController.signal,
|
||||
},
|
||||
);
|
||||
|
||||
await expect(factory.handleOAuthRequiredForTest()).resolves.toEqual(
|
||||
expect.objectContaining({ tokens: null, error: abortReason }),
|
||||
);
|
||||
|
||||
rejectSharedFlow?.(new Error('shared flow later failed'));
|
||||
await new Promise((resolve) => setImmediate(resolve));
|
||||
expect(mockFlowManager.deleteFlow).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
describe('static create method', () => {
|
||||
it('should create a basic connection without OAuth', async () => {
|
||||
const basicOptions = {
|
||||
|
|
@ -1250,12 +1342,7 @@ describe('MCPConnectionFactory', () => {
|
|||
expect(initCallOrder).toBeLessThan(createCallOrder);
|
||||
|
||||
// createFlow should receive {} since initFlow already persisted metadata
|
||||
expect(mockFlowManager.createFlow).toHaveBeenCalledWith(
|
||||
'flow123',
|
||||
'mcp_oauth',
|
||||
{},
|
||||
undefined,
|
||||
);
|
||||
expect(mockFlowManager.createFlow).toHaveBeenCalledWith('flow123', 'mcp_oauth', {});
|
||||
});
|
||||
|
||||
it('should delete stale flow and create new OAuth flow when existing flow is COMPLETED', async () => {
|
||||
|
|
@ -1342,7 +1429,6 @@ describe('MCPConnectionFactory', () => {
|
|||
'user123:test-server',
|
||||
'mcp_oauth',
|
||||
{},
|
||||
undefined,
|
||||
);
|
||||
});
|
||||
|
||||
|
|
@ -1416,7 +1502,7 @@ describe('MCPConnectionFactory', () => {
|
|||
}),
|
||||
);
|
||||
expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(refreshedTokens);
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled');
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'silent-refresh');
|
||||
// Silent refresh succeeded — interactive flow must NOT be initiated.
|
||||
expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled();
|
||||
});
|
||||
|
|
@ -1561,7 +1647,7 @@ describe('MCPConnectionFactory', () => {
|
|||
await oauthRequiredHandler!({ serverUrl: 'https://api.example.com' });
|
||||
|
||||
expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(refreshedTokens);
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled');
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'silent-refresh');
|
||||
// returnOnOAuth interactive path must NOT trigger when silent refresh succeeds.
|
||||
expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled();
|
||||
expect(oauthOptions.oauthStart).not.toHaveBeenCalled();
|
||||
|
|
@ -1743,7 +1829,7 @@ describe('MCPConnectionFactory', () => {
|
|||
// cached ones — that's the whole point of the fix.
|
||||
expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(freshlyRefreshedTokens);
|
||||
expect(mockConnectionInstance.setOAuthTokens).not.toHaveBeenCalledWith(staleCachedTokens);
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled');
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'silent-refresh');
|
||||
// The cached `mcp_get_tokens` flow state is dropped so the next
|
||||
// `getOAuthTokens` call reads the freshly persisted tokens from storage.
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith('flow123', 'mcp_get_tokens');
|
||||
|
|
@ -2345,6 +2431,88 @@ describe('MCPConnectionFactory', () => {
|
|||
cleanup();
|
||||
});
|
||||
|
||||
it('bounds request recovery to one silent refresh and one interactive flow', async () => {
|
||||
const sseConfig = {
|
||||
...mockServerConfig,
|
||||
url: 'https://api.example.com',
|
||||
type: 'sse' as const,
|
||||
} as t.SSEOptions;
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: sseConfig,
|
||||
};
|
||||
const refreshedTokens: MCPOAuthTokens = {
|
||||
access_token: 'refreshed-access',
|
||||
refresh_token: 'refresh-token',
|
||||
token_type: 'Bearer',
|
||||
obtained_at: Date.now(),
|
||||
};
|
||||
const interactiveTokens: MCPOAuthTokens = {
|
||||
access_token: 'interactive-access',
|
||||
token_type: 'Bearer',
|
||||
obtained_at: Date.now(),
|
||||
credential_set_id: 'interactive-generation',
|
||||
};
|
||||
mockProcessMCPEnv.mockReturnValue(sseConfig);
|
||||
mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow123');
|
||||
mockMCPTokenStorage.forceRefreshTokens.mockResolvedValueOnce(refreshedTokens);
|
||||
mockFlowManager.getFlowState.mockResolvedValue(null);
|
||||
mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValueOnce({
|
||||
authorizationUrl: 'https://auth.example.com',
|
||||
flowId: 'flow123',
|
||||
flowMetadata: {
|
||||
serverName: 'test-server',
|
||||
userId: 'user123',
|
||||
serverUrl: 'https://api.example.com',
|
||||
state: 'fresh-state',
|
||||
},
|
||||
});
|
||||
mockFlowManager.createFlow.mockResolvedValueOnce(interactiveTokens);
|
||||
|
||||
let requestOAuthHandler: ((data: Record<string, unknown>) => Promise<void>) | undefined;
|
||||
mockConnectionInstance.on.mockImplementation((event, handler) => {
|
||||
if (event === 'oauthReauthenticationRequired') {
|
||||
requestOAuthHandler = handler as (data: Record<string, unknown>) => Promise<void>;
|
||||
}
|
||||
return mockConnectionInstance;
|
||||
});
|
||||
const cleanup = MCPConnectionFactory.attachRequestOAuthHandler(
|
||||
basicOptions,
|
||||
{
|
||||
useOAuth: true,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
oauthStart: jest.fn(),
|
||||
tokenMethods: {
|
||||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
},
|
||||
mockConnectionInstance,
|
||||
);
|
||||
const challenge = {
|
||||
serverUrl: 'https://api.example.com',
|
||||
error: new Error('Non-200 status code (401)'),
|
||||
};
|
||||
|
||||
await Promise.all([requestOAuthHandler!(challenge), requestOAuthHandler!(challenge)]);
|
||||
await requestOAuthHandler!(challenge);
|
||||
await requestOAuthHandler!(challenge);
|
||||
|
||||
expect(mockMCPTokenStorage.forceRefreshTokens).toHaveBeenCalledTimes(1);
|
||||
expect(mockMCPOAuthHandler.initiateOAuthFlow).toHaveBeenCalledTimes(1);
|
||||
expect(mockConnectionInstance.setOAuthTokens).toHaveBeenNthCalledWith(1, refreshedTokens);
|
||||
expect(mockConnectionInstance.setOAuthTokens).toHaveBeenNthCalledWith(2, interactiveTokens);
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith(
|
||||
'oauthFailed',
|
||||
expect.objectContaining({ message: 'OAuth recovery phase budget exhausted' }),
|
||||
);
|
||||
|
||||
cleanup();
|
||||
});
|
||||
|
||||
it('should not reuse request-scoped OAuth callbacks after connection is cached', async () => {
|
||||
const sseConfig = {
|
||||
...mockServerConfig,
|
||||
|
|
@ -3273,7 +3441,7 @@ describe('MCPConnectionFactory', () => {
|
|||
|
||||
expect(mockMCPTokenStorage.storeTokens).not.toHaveBeenCalled();
|
||||
expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(callbackTokens);
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled');
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'interactive');
|
||||
});
|
||||
|
||||
it('rejects callback tokens that do not identify a persisted credential generation', async () => {
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -14,6 +14,7 @@ import type { OAuthTestServer } from './helpers/oauthTestServer';
|
|||
import type { MCPOAuthTokens } from '~/mcp/oauth';
|
||||
import { MCPTokenStorage, MCPOAuthHandler, ReauthenticationRequiredError } from '~/mcp/oauth';
|
||||
import { MockKeyv, createOAuthMCPServer } from './helpers/oauthTestServer';
|
||||
import { OAuthLifecycleRelay } from '~/mcp/oauth/pending';
|
||||
import { FlowStateManager } from '~/flow/manager';
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
|
|
@ -53,6 +54,44 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
});
|
||||
|
||||
describe('Fix 1: Connection mutex coalesces concurrent attempts', () => {
|
||||
it('does not overwrite a newer prompt while inspecting stored flow state', async () => {
|
||||
const ownerOAuthStart = jest.fn().mockResolvedValue(undefined);
|
||||
const waiterOAuthStart = jest.fn().mockResolvedValue(undefined);
|
||||
let resolveFlow: ((flow: object) => void) | undefined;
|
||||
const flowManager = {
|
||||
getFlowState: jest.fn(
|
||||
() =>
|
||||
new Promise<object>((resolve) => {
|
||||
resolveFlow = resolve;
|
||||
}),
|
||||
),
|
||||
};
|
||||
const relay = new OAuthLifecycleRelay({
|
||||
oauthStart: ownerOAuthStart,
|
||||
logPrefix: '[MCP][test]',
|
||||
});
|
||||
|
||||
await relay.start('https://auth.example.com/old');
|
||||
const addWaiter = relay.add({
|
||||
oauthStart: waiterOAuthStart,
|
||||
flowManager: flowManager as never,
|
||||
userId: 'user-1',
|
||||
serverName: 'test-server',
|
||||
});
|
||||
|
||||
expect(flowManager.getFlowState).toHaveBeenCalledTimes(1);
|
||||
await relay.start('https://auth.example.com/new');
|
||||
resolveFlow?.({
|
||||
createdAt: Date.now(),
|
||||
metadata: { authorizationUrl: 'https://auth.example.com/old' },
|
||||
status: 'PENDING',
|
||||
});
|
||||
await addWaiter;
|
||||
|
||||
expect(waiterOAuthStart).toHaveBeenCalledTimes(1);
|
||||
expect(waiterOAuthStart).toHaveBeenCalledWith('https://auth.example.com/new', undefined);
|
||||
});
|
||||
|
||||
it('should return the same pending promise for concurrent getUserConnection calls', async () => {
|
||||
const { UserConnectionManager } = await import('~/mcp/UserConnectionManager');
|
||||
|
||||
|
|
@ -283,6 +322,9 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
await oauthOptions.oauthStart?.(authorizationUrl);
|
||||
}
|
||||
await connectionReleased;
|
||||
if (oauthOptions && 'oauthEnd' in oauthOptions) {
|
||||
await oauthOptions.oauthEnd?.();
|
||||
}
|
||||
return mockConnection as never;
|
||||
});
|
||||
|
||||
|
|
@ -297,11 +339,13 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
await flowManager.initFlow(`${user.id}:${serverName}`, 'mcp_oauth', { authorizationUrl });
|
||||
|
||||
const firstOAuthStart = jest.fn().mockResolvedValue(undefined);
|
||||
const firstOAuthEnd = jest.fn().mockRejectedValue(new Error('owner response is stale'));
|
||||
const firstConnection = manager.getUserConnection({
|
||||
serverName,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
oauthStart: firstOAuthStart,
|
||||
oauthEnd: firstOAuthEnd,
|
||||
});
|
||||
for (let i = 0; i < 20 && firstOAuthStart.mock.calls.length === 0; i++) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
|
|
@ -309,21 +353,29 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
expect(firstOAuthStart).toHaveBeenCalledWith(authorizationUrl, undefined);
|
||||
|
||||
const joinedOAuthStart = jest.fn().mockResolvedValue(undefined);
|
||||
const joinedOAuthEnd = jest.fn().mockResolvedValue(undefined);
|
||||
const joinedConnection = manager.getUserConnection({
|
||||
serverName,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
oauthStart: joinedOAuthStart,
|
||||
oauthEnd: joinedOAuthEnd,
|
||||
});
|
||||
|
||||
for (let i = 0; i < 20 && joinedOAuthStart.mock.calls.length === 0; i++) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
}
|
||||
expect(joinedOAuthStart).toHaveBeenCalledWith(
|
||||
authorizationUrl,
|
||||
expect.objectContaining({ expiresAt: expect.any(Number) }),
|
||||
);
|
||||
|
||||
releaseConnection();
|
||||
const [conn1, conn2] = await Promise.all([firstConnection, joinedConnection]);
|
||||
|
||||
expect(conn1).toBe(conn2);
|
||||
expect(joinedOAuthStart).toHaveBeenCalledWith(
|
||||
authorizationUrl,
|
||||
expect.objectContaining({ expiresAt: expect.any(Number) }),
|
||||
);
|
||||
expect(firstOAuthEnd).toHaveBeenCalledTimes(1);
|
||||
expect(joinedOAuthEnd).toHaveBeenCalledTimes(1);
|
||||
expect(createSpy).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
releaseConnection();
|
||||
|
|
|
|||
|
|
@ -70,6 +70,10 @@ export interface OAuthTestServerOptions {
|
|||
scopesSupported?: string[];
|
||||
/** When true, /authorize and /token reject requests that omit the MCP resource parameter. */
|
||||
requireResourceParameter?: boolean;
|
||||
/** Number of refresh-grant access tokens the MCP resource should reject after issuance. */
|
||||
rejectRefreshTokens?: number;
|
||||
/** Optional test hook for controlling echo-tool completion. */
|
||||
echoHandler?: (message: string) => string | Promise<string>;
|
||||
}
|
||||
|
||||
export interface OAuthTokenRequestRecord {
|
||||
|
|
@ -136,6 +140,8 @@ export async function createOAuthMCPServer(
|
|||
requiredScopes = [],
|
||||
scopesSupported = [...new Set([...tokenScopes, ...requiredScopes])],
|
||||
requireResourceParameter = false,
|
||||
rejectRefreshTokens = 0,
|
||||
echoHandler,
|
||||
} = options;
|
||||
|
||||
const sessions = new Map<string, StreamableHTTPServerTransport>();
|
||||
|
|
@ -157,6 +163,7 @@ export async function createOAuthMCPServer(
|
|||
}
|
||||
>();
|
||||
const registeredClients = new Map<string, { client_id: string; client_secret: string }>();
|
||||
let rejectedRefreshTokensRemaining = rejectRefreshTokens;
|
||||
|
||||
let port = 0;
|
||||
const getBaseUrl = () => `http://127.0.0.1:${port}`;
|
||||
|
|
@ -410,7 +417,11 @@ export async function createOAuthMCPServer(
|
|||
const scopes = params.has('scope')
|
||||
? parseScopes(params.get('scope'))
|
||||
: (refreshTokenScopes.get(refreshToken) ?? tokenScopes);
|
||||
issuedTokens.add(newAccessToken);
|
||||
if (rejectedRefreshTokensRemaining > 0) {
|
||||
rejectedRefreshTokensRemaining -= 1;
|
||||
} else {
|
||||
issuedTokens.add(newAccessToken);
|
||||
}
|
||||
tokenIssueTimes.set(newAccessToken, Date.now());
|
||||
accessTokenScopes.set(newAccessToken, scopes);
|
||||
|
||||
|
|
@ -475,9 +486,10 @@ export async function createOAuthMCPServer(
|
|||
sessionIdGenerator: () => randomUUID(),
|
||||
});
|
||||
const mcp = new McpServer({ name: 'oauth-test-server', version: '0.0.1' });
|
||||
mcp.tool('echo', { message: z.string() }, async (args) => ({
|
||||
content: [{ type: 'text' as const, text: `echo: ${args.message}` }],
|
||||
}));
|
||||
mcp.tool('echo', { message: z.string() }, async (args) => {
|
||||
const text = echoHandler ? await echoHandler(args.message) : `echo: ${args.message}`;
|
||||
return { content: [{ type: 'text' as const, text }] };
|
||||
});
|
||||
await mcp.connect(transport);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -66,17 +66,42 @@ describe('MCP request context', () => {
|
|||
const res = createResponse();
|
||||
const context = getMCPRequestContext(req, res);
|
||||
const disconnect = jest.fn().mockResolvedValue(undefined);
|
||||
const dispose = jest.fn().mockResolvedValue(undefined);
|
||||
const pendingDisconnect = jest.fn().mockResolvedValue(undefined);
|
||||
const pendingDispose = jest.fn().mockResolvedValue(undefined);
|
||||
|
||||
context?.connections.set('server', { disconnect });
|
||||
context?.pending.set('pending-server', Promise.resolve({ disconnect: pendingDisconnect }));
|
||||
context?.connections.set('server', { disconnect, dispose });
|
||||
context?.pending.set(
|
||||
'pending-server',
|
||||
Promise.resolve({ disconnect: pendingDisconnect, dispose: pendingDispose }),
|
||||
);
|
||||
|
||||
res.emit('finish');
|
||||
await nextTick();
|
||||
|
||||
expect(disconnect).toHaveBeenCalledTimes(1);
|
||||
expect(pendingDisconnect).toHaveBeenCalledTimes(1);
|
||||
expect(dispose).toHaveBeenCalledTimes(1);
|
||||
expect(pendingDispose).toHaveBeenCalledTimes(1);
|
||||
expect(disconnect).not.toHaveBeenCalled();
|
||||
expect(pendingDisconnect).not.toHaveBeenCalled();
|
||||
expect(context?.connections.size).toBe(0);
|
||||
expect(context?.pending.size).toBe(0);
|
||||
});
|
||||
|
||||
it('uses the lifecycle disposer supplied by the connection manager', async () => {
|
||||
const req = {};
|
||||
const res = createResponse();
|
||||
const context = getMCPRequestContext(req, res);
|
||||
const connection = { disconnect: jest.fn().mockResolvedValue(undefined) };
|
||||
const disposeConnection = jest.fn().mockResolvedValue(undefined);
|
||||
if (context) {
|
||||
context.disposeConnection = disposeConnection;
|
||||
context.connections.set('user:server', connection);
|
||||
}
|
||||
|
||||
res.emit('close');
|
||||
await nextTick();
|
||||
|
||||
expect(disposeConnection).toHaveBeenCalledWith('user:server', connection);
|
||||
expect(connection.disconnect).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ import type * as t from './types';
|
|||
import { createSSRFSafeUndiciConnect, isSSRFTarget, resolveHostnameSSRF } from '~/auth';
|
||||
import { reserveMCPToolsChangedRevision } from './toolsChanged';
|
||||
import { isOAuthServer, sanitizeUrlForLogging } from './utils';
|
||||
import { isOAuthAuthenticationError } from './errors';
|
||||
import { runOutsideTracing } from '~/utils/tracing';
|
||||
import { isAddressAllowed } from '~/auth/domain';
|
||||
import { withTimeout } from '~/utils/promise';
|
||||
|
|
@ -1151,6 +1152,7 @@ export class MCPConnection extends EventEmitter {
|
|||
private readonly userId?: string;
|
||||
private lastPingTime: number;
|
||||
private lastConnectionCheckAt: number = 0;
|
||||
private lastConnectionCheckError?: unknown;
|
||||
private oauthTokens?: MCPOAuthTokens | null;
|
||||
private requestHeaders?: Record<string, string> | null;
|
||||
private oauthRequired = false;
|
||||
|
|
@ -1777,6 +1779,7 @@ export class MCPConnection extends EventEmitter {
|
|||
this.on('connectionChange', (state: t.ConnectionState) => {
|
||||
this.connectionState = state;
|
||||
if (state === 'connected') {
|
||||
this.lastConnectionCheckError = undefined;
|
||||
const isReconnect = this.hasConnected;
|
||||
this.hasConnected = true;
|
||||
this.toolListRefreshSuspended = false;
|
||||
|
|
@ -2114,7 +2117,7 @@ export class MCPConnection extends EventEmitter {
|
|||
}
|
||||
|
||||
// Check if it's an OAuth authentication error
|
||||
if (this.isOAuthError(error)) {
|
||||
if (isOAuthAuthenticationError(error)) {
|
||||
logger.warn(`${this.getLogPrefix()} OAuth authentication required`);
|
||||
this.oauthRequired = true;
|
||||
const serverUrl = this.url;
|
||||
|
|
@ -2310,9 +2313,12 @@ export class MCPConnection extends EventEmitter {
|
|||
}
|
||||
|
||||
// Check if it's an OAuth authentication error
|
||||
if (this.isOAuthError(error)) {
|
||||
if (isOAuthAuthenticationError(error)) {
|
||||
logger.warn(`${this.getLogPrefix()} OAuth authentication error detected`);
|
||||
this.lastConnectionCheckError = error;
|
||||
this.connectionState = 'error';
|
||||
this.emit('oauthError', error);
|
||||
return;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -2644,6 +2650,7 @@ export class MCPConnection extends EventEmitter {
|
|||
return true;
|
||||
}
|
||||
this.lastConnectionCheckAt = now;
|
||||
this.lastConnectionCheckError = undefined;
|
||||
|
||||
try {
|
||||
// Try ping first as it's the lightest check
|
||||
|
|
@ -2660,6 +2667,7 @@ export class MCPConnection extends EventEmitter {
|
|||
(error as Error)?.message.includes('method not found'));
|
||||
|
||||
if (!pingUnsupported) {
|
||||
this.lastConnectionCheckError = error;
|
||||
logger.error(`${this.getLogPrefix()} Ping failed:`, error);
|
||||
return false;
|
||||
}
|
||||
|
|
@ -2692,6 +2700,7 @@ export class MCPConnection extends EventEmitter {
|
|||
}
|
||||
} catch (capabilityError) {
|
||||
// If capability check fails, the connection is likely broken
|
||||
this.lastConnectionCheckError = capabilityError;
|
||||
logger.error(`${this.getLogPrefix()} Connection verification failed:`, capabilityError);
|
||||
return false;
|
||||
}
|
||||
|
|
@ -2707,6 +2716,14 @@ export class MCPConnection extends EventEmitter {
|
|||
return isOAuthServer(this.options);
|
||||
}
|
||||
|
||||
public isOAuthAuthenticationError(error: unknown): boolean {
|
||||
return isOAuthAuthenticationError(error);
|
||||
}
|
||||
|
||||
public getLastConnectionCheckError(): unknown {
|
||||
return this.lastConnectionCheckError;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if this connection is stale compared to config update time.
|
||||
* A connection is stale if it was created before the config was updated.
|
||||
|
|
@ -2718,47 +2735,6 @@ export class MCPConnection extends EventEmitter {
|
|||
return this.createdAt < configUpdatedAt;
|
||||
}
|
||||
|
||||
private isOAuthError(error: unknown): boolean {
|
||||
if (!error || typeof error !== 'object') {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check for error code
|
||||
if ('code' in error) {
|
||||
const code = (error as { code?: number }).code;
|
||||
if (code === 401 || code === 403) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// Check message for various auth error indicators
|
||||
if ('message' in error && typeof error.message === 'string') {
|
||||
const message = error.message.toLowerCase();
|
||||
// Check for 401 status
|
||||
if (message.includes('401') || message.includes('non-200 status code (401)')) {
|
||||
return true;
|
||||
}
|
||||
// Check for invalid_token (OAuth servers return this for expired/revoked tokens)
|
||||
if (message.includes('invalid_token')) {
|
||||
return true;
|
||||
}
|
||||
// Check for invalid_grant (OAuth servers return this for expired/revoked grants)
|
||||
if (message.includes('invalid_grant')) {
|
||||
return true;
|
||||
}
|
||||
// Check for authentication required
|
||||
if (message.includes('authentication required') || message.includes('unauthorized')) {
|
||||
return true;
|
||||
}
|
||||
// Check for missing authorization values (e.g., Amazon Ads MCP returns HTTP 400 with this)
|
||||
if (message.includes('no authorization')) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if an error indicates rate limiting (HTTP 429).
|
||||
* Rate limited requests should stop reconnection attempts to avoid making the situation worse.
|
||||
|
|
|
|||
|
|
@ -10,6 +10,47 @@ export const MCPErrorCodes = {
|
|||
|
||||
export type MCPErrorCode = (typeof MCPErrorCodes)[keyof typeof MCPErrorCodes];
|
||||
|
||||
interface OAuthErrorLike {
|
||||
code?: number;
|
||||
status?: number;
|
||||
statusCode?: number;
|
||||
message?: string;
|
||||
}
|
||||
|
||||
const OAUTH_HTTP_STATUS_PATTERN =
|
||||
/(?:\bhttp\s+(?:401|403)\b|\bnon-2\d\d\s+status\s+code\s*\((?:401|403)\)|^(?:error:\s*)?(?:401|403)\b|\bunauthorized\s*\(\s*401\s*\)|\bforbidden\s*\(\s*403\s*\))/i;
|
||||
const MISSING_AUTHORIZATION_PATTERN = /\bno authorization (?:headers?|values?)\b/i;
|
||||
|
||||
/** Detects HTTP authentication failures and OAuth protocol errors without matching unrelated IDs. */
|
||||
export function isOAuthAuthenticationError(error: unknown): boolean {
|
||||
if (!error || typeof error !== 'object') {
|
||||
return false;
|
||||
}
|
||||
|
||||
const candidate = error as OAuthErrorLike;
|
||||
if (
|
||||
[candidate.status, candidate.statusCode, candidate.code].some(
|
||||
(status) => status === 401 || status === 403,
|
||||
)
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (typeof candidate.message !== 'string') {
|
||||
return false;
|
||||
}
|
||||
|
||||
const message = candidate.message.toLowerCase();
|
||||
return (
|
||||
OAUTH_HTTP_STATUS_PATTERN.test(message) ||
|
||||
message.includes('invalid_token') ||
|
||||
message.includes('invalid_grant') ||
|
||||
message.includes('insufficient_scope') ||
|
||||
message.includes('authentication required') ||
|
||||
MISSING_AUTHORIZATION_PATTERN.test(message)
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Custom error for MCP domain restriction violations.
|
||||
* Thrown when a user attempts to connect to an MCP server whose domain is not in the allowlist.
|
||||
|
|
|
|||
|
|
@ -68,3 +68,142 @@ export async function getReplayablePendingMCPOAuthStart({
|
|||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
type OAuthEndHandler = () => Promise<void>;
|
||||
|
||||
export class OAuthLifecycleRelay {
|
||||
private readonly oauthStarts = new Set<t.OAuthStartHandler>();
|
||||
private readonly oauthEnds = new Set<OAuthEndHandler>();
|
||||
private readonly emittedAuthUrls = new WeakMap<t.OAuthStartHandler, string>();
|
||||
private readonly emittedOAuthEnds = new WeakSet<OAuthEndHandler>();
|
||||
private lastOAuthStart?: PendingOAuthStart;
|
||||
private oauthEnded = false;
|
||||
|
||||
constructor({
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
logPrefix,
|
||||
}: {
|
||||
oauthStart?: t.OAuthStartHandler;
|
||||
oauthEnd?: OAuthEndHandler;
|
||||
logPrefix: string;
|
||||
}) {
|
||||
this.logPrefix = logPrefix;
|
||||
if (oauthStart) {
|
||||
this.oauthStarts.add(oauthStart);
|
||||
}
|
||||
if (oauthEnd) {
|
||||
this.oauthEnds.add(oauthEnd);
|
||||
}
|
||||
}
|
||||
|
||||
private readonly logPrefix: string;
|
||||
|
||||
public readonly start: t.OAuthStartHandler = async (authURL, options) => {
|
||||
this.lastOAuthStart = { authURL, options };
|
||||
const errors: unknown[] = [];
|
||||
let delivered = false;
|
||||
|
||||
for (const oauthStart of Array.from(this.oauthStarts)) {
|
||||
try {
|
||||
await this.emit(oauthStart, authURL, options);
|
||||
delivered = true;
|
||||
} catch (error) {
|
||||
errors.push(error);
|
||||
logger.warn(`${this.logPrefix} Failed to notify OAuth prompt listener`, error);
|
||||
}
|
||||
}
|
||||
|
||||
if (!delivered && errors.length > 0) {
|
||||
throw errors[0];
|
||||
}
|
||||
};
|
||||
|
||||
/** Completion notifications are best effort and cannot invalidate received OAuth tokens. */
|
||||
public readonly end: OAuthEndHandler = async () => {
|
||||
this.oauthEnded = true;
|
||||
for (const oauthEnd of Array.from(this.oauthEnds)) {
|
||||
try {
|
||||
await this.emitEnd(oauthEnd);
|
||||
} catch (error) {
|
||||
logger.warn(`${this.logPrefix} Failed to notify OAuth completion listener`, error);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
public async add({
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
flowManager,
|
||||
userId,
|
||||
serverName,
|
||||
}: ReplayablePendingMCPOAuthStartOptions & {
|
||||
oauthStart?: t.OAuthStartHandler;
|
||||
oauthEnd?: OAuthEndHandler;
|
||||
}): Promise<void> {
|
||||
if (oauthStart) {
|
||||
this.oauthStarts.add(oauthStart);
|
||||
}
|
||||
if (oauthEnd) {
|
||||
this.oauthEnds.add(oauthEnd);
|
||||
}
|
||||
if (this.oauthEnded) {
|
||||
if (oauthEnd) {
|
||||
try {
|
||||
await this.emitEnd(oauthEnd);
|
||||
} catch (error) {
|
||||
logger.warn(`${this.logPrefix} Failed to re-issue OAuth completion`, error);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (!oauthStart) {
|
||||
return;
|
||||
}
|
||||
const lastOAuthStart = this.lastOAuthStart;
|
||||
const storedOAuthStart =
|
||||
!lastOAuthStart || lastOAuthStart.options?.expiresAt == null
|
||||
? await getReplayablePendingMCPOAuthStart({ flowManager, userId, serverName })
|
||||
: undefined;
|
||||
const currentOAuthStart = this.lastOAuthStart;
|
||||
const replayOAuthStart =
|
||||
storedOAuthStart &&
|
||||
(!currentOAuthStart || storedOAuthStart.authURL === currentOAuthStart.authURL)
|
||||
? storedOAuthStart
|
||||
: currentOAuthStart;
|
||||
if (!replayOAuthStart) {
|
||||
return;
|
||||
}
|
||||
if (this.oauthEnded) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.lastOAuthStart = replayOAuthStart;
|
||||
try {
|
||||
await this.emit(oauthStart, replayOAuthStart.authURL, replayOAuthStart.options);
|
||||
} catch (error) {
|
||||
logger.warn(`${this.logPrefix} Failed to re-issue pending OAuth URL`, error);
|
||||
}
|
||||
}
|
||||
|
||||
private async emit(
|
||||
oauthStart: t.OAuthStartHandler,
|
||||
authURL: string,
|
||||
options?: t.OAuthStartOptions,
|
||||
): Promise<void> {
|
||||
if (this.emittedAuthUrls.get(oauthStart) === authURL) {
|
||||
return;
|
||||
}
|
||||
this.emittedAuthUrls.set(oauthStart, authURL);
|
||||
await oauthStart(authURL, options);
|
||||
}
|
||||
|
||||
private async emitEnd(oauthEnd: OAuthEndHandler): Promise<void> {
|
||||
if (this.emittedOAuthEnds.has(oauthEnd)) {
|
||||
return;
|
||||
}
|
||||
this.emittedOAuthEnds.add(oauthEnd);
|
||||
await oauthEnd();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ interface MCPResponseLike {
|
|||
|
||||
interface Disconnectable {
|
||||
disconnect: () => Promise<unknown> | unknown;
|
||||
dispose?: () => Promise<unknown> | unknown;
|
||||
}
|
||||
|
||||
const contexts = new WeakMap<object, MCPRequestContext>();
|
||||
|
|
@ -50,29 +51,36 @@ export async function cleanupMCPRequestContext(context?: MCPRequestContext): Pro
|
|||
}
|
||||
context.cleanupStarted = true;
|
||||
|
||||
const connections = new Set<Disconnectable>();
|
||||
for (const connection of context.connections.values()) {
|
||||
const connections = new Map<Disconnectable, string>();
|
||||
for (const [connectionKey, connection] of context.connections) {
|
||||
if (isDisconnectable(connection)) {
|
||||
connections.add(connection);
|
||||
connections.set(connection, connectionKey);
|
||||
}
|
||||
}
|
||||
|
||||
const pending = Array.from(context.pending.values());
|
||||
const pending = Array.from(context.pending.entries());
|
||||
if (pending.length > 0) {
|
||||
const settled = await Promise.allSettled(pending);
|
||||
for (const result of settled) {
|
||||
const settled = await Promise.allSettled(pending.map(([, promise]) => promise));
|
||||
for (let index = 0; index < settled.length; index++) {
|
||||
const result = settled[index];
|
||||
if (result.status === 'fulfilled' && isDisconnectable(result.value)) {
|
||||
connections.add(result.value);
|
||||
connections.set(result.value, pending[index][0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
await Promise.allSettled(
|
||||
Array.from(connections).map(async (connection) => {
|
||||
Array.from(connections).map(async ([connection, connectionKey]) => {
|
||||
try {
|
||||
await connection.disconnect();
|
||||
if (context.disposeConnection) {
|
||||
await context.disposeConnection(connectionKey, connection);
|
||||
} else if (connection.dispose) {
|
||||
await connection.dispose();
|
||||
} else {
|
||||
await connection.disconnect();
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn('[MCP Request Context] Failed to disconnect request-scoped connection', error);
|
||||
logger.warn('[MCP Request Context] Failed to dispose request-scoped connection', error);
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -62,6 +62,8 @@ export interface MCPPrompt {
|
|||
|
||||
export type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'error';
|
||||
|
||||
export type OAuthHandledSource = 'silent-refresh' | 'interactive';
|
||||
|
||||
export type MCPTool = Tool;
|
||||
export type MCPToolListResponse = ListToolsResult;
|
||||
export type ToolContentPart = TextContent | ImageContent | EmbeddedResource | AudioContent;
|
||||
|
|
@ -211,6 +213,7 @@ export interface UserConnectionContext {
|
|||
export interface RequestScopedMCPConnectionStore {
|
||||
connections: Map<string, unknown>;
|
||||
pending: Map<string, Promise<unknown>>;
|
||||
disposeConnection?: (connectionKey: string, connection: unknown) => Promise<void>;
|
||||
}
|
||||
|
||||
export interface OAuthStartOptions {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue