From 1bccc2bc1806a9bc94dd6b90522e31930f6d2c84 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 8 Aug 2026 13:50:21 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=A1=20fix:=20Refresh=20MCP=20Tools=20A?= =?UTF-8?q?fter=20List-Changed=20Notifications=20(#14686)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): handle dynamic tool list changes Co-authored-by: Pascal Garber * test(mcp): fix CI validation * fix(mcp): keep dynamic tool catalogs live * fix(mcp): harden dynamic catalog lifecycle * test(mcp): use typed startup connection * test(mcp): isolate dynamic e2e fixtures * fix(mcp): refresh tools after reconnect * fix(mcp): close dynamic catalog cache gaps * test(mcp): update OAuth connection mocks * fix(mcp): preserve app snapshot ownership * style(mcp): sort connection imports * fix(mcp): close review race conditions * fix(mcp): preserve cache ownership edges * fix(mcp): harden recovery lifecycle * fix(mcp): guard tool-less app refresh * fix(mcp): fence distributed cache races * fix(mcp): retire stale connection state * fix(mcp): keep tool snapshots authoritative * fix(mcp): fence stale app tool publications * style(mcp): sort repository test imports * test(mcp): mock empty startup publication * fix(mcp): preserve app publication generations * fix(mcp): harden publication recovery races * fix(mcp): address tool catalogs by runtime config * fix(mcp): load scoped catalogs for assistant writes * fix(mcp): harden catalog publication recovery * fix(mcp): serialize forced connection replacement * fix(mcp): serialize ordinary creation with replacements * fix(mcp): harden catalog fallback boundaries * fix(mcp): close lifecycle fencing gaps * fix(mcp): preserve catalog authority on failures * fix(mcp): compensate failed catalog mutations * fix(mcp): fence catalog refresh ordering * style(mcp): sort agent loader imports * fix(mcp): cancel stale connection creation * fix(mcp): fence catalog coordination * fix(mcp): close catalog race windows * fix(mcp): harden cross-pod catalog fencing * fix(mcp): close catalog lifecycle edges * style(mcp): sort assistant imports * fix(mcp): reject stale recovery authority * fix(mcp): restore static catalog on every startup * fix(mcp): order app catalog publications * style(mcp): sort catalog revision imports * fix(mcp): separate catalog allocation and commit fences --------- Co-authored-by: Pascal Garber --- .github/workflows/playwright-mock.yml | 118 ++- api/server/controllers/UserController.js | 50 +- .../__tests__/UserController.mcpOAuth.spec.js | 32 + .../__tests__/deleteUserMcpServers.spec.js | 40 + .../controllers/__tests__/mcp.servers.spec.js | 235 ++++- api/server/controllers/assistants/v1.js | 7 +- api/server/controllers/assistants/v2.js | 7 +- api/server/controllers/mcp.js | 137 ++- api/server/index.metrics.spec.js | 1 + api/server/index.spec.js | 1 + api/server/routes/__tests__/mcp.spec.js | 260 ++++- api/server/routes/mcp.js | 32 +- .../__tests__/getCachedTools.lock.spec.js | 341 +++++++ .../Config/__tests__/getCachedTools.spec.js | 481 +++++++++- api/server/services/Config/getCachedTools.js | 103 +- api/server/services/Config/mcp.js | 42 +- api/server/services/MCP.js | 63 +- api/server/services/Tools/mcp.js | 52 +- api/server/services/Tools/mcp.spec.js | 77 +- api/server/services/__tests__/MCP.spec.js | 111 ++- api/server/services/initializeMCPs.js | 60 +- api/server/services/initializeMCPs.spec.js | 151 ++- e2e/config/librechat.e2e.yaml | 3 + e2e/playwright.config.mock.ts | 90 +- e2e/setup/dynamic-mcp-tools.js | 130 +++ e2e/setup/env.ts | 16 + e2e/setup/fake-mcp-dynamic-network-server.js | 125 +++ e2e/setup/fake-mcp-server.js | 14 + e2e/setup/start-server-cluster.js | 165 ++++ e2e/specs/mock/mcp-fixture-isolation.spec.ts | 17 + e2e/specs/mock/mcp-tool-list-changed.spec.ts | 161 ++++ .../api/src/agents/__tests__/load.spec.ts | 59 +- packages/api/src/agents/added.ts | 13 +- packages/api/src/agents/load.ts | 15 +- packages/api/src/index.ts | 3 + packages/api/src/mcp/ConnectionsRepository.ts | 167 +++- packages/api/src/mcp/MCPConnectionFactory.ts | 21 +- packages/api/src/mcp/MCPManager.ts | 126 ++- packages/api/src/mcp/UserConnectionManager.ts | 349 ++++++- .../__tests__/ConnectionsRepository.test.ts | 224 ++++- .../__tests__/MCPConnectionFactory.test.ts | 31 +- .../__tests__/MCPConnectionFetchTools.test.ts | 51 +- .../api/src/mcp/__tests__/MCPManager.test.ts | 712 +++++++++++++- .../__tests__/MCPOAuthRaceCondition.test.ts | 171 ++++ .../toolListChanged.integration.test.ts | 362 +++++++ packages/api/src/mcp/__tests__/utils.test.ts | 24 + packages/api/src/mcp/assistants.spec.ts | 199 ++++ packages/api/src/mcp/assistants.ts | 183 ++++ packages/api/src/mcp/catalog/store.ts | 886 ++++++++++++++++++ packages/api/src/mcp/connection.ts | 260 ++++- .../src/mcp/registry/MCPServerInspector.ts | 6 +- .../src/mcp/registry/MCPServersRegistry.ts | 42 +- .../__tests__/MCPServerInspector.test.ts | 76 +- .../__tests__/MCPServersRegistry.test.ts | 43 + .../__tests__/mcpConnectionsMock.helper.ts | 3 + packages/api/src/mcp/tools.spec.ts | 741 ++++++++------- packages/api/src/mcp/tools.ts | 440 +++++++-- packages/api/src/mcp/toolsChanged.spec.ts | 258 +++++ packages/api/src/mcp/toolsChanged.ts | 265 ++++++ packages/api/src/mcp/utils.ts | 25 +- 60 files changed, 8046 insertions(+), 831 deletions(-) create mode 100644 api/server/services/Config/__tests__/getCachedTools.lock.spec.js create mode 100644 e2e/setup/dynamic-mcp-tools.js create mode 100644 e2e/setup/fake-mcp-dynamic-network-server.js create mode 100644 e2e/setup/start-server-cluster.js create mode 100644 e2e/specs/mock/mcp-fixture-isolation.spec.ts create mode 100644 e2e/specs/mock/mcp-tool-list-changed.spec.ts create mode 100644 packages/api/src/mcp/__tests__/toolListChanged.integration.test.ts create mode 100644 packages/api/src/mcp/assistants.spec.ts create mode 100644 packages/api/src/mcp/assistants.ts create mode 100644 packages/api/src/mcp/catalog/store.ts create mode 100644 packages/api/src/mcp/toolsChanged.spec.ts create mode 100644 packages/api/src/mcp/toolsChanged.ts diff --git a/.github/workflows/playwright-mock.yml b/.github/workflows/playwright-mock.yml index 10c9f4101c..16779dff04 100644 --- a/.github/workflows/playwright-mock.yml +++ b/.github/workflows/playwright-mock.yml @@ -164,6 +164,116 @@ jobs: retention-days: 7 if-no-files-found: ignore + mcp_tool_list_changed: + name: MCP list_changed (replica count ${{ matrix.replicas }}) + if: >- + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + github.event.pull_request != null && + contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.pull_request.author_association)) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + replicas: [1, 2] + env: + CI: 'true' + E2E_CHROMIUM_CHANNEL: chrome + E2E_MCP_LIST_CHANGED: 'true' + E2E_REPLICAS: ${{ matrix.replicas }} + steps: + - uses: actions/checkout@v4 + + - name: Use Node.js 24.16.0 + uses: actions/setup-node@v4 + with: + node-version: '24.16.0' + + - name: Restore node_modules cache + id: cache-node-modules + uses: actions/cache@v4 + with: + path: | + node_modules + client/node_modules + packages/client/node_modules + packages/data-provider/node_modules + packages/data-schemas/node_modules + packages/api/node_modules + api/node_modules + key: node-modules-e2e-${{ runner.os }}-24.16.0-${{ hashFiles('package-lock.json') }} + + - name: Install dependencies + if: steps.cache-node-modules.outputs.cache-hit != 'true' + run: npm ci + + - name: Build e2e dependencies + run: npm run e2e:prepare + + - name: Install Playwright and Redis runtime dependencies + timeout-minutes: 5 + run: | + google-chrome --version + npx playwright install-deps chrome + sudo apt-get update + sudo apt-get install -y redis-server redis-tools + + - name: Start standalone Redis and Redis Cluster + run: | + redis-server --daemonize yes --port 6379 + redis-cli -p 6379 ping + chmod +x redis-config/start-cluster.sh redis-config/stop-cluster.sh + ./redis-config/start-cluster.sh + redis-cli -p 7001 cluster info + + - name: Test MCP notifications with in-memory cache + env: + E2E_STREAM_STORE: memory + run: >- + npx playwright test --config=e2e/playwright.config.mock.ts + mcp-tool-list-changed.spec.ts --retries=0 + + - name: Test MCP notifications with standalone Redis cache + env: + E2E_STREAM_STORE: redis + REDIS_URI: redis://127.0.0.1:6379 + run: >- + npx playwright test --config=e2e/playwright.config.mock.ts + mcp-tool-list-changed.spec.ts --retries=0 + + - name: Test MCP notifications with Redis Cluster cache + env: + E2E_STREAM_STORE: redis-cluster + REDIS_URI: redis://127.0.0.1:7001,redis://127.0.0.1:7002,redis://127.0.0.1:7003 + run: >- + npx playwright test --config=e2e/playwright.config.mock.ts + mcp-tool-list-changed.spec.ts --retries=0 + + - name: Upload Playwright HTML report + if: ${{ !cancelled() }} + uses: actions/upload-artifact@v4 + with: + name: playwright-mcp-list-changed-${{ matrix.replicas }}-replicas + path: e2e/playwright-report/** + retention-days: 7 + if-no-files-found: ignore + + - name: Upload traces & screenshots + if: failure() + uses: actions/upload-artifact@v4 + with: + name: playwright-mcp-list-changed-results-${{ matrix.replicas }}-replicas + path: e2e/specs/.test-results/** + retention-days: 7 + if-no-files-found: ignore + + - name: Stop Redis processes + if: always() + run: | + ./redis-config/stop-cluster.sh || true + redis-cli -p 6379 shutdown || true + e2e: name: e2e if: >- @@ -172,9 +282,11 @@ jobs: (github.event_name == 'pull_request' && github.event.pull_request != null && contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.pull_request.author_association))) - needs: e2e_shards + needs: [e2e_shards, mcp_tool_list_changed] runs-on: ubuntu-latest steps: - - name: Verify every Playwright shard passed - if: needs.e2e_shards.result != 'success' + - name: Verify every Playwright job passed + if: >- + needs.e2e_shards.result != 'success' || + needs.mcp_tool_list_changed.result != 'success' run: exit 1 diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index d341b182ee..8cc32e4a85 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -177,14 +177,24 @@ const deleteUserMcpServers = async (userId) => { const allServersToDelete = [...aclOwnedServers, ...legacyServers]; const mcpManager = getMCPManager(); - if (mcpManager) { - await Promise.all( - allServersToDelete.map(async (s) => { - await mcpManager.disconnectUserConnection(userId, s.serverName); + await Promise.allSettled( + allServersToDelete.map(async (s) => { + try { await invalidateCachedTools({ userId, serverName: s.serverName }); - }), - ); - } + } catch (error) { + logger.warn( + `[deleteUserMcpServers] Failed to invalidate tools for ${s.serverName}:`, + error, + ); + } finally { + try { + await mcpManager?.disconnectUserConnection(userId, s.serverName); + } catch (error) { + logger.warn(`[deleteUserMcpServers] Failed to disconnect ${s.serverName}:`, error); + } + } + }), + ); await AclEntry.deleteMany({ resourceType: ResourceType.MCPSERVER, @@ -295,21 +305,37 @@ const updateUserPluginsController = async (req, res) => { if (pluginKey.startsWith(Constants.mcp_prefix)) { try { const mcpManager = getMCPManager(); + // Extract server name from pluginKey (format: "mcp_") + const serverName = pluginKey.replace(Constants.mcp_prefix, ''); if (mcpManager) { - // Extract server name from pluginKey (format: "mcp_") - const serverName = pluginKey.replace(Constants.mcp_prefix, ''); logger.info( `[updateUserPluginsController] Attempting disconnect of MCP server "${serverName}" for user ${user.id} after plugin auth update.`, ); - await mcpManager.disconnectUserConnection(user.id, serverName); + } + let invalidationError; + try { await invalidateCachedTools({ userId: user.id, serverName }); + } catch (error) { + invalidationError = error; + } + try { + await mcpManager?.disconnectUserConnection(user.id, serverName); + } catch (error) { + logger.error( + `[updateUserPluginsController] Error disconnecting MCP connection for user ${user.id} after plugin auth update:`, + error, + ); + } + if (invalidationError) { + throw invalidationError; } } catch (disconnectError) { logger.error( - `[updateUserPluginsController] Error disconnecting MCP connection for user ${user.id} after plugin auth update:`, + `[updateUserPluginsController] Error fencing MCP connection for user ${user.id} after plugin auth update:`, disconnectError, ); - // Do not fail the request for this, but log it. + // A credential mutation is not safely published until the shared generation fence moves. + throw disconnectError; } } return res.status(status).send(); diff --git a/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js b/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js index 3fb7196266..3a4b365646 100644 --- a/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js +++ b/api/server/controllers/__tests__/UserController.mcpOAuth.spec.js @@ -174,6 +174,38 @@ beforeEach(() => { }); describe('updateUserPluginsController MCP OAuth cleanup', () => { + it('invalidates the shared tool generation even when local disconnect fails', async () => { + const { mcpManager } = setupMCPMocks(); + mcpManager.disconnectUserConnection.mockRejectedValue(new Error('local dispose failed')); + MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null); + + const res = createResponse(); + await updateUserPluginsController(createRequest(), res); + + expect(mockInvalidateCachedTools).toHaveBeenCalledWith({ + userId: 'user-1', + serverName: 'test-server', + }); + expect(mockInvalidateCachedTools.mock.invocationCallOrder[0]).toBeLessThan( + mcpManager.disconnectUserConnection.mock.invocationCallOrder[0], + ); + expect(res.status).toHaveBeenCalledWith(200); + }); + + it('fails the credential update response when the shared generation fence cannot move', async () => { + const { mcpManager } = setupMCPMocks(); + const fenceError = new Error('Redis unavailable'); + mockInvalidateCachedTools.mockRejectedValue(fenceError); + MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null); + + const res = createResponse(); + await updateUserPluginsController(createRequest(), res); + + expect(mcpManager.disconnectUserConnection).toHaveBeenCalledWith('user-1', 'test-server'); + expect(res.status).toHaveBeenCalledWith(500); + expect(logger.error).toHaveBeenCalledWith('[updateUserPluginsController]', fenceError); + }); + it('clears stored OAuth token state when client metadata is missing', async () => { const { flowManager, mcpManager } = setupMCPMocks(); MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null); diff --git a/api/server/controllers/__tests__/deleteUserMcpServers.spec.js b/api/server/controllers/__tests__/deleteUserMcpServers.spec.js index fcb3211f24..195de09538 100644 --- a/api/server/controllers/__tests__/deleteUserMcpServers.spec.js +++ b/api/server/controllers/__tests__/deleteUserMcpServers.spec.js @@ -130,6 +130,42 @@ describe('deleteUserMcpServers', () => { }); }); + test('should delete owned servers when cache invalidation fails', async () => { + const userId = new mongoose.Types.ObjectId(); + const server = await MCPServer.create({ + serverName: 'cache-failure-server', + config: { title: 'Cache Failure Server' }, + author: userId, + }); + + await permissionService.grantPermission({ + principalType: PrincipalType.USER, + principalId: userId, + resourceType: ResourceType.MCPSERVER, + resourceId: server._id, + accessRoleId: AccessRoleIds.MCPSERVER_OWNER, + grantedBy: userId, + }); + + const disconnectUserConnection = jest.fn().mockResolvedValue(undefined); + mockGetMCPManager.mockReturnValue({ disconnectUserConnection }); + mockInvalidateCachedTools.mockRejectedValueOnce(new Error('Redis unavailable')); + + await deleteUserMcpServers(userId.toString()); + + expect(disconnectUserConnection).toHaveBeenCalledWith( + userId.toString(), + 'cache-failure-server', + ); + expect(await MCPServer.findById(server._id)).toBeNull(); + await expect( + AclEntry.countDocuments({ + resourceType: ResourceType.MCPSERVER, + resourceId: server._id, + }), + ).resolves.toBe(0); + }); + test('should preserve multi-owned MCP servers', async () => { const deletingUserId = new mongoose.Types.ObjectId(); const otherOwnerId = new mongoose.Types.ObjectId(); @@ -263,6 +299,10 @@ describe('deleteUserMcpServers', () => { await deleteUserMcpServers(userId.toString()); expect(await MCPServer.findById(server._id)).toBeNull(); + expect(mockInvalidateCachedTools).toHaveBeenCalledWith({ + userId: userId.toString(), + serverName: 'no-manager-server', + }); }); test('should delete legacy MCP servers that have author but no ACL entries', async () => { diff --git a/api/server/controllers/__tests__/mcp.servers.spec.js b/api/server/controllers/__tests__/mcp.servers.spec.js index e8f3f6abda..c46cc07615 100644 --- a/api/server/controllers/__tests__/mcp.servers.spec.js +++ b/api/server/controllers/__tests__/mcp.servers.spec.js @@ -24,11 +24,16 @@ jest.mock('~/server/services/GraphApiService', () => ({ const mockRegistryInstance = { getServerConfig: jest.fn(), + inspectServerUpdate: jest.fn(), + commitServerUpdate: jest.fn(), + updateServer: jest.fn(), + removeServer: jest.fn(), }; +const mockMcpManager = { disconnectUserConnection: jest.fn() }; jest.mock('~/config', () => ({ logger: { debug: jest.fn(), info: jest.fn(), warn: jest.fn(), error: jest.fn() }, - getMCPManager: jest.fn(), + getMCPManager: jest.fn(() => mockMcpManager), getMCPServersRegistry: jest.fn(() => mockRegistryInstance), })); @@ -41,10 +46,17 @@ jest.mock('~/server/services/MCP', () => ({ jest.mock('~/server/services/Config', () => ({ cacheMCPServerTools: jest.fn(), + getMCPToolsCacheGeneration: jest.fn().mockResolvedValue('test-generation'), getMCPServerTools: jest.fn(), + invalidateCachedTools: jest.fn(), })); -const { getMCPServersList, getMCPServerById } = require('~/server/controllers/mcp'); +const { + getMCPServersList, + getMCPServerById, + updateMCPServerController, + deleteMCPServerController, +} = require('~/server/controllers/mcp'); const { grantPermission } = require('~/server/services/PermissionService'); const { seedDefaultRoles } = require('~/models'); @@ -108,6 +120,16 @@ beforeEach(async () => { await User.deleteMany({}); mockResolveAllMcpConfigs.mockReset(); mockRegistryInstance.getServerConfig.mockReset(); + mockRegistryInstance.inspectServerUpdate.mockReset(); + mockRegistryInstance.commitServerUpdate.mockReset(); + mockRegistryInstance.updateServer.mockReset(); + mockRegistryInstance.removeServer.mockReset(); + mockMcpManager.disconnectUserConnection.mockReset().mockResolvedValue(undefined); + const cacheService = require('~/server/services/Config'); + cacheService.invalidateCachedTools.mockReset().mockResolvedValue(undefined); + cacheService.getMCPServerTools.mockReset().mockResolvedValue({ retained: {} }); + cacheService.getMCPToolsCacheGeneration.mockReset().mockResolvedValue('restored-generation'); + cacheService.cacheMCPServerTools.mockReset().mockResolvedValue(undefined); existsSpy = jest.spyOn(SystemGrant, 'exists'); }); @@ -256,3 +278,212 @@ describe('getMCPServerById', () => { expect(payload.oauth.authorization_url).toBeUndefined(); }); }); + +describe('DB-backed server mutation fencing', () => { + const updatedConfig = { + type: 'streamable-http', + url: 'https://updated.example.com/mcp', + source: 'user', + }; + + it('inspects, fences, commits, fences cross-replica creations, and disconnects', async () => { + const user = await createUser(); + mockRegistryInstance.getServerConfig.mockResolvedValue( + createDbConfig(new mongoose.Types.ObjectId()), + ); + mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig); + mockRegistryInstance.commitServerUpdate.mockResolvedValue(updatedConfig); + const res = createRes(); + + await updateMCPServerController( + { user, params: { serverName: 'github' }, body: { config: updatedConfig } }, + res, + ); + + const { invalidateCachedTools } = require('~/server/services/Config'); + expect(invalidateCachedTools).toHaveBeenCalledWith({ userId: user.id, serverName: 'github' }); + expect(invalidateCachedTools).toHaveBeenCalledTimes(2); + expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(user.id, 'github'); + expect(mockRegistryInstance.inspectServerUpdate.mock.invocationCallOrder[0]).toBeLessThan( + invalidateCachedTools.mock.invocationCallOrder[0], + ); + expect(invalidateCachedTools.mock.invocationCallOrder[0]).toBeLessThan( + mockRegistryInstance.commitServerUpdate.mock.invocationCallOrder[0], + ); + expect(mockRegistryInstance.commitServerUpdate.mock.invocationCallOrder[0]).toBeLessThan( + invalidateCachedTools.mock.invocationCallOrder[1], + ); + expect(invalidateCachedTools.mock.invocationCallOrder[1]).toBeLessThan( + mockMcpManager.disconnectUserConnection.mock.invocationCallOrder[0], + ); + expect(res.status).toHaveBeenCalledWith(200); + }); + + it('does not fence the valid catalog when update inspection or persistence fails', async () => { + const user = await createUser(); + const updateError = new Error('inspection failed'); + mockRegistryInstance.getServerConfig.mockResolvedValue( + createDbConfig(new mongoose.Types.ObjectId()), + ); + mockRegistryInstance.inspectServerUpdate.mockRejectedValue(updateError); + const res = createRes(); + + await updateMCPServerController( + { user, params: { serverName: 'github' }, body: { config: updatedConfig } }, + res, + ); + + expect(require('~/server/services/Config').invalidateCachedTools).not.toHaveBeenCalled(); + expect(mockRegistryInstance.commitServerUpdate).not.toHaveBeenCalled(); + expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(500); + }); + + it('does not commit an inspected update when the distributed fence fails', async () => { + const user = await createUser(); + mockRegistryInstance.getServerConfig.mockResolvedValue( + createDbConfig(new mongoose.Types.ObjectId()), + ); + mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig); + require('~/server/services/Config').invalidateCachedTools.mockRejectedValue( + new Error('Redis unavailable'), + ); + const res = createRes(); + + await updateMCPServerController( + { user, params: { serverName: 'github' }, body: { config: updatedConfig } }, + res, + ); + + expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled(); + expect(mockRegistryInstance.commitServerUpdate).not.toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(500); + }); + + it('restores the retained catalog when update persistence fails after fencing', async () => { + const user = await createUser(); + const existingConfig = createDbConfig(new mongoose.Types.ObjectId()); + const retainedTools = { retained: { function: { name: 'retained' } } }; + mockRegistryInstance.getServerConfig.mockResolvedValue(existingConfig); + mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig); + mockRegistryInstance.commitServerUpdate.mockRejectedValue(new Error('database unavailable')); + require('~/server/services/Config').getMCPServerTools.mockResolvedValue(retainedTools); + const res = createRes(); + + await updateMCPServerController( + { user, params: { serverName: 'github' }, body: { config: updatedConfig } }, + res, + ); + + expect(require('~/server/services/Config').cacheMCPServerTools).toHaveBeenCalledWith({ + userId: user.id, + serverName: 'github', + serverConfig: existingConfig, + serverTools: retainedTools, + publicationGeneration: 'restored-generation', + }); + expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(500); + }); + + it('continues an update when only local disconnect cleanup fails', async () => { + const user = await createUser(); + mockRegistryInstance.getServerConfig.mockResolvedValue( + createDbConfig(new mongoose.Types.ObjectId()), + ); + mockMcpManager.disconnectUserConnection.mockRejectedValue(new Error('dispose failed')); + mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig); + mockRegistryInstance.commitServerUpdate.mockResolvedValue(updatedConfig); + const res = createRes(); + + await updateMCPServerController( + { user, params: { serverName: 'github' }, body: { config: updatedConfig } }, + res, + ); + + expect(mockRegistryInstance.commitServerUpdate).toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(200); + }); + + it('retries a transient post-commit fence failure before returning success', async () => { + const user = await createUser(); + mockRegistryInstance.getServerConfig.mockResolvedValue( + createDbConfig(new mongoose.Types.ObjectId()), + ); + mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig); + mockRegistryInstance.commitServerUpdate.mockResolvedValue(updatedConfig); + require('~/server/services/Config') + .invalidateCachedTools.mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error('Redis MOVED')) + .mockResolvedValueOnce(undefined); + const res = createRes(); + + await updateMCPServerController( + { user, params: { serverName: 'github' }, body: { config: updatedConfig } }, + res, + ); + + expect(require('~/server/services/Config').invalidateCachedTools).toHaveBeenCalledTimes(3); + expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(user.id, 'github'); + expect(res.status).toHaveBeenCalledWith(200); + }); + + it('fences before deletion and fences cross-replica creations before disconnecting', async () => { + const user = await createUser(); + mockRegistryInstance.removeServer.mockResolvedValue(undefined); + const res = createRes(); + + await deleteMCPServerController({ user, params: { serverName: 'github' } }, res); + + const { invalidateCachedTools } = require('~/server/services/Config'); + expect(invalidateCachedTools).toHaveBeenCalledWith({ userId: user.id, serverName: 'github' }); + expect(invalidateCachedTools).toHaveBeenCalledTimes(2); + expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(user.id, 'github'); + expect(invalidateCachedTools.mock.invocationCallOrder[0]).toBeLessThan( + mockRegistryInstance.removeServer.mock.invocationCallOrder[0], + ); + expect(mockRegistryInstance.removeServer.mock.invocationCallOrder[0]).toBeLessThan( + invalidateCachedTools.mock.invocationCallOrder[1], + ); + expect(invalidateCachedTools.mock.invocationCallOrder[1]).toBeLessThan( + mockMcpManager.disconnectUserConnection.mock.invocationCallOrder[0], + ); + expect(res.status).toHaveBeenCalledWith(200); + }); + + it('does not delete the registry entry when the distributed fence fails', async () => { + const user = await createUser(); + require('~/server/services/Config').invalidateCachedTools.mockRejectedValue( + new Error('Redis unavailable'), + ); + const res = createRes(); + + await deleteMCPServerController({ user, params: { serverName: 'github' } }, res); + + expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled(); + expect(mockRegistryInstance.removeServer).not.toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(500); + }); + + it('restores the retained catalog when deletion persistence fails after fencing', async () => { + const user = await createUser(); + const existingConfig = createDbConfig(new mongoose.Types.ObjectId()); + const retainedTools = { retained: { function: { name: 'retained' } } }; + mockRegistryInstance.getServerConfig.mockResolvedValue(existingConfig); + mockRegistryInstance.removeServer.mockRejectedValue(new Error('Deletion failed')); + require('~/server/services/Config').getMCPServerTools.mockResolvedValue(retainedTools); + const res = createRes(); + + await deleteMCPServerController({ user, params: { serverName: 'github' } }, res); + + expect(require('~/server/services/Config').cacheMCPServerTools).toHaveBeenCalledWith({ + userId: user.id, + serverName: 'github', + serverConfig: existingConfig, + serverTools: retainedTools, + publicationGeneration: 'restored-generation', + }); + expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(500); + }); +}); diff --git a/api/server/controllers/assistants/v1.js b/api/server/controllers/assistants/v1.js index 460734c01a..926ab7db4d 100644 --- a/api/server/controllers/assistants/v1.js +++ b/api/server/controllers/assistants/v1.js @@ -7,8 +7,7 @@ const validateAuthor = require('~/server/middleware/assistants/validateAuthor'); const { getStrategyFunctions } = require('~/server/services/Files/strategies'); const { deleteAssistantActions } = require('~/server/services/ActionService'); const { getOpenAIClient, fetchAssistants } = require('./helpers'); -const { healMcpToolNames } = require('~/server/services/MCP'); -const { getCachedTools } = require('~/server/services/Config'); +const { healMcpToolNames, getAssistantToolDefinitions } = require('~/server/services/MCP'); const { manifestToolMap, isAgentsOnlyTool } = require('~/app/clients/tools'); /** @@ -31,7 +30,7 @@ const createAssistant = async (req, res) => { delete assistantData.conversation_starters; delete assistantData.append_current_datetime; - const toolDefinitions = (await getCachedTools()) ?? {}; + const toolDefinitions = await getAssistantToolDefinitions({ req, tools }); const healedTools = await healMcpToolNames({ req, tools, toolDefinitions }); assistantData.tools = healedTools @@ -146,7 +145,7 @@ const patchAssistant = async (req, res) => { ...updateData } = req.body; - const toolDefinitions = (await getCachedTools()) ?? {}; + const toolDefinitions = await getAssistantToolDefinitions({ req, tools: updateData.tools }); const healedTools = await healMcpToolNames({ req, tools: updateData.tools, toolDefinitions }); updateData.tools = healedTools diff --git a/api/server/controllers/assistants/v2.js b/api/server/controllers/assistants/v2.js index 221ebc03c1..a436ed611d 100644 --- a/api/server/controllers/assistants/v2.js +++ b/api/server/controllers/assistants/v2.js @@ -2,8 +2,7 @@ const { logger } = require('@librechat/data-schemas'); const { ToolCallTypes } = require('librechat-data-provider'); const validateAuthor = require('~/server/middleware/assistants/validateAuthor'); const { validateAndUpdateTool } = require('~/server/services/ActionService'); -const { healMcpToolNames } = require('~/server/services/MCP'); -const { getCachedTools } = require('~/server/services/Config'); +const { healMcpToolNames, getAssistantToolDefinitions } = require('~/server/services/MCP'); const { manifestToolMap, isAgentsOnlyTool } = require('~/app/clients/tools'); const { updateAssistantDoc } = require('~/models'); const { getOpenAIClient } = require('./helpers'); @@ -29,7 +28,7 @@ const createAssistant = async (req, res) => { delete assistantData.conversation_starters; delete assistantData.append_current_datetime; - const toolDefinitions = (await getCachedTools()) ?? {}; + const toolDefinitions = await getAssistantToolDefinitions({ req, tools }); const healedTools = await healMcpToolNames({ req, tools, toolDefinitions }); assistantData.tools = healedTools @@ -135,7 +134,7 @@ const updateAssistant = async ({ req, openai, assistant_id, updateData }) => { } let hasFileSearch = false; - const toolDefinitions = (await getCachedTools()) ?? {}; + const toolDefinitions = await getAssistantToolDefinitions({ req, tools: updateData.tools }); const healedTools = await healMcpToolNames({ req, tools: updateData.tools, toolDefinitions }); for (const tool of healedTools) { /** Agents-runtime-only tools (e.g. ask_user_question) cannot execute on diff --git a/api/server/controllers/mcp.js b/api/server/controllers/mcp.js index 10eba42768..f150c4ccaf 100644 --- a/api/server/controllers/mcp.js +++ b/api/server/controllers/mcp.js @@ -35,7 +35,12 @@ const { resolveMcpConfigNames, resolveAllMcpConfigs, } = require('~/server/services/MCP'); -const { cacheMCPServerTools, getMCPServerTools } = require('~/server/services/Config'); +const { + cacheMCPServerTools, + getMCPServerTools, + getMCPToolsCacheGeneration, + invalidateCachedTools, +} = require('~/server/services/Config'); const { getResourcePermissionsMap } = require('~/server/services/PermissionService'); const { hasCapability } = require('~/server/middleware/roles/capabilities'); const { getMCPManager, getMCPServersRegistry } = require('~/config'); @@ -94,6 +99,68 @@ function handleMCPError(error, res) { return null; } +/** Disposes a stale local connection after its DB-backed config has changed. */ +async function disconnectLocalMCPServer(userId, serverName) { + try { + await getMCPManager()?.disconnectUserConnection(userId, serverName); + } catch (error) { + logger.warn( + `[MCP Cache] Failed to disconnect the local connection for ${serverName} (user: ${userId}):`, + error, + ); + } +} + +const POST_COMMIT_FENCE_RETRY_DELAYS_MS = [0, 50, 200]; + +/** Retries the shared fence after persistence; config-bound connections remain a durable + * fallback if Redis stays unavailable, so an old connection cannot serve the new config. */ +async function fenceCommittedMCPMutation({ userId, serverName }) { + let lastError; + for (const delay of POST_COMMIT_FENCE_RETRY_DELAYS_MS) { + if (delay > 0) { + await new Promise((resolve) => setTimeout(resolve, delay)); + } + try { + await invalidateCachedTools({ userId, serverName }); + return; + } catch (error) { + lastError = error; + logger.warn( + `[MCP Cache] Failed to fence committed mutation for ${serverName} (user: ${userId}); retrying:`, + error, + ); + } + } + throw lastError; +} + +/** + * Republishes the pre-mutation catalog under the new fence when persistence + * fails. The retained connection will reacquire that generation on its next + * use; this snapshot keeps every replica authoritative in the meantime. + */ +async function restoreRetainedServerCatalog({ userId, serverName, serverConfig, serverTools }) { + if (serverTools == null) { + return; + } + try { + const publicationGeneration = await getMCPToolsCacheGeneration({ userId, serverName }); + await cacheMCPServerTools({ + userId, + serverName, + serverConfig, + serverTools, + publicationGeneration, + }); + } catch (error) { + logger.error( + `[MCP Cache] Failed to restore the retained catalog for ${serverName} (user: ${userId}):`, + error, + ); + } +} + /** * Get all MCP tools available to the user. */ @@ -150,8 +217,14 @@ const getMCPTools = async (req, res) => { } let serverTools; + let publicationGeneration; try { - serverTools = await mcpManager.getServerToolFunctions(userId, serverName); + ({ tools: serverTools, publicationGeneration } = + await mcpManager.getServerToolFunctionsSnapshot( + userId, + serverName, + mcpConfig[serverName], + )); } catch (error) { logger.error(`[getMCPTools] Error fetching tools for server ${serverName}:`, error); continue; @@ -162,17 +235,16 @@ const getMCPTools = async (req, res) => { } serverToolsMap.set(serverName, serverTools); - if (Object.keys(serverTools).length > 0) { - // Cache asynchronously without blocking - cacheMCPServerTools({ - userId, - serverName, - serverTools, - serverConfig: mcpConfig[serverName], - }).catch((err) => - logger.error(`[getMCPTools] Failed to cache tools for ${serverName}:`, err), - ); - } + // Empty is an authoritative catalog too; re-cache it after TTL expiry to avoid polling. + cacheMCPServerTools({ + userId, + serverName, + serverTools, + serverConfig: mcpConfig[serverName], + publicationGeneration, + }).catch((err) => + logger.error(`[getMCPTools] Failed to cache tools for ${serverName}:`, err), + ); } // Process each configured server @@ -509,12 +581,30 @@ const updateMCPServerController = async (req, res) => { .json({ message: 'Forbidden: Insufficient permissions to configure OBO' }); } - const parsedConfig = await getMCPServersRegistry().updateServer( + const registry = getMCPServersRegistry(); + const parsedConfig = await registry.inspectServerUpdate( serverName, validation.data, 'DB', userId, ); + const retainedTools = await getMCPServerTools(userId, serverName, existingConfig); + await invalidateCachedTools({ userId, serverName }); + try { + await registry.commitServerUpdate(serverName, parsedConfig, 'DB', userId); + } catch (error) { + await restoreRetainedServerCatalog({ + userId, + serverName, + serverConfig: existingConfig, + serverTools: retainedTools, + }); + throw error; + } + /** Fence connections another replica could have created from the old DB + * config between the pre-commit fence and the committed update. */ + await fenceCommittedMCPMutation({ userId, serverName }); + await disconnectLocalMCPServer(userId, serverName); res.status(200).json(redactServerSecrets(parsedConfig, { canEdit: true })); } catch (error) { @@ -535,7 +625,24 @@ const deleteMCPServerController = async (req, res) => { try { const userId = req.user?.id; const { serverName } = req.params; - await getMCPServersRegistry().removeServer(serverName, 'DB', userId); + const registry = getMCPServersRegistry(); + const existingConfig = await registry.getServerConfig(serverName, userId); + const retainedTools = await getMCPServerTools(userId, serverName, existingConfig); + await invalidateCachedTools({ userId, serverName }); + try { + await registry.removeServer(serverName, 'DB', userId); + } catch (error) { + await restoreRetainedServerCatalog({ + userId, + serverName, + serverConfig: existingConfig, + serverTools: retainedTools, + }); + throw error; + } + /** Fence connections another replica could have created before deletion committed. */ + await fenceCommittedMCPMutation({ userId, serverName }); + await disconnectLocalMCPServer(userId, serverName); res.status(200).json({ message: 'MCP server deleted successfully' }); } catch (error) { logger.error('[deleteMCPServer]', error); diff --git a/api/server/index.metrics.spec.js b/api/server/index.metrics.spec.js index c907aca3b6..83ca7ba830 100644 --- a/api/server/index.metrics.spec.js +++ b/api/server/index.metrics.spec.js @@ -15,6 +15,7 @@ jest.mock('~/server/services/Config', () => ({ fileStrategy: 'local', imageOutputType: 'PNG', }), + mergeAppTools: jest.fn().mockResolvedValue(undefined), setCachedTools: jest.fn(), })); diff --git a/api/server/index.spec.js b/api/server/index.spec.js index f13faa54c0..c4c1486b9e 100644 --- a/api/server/index.spec.js +++ b/api/server/index.spec.js @@ -5,6 +5,7 @@ const { MongoMemoryServer } = require('mongodb-memory-server'); const mongoose = require('mongoose'); jest.mock('~/server/services/Config', () => ({ + mergeAppTools: jest.fn().mockResolvedValue(undefined), loadCustomConfig: jest.fn(() => Promise.resolve({})), getAppConfig: jest.fn().mockResolvedValue({ paths: { diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 83dfb7874b..ed919ba703 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -20,6 +20,8 @@ const mockRegistryInstance = { getAllServerConfigs: jest.fn(), ensureConfigServers: jest.fn().mockResolvedValue({}), addServer: jest.fn(), + inspectServerUpdate: jest.fn(), + commitServerUpdate: jest.fn(), updateServer: jest.fn(), removeServer: jest.fn(), getAllowedDomains: jest.fn().mockReturnValue(null), @@ -128,6 +130,9 @@ jest.mock('~/models', () => ({ jest.mock('~/server/services/Config', () => ({ setCachedTools: jest.fn(), getCachedTools: jest.fn(), + cacheMCPServerTools: jest.fn(), + getMCPToolsCacheGeneration: jest.fn().mockResolvedValue('test-generation'), + invalidateCachedTools: jest.fn(), getMCPServerTools: jest.fn(), loadCustomConfig: jest.fn(), getAppConfig: jest.fn().mockResolvedValue({ mcpConfig: {} }), @@ -272,8 +277,15 @@ describe('MCP Routes', () => { */ mockRegistryInstance.getServerConfig.mockReset().mockResolvedValue(undefined); mockRegistryInstance.addServer.mockReset(); + mockRegistryInstance.inspectServerUpdate.mockReset(); + mockRegistryInstance.commitServerUpdate.mockReset(); mockRegistryInstance.updateServer.mockReset(); mockRegistryInstance.removeServer.mockReset(); + const cacheService = require('~/server/services/Config'); + cacheService.getMCPServerTools.mockReset().mockResolvedValue(null); + cacheService.getMCPToolsCacheGeneration.mockReset().mockResolvedValue('test-generation'); + cacheService.cacheMCPServerTools.mockReset().mockResolvedValue(undefined); + cacheService.invalidateCachedTools.mockReset().mockResolvedValue(undefined); }); describe('GET /:serverName/oauth/initiate', () => { @@ -843,7 +855,7 @@ describe('MCP Routes', () => { const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -905,7 +917,7 @@ describe('MCP Routes', () => { const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -958,9 +970,12 @@ describe('MCP Routes', () => { mockRegistryInstance.getServerConfig.mockResolvedValue({}); mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'test-server': mergedServerConfig }); + const fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: fetchedTools, complete: true }); const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue(fetchedTools), + fetchOrderedToolsSnapshot, }), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -976,6 +991,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(302); expect(mockResolveAllMcpConfigs).toHaveBeenCalledWith('test-user-id'); + expect(fetchOrderedToolsSnapshot).toHaveBeenCalledTimes(1); expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith( expect.objectContaining({ serverConfig: mergedServerConfig }), ); @@ -1030,7 +1046,9 @@ describe('MCP Routes', () => { const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue(fetchedTools), + fetchToolsSnapshot: jest + .fn() + .mockResolvedValue({ tools: fetchedTools, complete: true }), }), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -1093,7 +1111,9 @@ describe('MCP Routes', () => { const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue(fetchedTools), + fetchToolsSnapshot: jest + .fn() + .mockResolvedValue({ tools: fetchedTools, complete: true }), }), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -1209,13 +1229,16 @@ describe('MCP Routes', () => { require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); const mockUserConnection = { - fetchTools: jest.fn().mockResolvedValue([ - { - name: 'test-tool', - description: 'A test tool', - inputSchema: { type: 'object' }, - }, - ]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ + tools: [ + { + name: 'test-tool', + description: 'A test tool', + inputSchema: { type: 'object' }, + }, + ], + complete: true, + }), }; const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue(mockUserConnection), @@ -1310,7 +1333,7 @@ describe('MCP Routes', () => { }); require('~/config').getMCPManager.mockReturnValue({ getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); @@ -1385,7 +1408,7 @@ describe('MCP Routes', () => { }); require('~/config').getMCPManager.mockReturnValue({ getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); @@ -1443,7 +1466,7 @@ describe('MCP Routes', () => { }); require('~/config').getMCPManager.mockReturnValue({ getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); @@ -1497,7 +1520,7 @@ describe('MCP Routes', () => { }); require('~/config').getMCPManager.mockReturnValue({ getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); @@ -1730,7 +1753,7 @@ describe('MCP Routes', () => { require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); const mockUserConnection = { - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }; const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue(mockUserConnection), @@ -2389,10 +2412,13 @@ describe('MCP Routes', () => { it('should successfully reinitialize server and cache tools', async () => { const mockUserConnection = { - fetchTools: jest.fn().mockResolvedValue([ - { name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } }, - { name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } }, - ]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ + tools: [ + { name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } }, + { name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } }, + ], + complete: true, + }), }; const mockMcpManager = { @@ -2435,11 +2461,18 @@ describe('MCP Routes', () => { 'test-user-id', 'test-server', ); + expect(require('~/server/services/Config').invalidateCachedTools).toHaveBeenCalledWith({ + userId: 'test-user-id', + serverName: 'test-server', + }); + expect( + require('~/server/services/Config').invalidateCachedTools.mock.invocationCallOrder[0], + ).toBeLessThan(mockMcpManager.disconnectUserConnection.mock.invocationCallOrder[0]); }); it('should handle server with custom user variables', async () => { const mockUserConnection = { - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }; const mockMcpManager = { @@ -2896,7 +2929,7 @@ describe('MCP Routes', () => { const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest.fn().mockResolvedValue([]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -2949,10 +2982,12 @@ describe('MCP Routes', () => { const mockMcpManager = { getUserConnection: jest.fn().mockResolvedValue({ - fetchTools: jest - .fn() - .mockResolvedValue([{ name: 'test-tool', description: 'Test tool' }]), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ + tools: [{ name: 'test-tool', description: 'Test tool' }], + complete: true, + }), }), + getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -2967,6 +3002,60 @@ describe('MCP Routes', () => { const basePath = getBasePath(); expect(response.headers.location).toContain(`${basePath}/oauth/success`); + expect(require('~/server/services/Config/mcp').updateMCPServerTools).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'test-user-id', + serverName: 'test-server', + tools: [{ name: 'test-tool', description: 'Test tool' }], + publicationGeneration: 'oauth-connection-generation', + }), + ); + }); + + it('preserves cached tools when the post-OAuth snapshot is incomplete', async () => { + const { logger } = require('@librechat/data-schemas'); + const { MCPOAuthHandler, MCPTokenStorage } = require('@librechat/api'); + const mockTokens = { + access_token: 'edge-access-token', + refresh_token: 'edge-refresh-token', + }; + const mockFlowManager = { + getFlowState: jest.fn(), + completeFlow: jest.fn(), + }; + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + MCPOAuthHandler.getFlowState.mockResolvedValue({ + state: 'test-user-id:test-server', + serverName: 'test-server', + userId: 'test-user-id', + metadata: { serverUrl: 'https://example.com', oauth: {} }, + clientInfo: {}, + codeVerifier: 'test-verifier', + }); + 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'), + }); + + const flowId = 'test-user-id:test-server'; + const csrfToken = generateTestCsrfToken(flowId); + await request(app) + .get(`/api/mcp/test-server/oauth/callback?code=test-code&state=${flowId}`) + .set('Cookie', [`oauth_csrf=${csrfToken}`]) + .expect(302); + + expect(require('~/server/services/Config/mcp').updateMCPServerTools).not.toHaveBeenCalled(); + expect(logger.warn).toHaveBeenCalledWith( + '[MCP OAuth] Preserving cached tools for test-server because tools/list returned an incomplete snapshot', + ); }); }); @@ -3064,6 +3153,68 @@ describe('MCP Routes', () => { expect(mockResolveAllMcpConfigs).not.toHaveBeenCalled(); }); + it('caches a live user snapshot with its connection-bound publication generation', async () => { + const { Constants } = require('librechat-data-provider'); + const { cacheMCPServerTools, getMCPServerTools } = require('~/server/services/Config'); + const pluginKey = `search${Constants.mcp_delimiter}user-server`; + const serverTools = { + [pluginKey]: { + type: 'function', + function: { + name: pluginKey, + description: 'Search', + parameters: { type: 'object' }, + }, + }, + }; + const serverConfig = { type: 'sse', url: 'https://user.example.com/sse' }; + mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'user-server': serverConfig }); + getMCPServerTools.mockResolvedValueOnce(null); + cacheMCPServerTools.mockResolvedValueOnce(); + require('~/config').getMCPManager.mockReturnValue({ + getServerToolFunctionsSnapshot: jest.fn().mockResolvedValue({ + tools: serverTools, + publicationGeneration: 'connection-generation', + }), + }); + + const response = await request(app).get('/api/mcp/tools'); + + expect(response.status).toBe(200); + expect(cacheMCPServerTools).toHaveBeenCalledWith({ + userId: 'test-user-id', + serverName: 'user-server', + serverTools, + serverConfig, + publicationGeneration: 'connection-generation', + }); + }); + + it('re-caches an authoritative empty live snapshot after a cache miss', async () => { + const { cacheMCPServerTools, getMCPServerTools } = require('~/server/services/Config'); + const serverConfig = { type: 'sse', url: 'https://empty.example.com/sse' }; + mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'empty-server': serverConfig }); + getMCPServerTools.mockResolvedValueOnce(null); + cacheMCPServerTools.mockResolvedValueOnce(); + const getServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({ + tools: {}, + publicationGeneration: undefined, + }); + require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot }); + + const response = await request(app).get('/api/mcp/tools'); + + expect(response.status).toBe(200); + expect(response.body.servers['empty-server'].tools).toEqual([]); + expect(cacheMCPServerTools).toHaveBeenCalledWith({ + userId: 'test-user-id', + serverName: 'empty-server', + serverTools: {}, + serverConfig, + publicationGeneration: undefined, + }); + }); + it('should continue returning MCP tools when one server cache lookup fails', async () => { const { Constants } = require('librechat-data-provider'); const { logger } = require('@librechat/data-schemas'); @@ -3095,9 +3246,12 @@ describe('MCP Routes', () => { }, }); - const mockGetServerToolFunctions = jest.fn().mockResolvedValue(null); + const mockGetServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({ + tools: null, + publicationGeneration: 'test-generation', + }); require('~/config').getMCPManager.mockReturnValue({ - getServerToolFunctions: mockGetServerToolFunctions, + getServerToolFunctionsSnapshot: mockGetServerToolFunctionsSnapshot, }); const response = await request(app).get('/api/mcp/tools'); @@ -3107,7 +3261,14 @@ describe('MCP Routes', () => { '[getMCPTools] Error fetching cached tools for bad-server:', expect.any(Error), ); - expect(mockGetServerToolFunctions).toHaveBeenCalledWith('test-user-id', 'bad-server'); + expect(mockGetServerToolFunctionsSnapshot).toHaveBeenCalledWith( + 'test-user-id', + 'bad-server', + { + type: 'sse', + url: 'https://bad.example.com/sse', + }, + ); expect(response.body.servers['good-server']).toMatchObject({ name: 'good-server', icon: '/icons/good.svg', @@ -3142,9 +3303,12 @@ describe('MCP Routes', () => { getMCPServerTools.mockRejectedValue(new Error('cache unavailable')); - const mockGetServerToolFunctions = jest.fn().mockResolvedValue(null); + const mockGetServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({ + tools: null, + publicationGeneration: 'test-generation', + }); require('~/config').getMCPManager.mockReturnValue({ - getServerToolFunctions: mockGetServerToolFunctions, + getServerToolFunctionsSnapshot: mockGetServerToolFunctionsSnapshot, }); const response = await request(app).get('/api/mcp/tools'); @@ -3159,7 +3323,7 @@ describe('MCP Routes', () => { tools: [], }); expect(logger.error).toHaveBeenCalledTimes(2); - expect(mockGetServerToolFunctions).toHaveBeenCalledTimes(2); + expect(mockGetServerToolFunctionsSnapshot).toHaveBeenCalledTimes(2); }); }); @@ -3567,7 +3731,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(403); expect(response.body.message).toMatch(/Insufficient permissions to configure OBO/); - expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled(); + expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled(); }); it('allows PATCH without CONFIGURE_OBO when OBO is unchanged', async () => { @@ -3589,10 +3753,11 @@ describe('MCP Routes', () => { ...oboConfig, obo: { scopes: 'api://mcp-server-id/Mcp.Tools.ReadWrite' }, }); - mockRegistryInstance.updateServer.mockResolvedValue({ + mockRegistryInstance.inspectServerUpdate.mockResolvedValue({ ...oboConfig, title: 'Renamed OBO Server', }); + mockRegistryInstance.commitServerUpdate.mockResolvedValue(undefined); const response = await request(app) .patch('/api/mcp/servers/obo-server') @@ -3605,7 +3770,7 @@ describe('MCP Routes', () => { }); expect(response.status).toBe(200); - expect(mockRegistryInstance.updateServer).toHaveBeenCalled(); + expect(mockRegistryInstance.commitServerUpdate).toHaveBeenCalled(); }); it('rejects PATCH that removes OBO from an existing OBO server without CONFIGURE_OBO', async () => { @@ -3639,7 +3804,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(403); expect(response.body.message).toMatch(/Insufficient permissions to configure OBO/); - expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled(); + expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled(); }); it('rejects PATCH that redirects the URL of an existing OBO server without CONFIGURE_OBO', async () => { @@ -3676,7 +3841,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(403); expect(response.body.message).toMatch(/Insufficient permissions to configure OBO/); - expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled(); + expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled(); }); }); @@ -3772,7 +3937,11 @@ describe('MCP Routes', () => { description: 'Updated description', }; - mockRegistryInstance.updateServer.mockResolvedValue({ ...updatedConfig, source: 'user' }); + mockRegistryInstance.inspectServerUpdate.mockResolvedValue({ + ...updatedConfig, + source: 'user', + }); + mockRegistryInstance.commitServerUpdate.mockResolvedValue(undefined); const response = await request(app) .patch('/api/mcp/servers/test-server') @@ -3782,7 +3951,7 @@ describe('MCP Routes', () => { expect(response.body.type).toBe('sse'); expect(response.body.url).toBe('https://updated-mcp-server.example.com/sse'); expect(response.body.title).toBe('Updated Server'); - expect(mockRegistryInstance.updateServer).toHaveBeenCalledWith( + expect(mockRegistryInstance.inspectServerUpdate).toHaveBeenCalledWith( 'test-server', expect.objectContaining({ type: 'sse', @@ -3800,13 +3969,14 @@ describe('MCP Routes', () => { title: 'Updated Server', }; - mockRegistryInstance.updateServer.mockResolvedValue({ + mockRegistryInstance.inspectServerUpdate.mockResolvedValue({ ...validConfig, apiKey: { source: 'admin', authorization_type: 'bearer', key: 'preserved-admin-key' }, oauth: { client_id: 'cid', client_secret: 'preserved-oauth-secret' }, headers: { Authorization: 'Bearer internal-token' }, env: { DATABASE_URL: 'postgres://admin:pass@localhost/db' }, }); + mockRegistryInstance.commitServerUpdate.mockResolvedValue(undefined); const response = await request(app) .patch('/api/mcp/servers/test-server') @@ -3832,7 +4002,7 @@ describe('MCP Routes', () => { statusCode: 400, }, ); - mockRegistryInstance.updateServer.mockRejectedValue(error); + mockRegistryInstance.inspectServerUpdate.mockRejectedValue(error); const response = await request(app) .patch('/api/mcp/servers/test-server') @@ -3884,7 +4054,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(400); expect(response.body.message).toBe('Invalid configuration'); - expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled(); + expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled(); }); it('should reject streamable-http URL containing env variable references', async () => { @@ -3899,7 +4069,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(400); expect(response.body.message).toBe('Invalid configuration'); - expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled(); + expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled(); }); it('should reject websocket URL containing env variable references', async () => { @@ -3914,7 +4084,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(400); expect(response.body.message).toBe('Invalid configuration'); - expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled(); + expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled(); }); it('should return 500 when registry throws error', async () => { @@ -3924,7 +4094,7 @@ describe('MCP Routes', () => { title: 'Test Server', }; - mockRegistryInstance.updateServer.mockRejectedValue(new Error('Update failed')); + mockRegistryInstance.inspectServerUpdate.mockRejectedValue(new Error('Update failed')); const response = await request(app) .patch('/api/mcp/servers/test-server') diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index e6a5eecf73..bb4be24618 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -46,6 +46,7 @@ const { } = require('~/server/services/MCP'); const { requireJwtAuth, canAccessMCPServerResource } = require('~/server/middleware'); const { getUserPluginAuthValue } = require('~/server/services/PluginService'); +const { invalidateCachedTools } = require('~/server/services/Config'); const { updateMCPServerTools } = require('~/server/services/Config/mcp'); const { reinitMCPServer } = require('~/server/services/Tools/mcp'); const { getLogStores } = require('~/cache'); @@ -542,13 +543,24 @@ router.get('/:serverName/oauth/callback', async (req, res) => { const oauthReconnectionManager = getOAuthReconnectionManager(); oauthReconnectionManager.clearReconnection(flowState.userId, serverName); - const tools = await userConnection.fetchTools(); - await updateMCPServerTools({ - userId: flowState.userId, - serverName, - tools, - serverConfig, - }); + 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, + tools: snapshot.tools, + serverConfig, + publicationGeneration, + }); + } else { + logger.warn( + `[MCP OAuth] Preserving cached tools for ${serverName} because tools/list returned an incomplete snapshot`, + ); + } } else { logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`); } @@ -791,7 +803,11 @@ router.post( }); } - await mcpManager.disconnectUserConnection(user.id, serverName); + try { + await invalidateCachedTools({ userId: user.id, serverName }); + } finally { + await mcpManager.disconnectUserConnection(user.id, serverName); + } logger.info( `[MCP Reinitialize] Disconnected existing user connection for server: ${serverName}`, ); diff --git a/api/server/services/Config/__tests__/getCachedTools.lock.spec.js b/api/server/services/Config/__tests__/getCachedTools.lock.spec.js new file mode 100644 index 0000000000..9417bbfcba --- /dev/null +++ b/api/server/services/Config/__tests__/getCachedTools.lock.spec.js @@ -0,0 +1,341 @@ +const { CacheKeys } = require('librechat-data-provider'); +const calculateSlot = require('cluster-key-slot'); + +const mockRedisClient = { + set: jest.fn(), + eval: jest.fn(), +}; +const mockKeyvRedisClient = { + eval: jest.fn(), +}; + +jest.mock('@librechat/api', () => ({ + ...jest.requireActual('@librechat/api'), + cacheConfig: { FORCED_IN_MEMORY_CACHE_NAMESPACES: [] }, + mcpConfig: { USER_CONNECTION_IDLE_TIMEOUT: 15 * 60 * 1000 }, + ioredisClient: mockRedisClient, + keyvRedisClient: mockKeyvRedisClient, +})); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { warn: jest.fn() }, +})); + +jest.mock('~/cache/getLogStores', () => jest.fn()); + +const getLogStores = require('~/cache/getLogStores'); +const mockCache = { get: jest.fn(), set: jest.fn(), delete: jest.fn() }; +getLogStores.mockReturnValue(mockCache); + +const { + getCachedTools, + updateCachedGlobalTools, + getMCPToolsCacheGeneration, + setCachedTools, + setCachedToolsIfCurrent, + runWithGlobalCacheLock, + invalidateCachedTools, + setCachedToolsWithinGlobalLock, + getNextAppToolsPublicationRevision, + setCachedAppServerTools, +} = require('../getCachedTools'); + +describe('global tool cache write lock', () => { + afterEach(() => { + jest.useRealTimers(); + }); + + beforeEach(() => { + jest.clearAllMocks(); + mockRedisClient.set.mockResolvedValue('OK'); + mockRedisClient.eval.mockResolvedValue(1); + mockKeyvRedisClient.eval.mockResolvedValue(1); + mockCache.set.mockResolvedValue(true); + mockCache.delete.mockResolvedValue(true); + }); + + it('acquires and safely releases the Redis lock around an aggregate update', async () => { + const operation = jest.fn().mockResolvedValue('updated'); + + await expect(runWithGlobalCacheLock(operation)).resolves.toBe('updated'); + + expect(mockRedisClient.set).toHaveBeenCalledWith( + `${CacheKeys.TOOL_CACHE}:tools:global:write-lock`, + expect.any(String), + 'PX', + 30_000, + 'NX', + ); + const token = mockRedisClient.set.mock.calls[0][1]; + expect(mockRedisClient.eval).toHaveBeenCalledWith( + expect.stringContaining("redis.call('GET'"), + 1, + `${CacheKeys.TOOL_CACHE}:tools:global:write-lock`, + token, + ); + const fenceKey = mockKeyvRedisClient.eval.mock.calls[0][1].keys[0]; + expect(calculateSlot(fenceKey)).toBe(calculateSlot(`${CacheKeys.TOOL_CACHE}:tools:global`)); + }); + + it('releases the Redis lock when the aggregate update fails', async () => { + const operation = jest.fn().mockRejectedValue(new Error('cache read failed')); + + await expect(runWithGlobalCacheLock(operation)).rejects.toThrow('cache read failed'); + + expect(mockRedisClient.eval).toHaveBeenCalledTimes(1); + }); + + it('serializes direct global writes and invalidation', async () => { + await setCachedTools({ builtin: {} }); + await invalidateCachedTools({ invalidateGlobal: true }); + + expect(mockRedisClient.set).toHaveBeenCalledTimes(2); + expect(mockRedisClient.eval).toHaveBeenCalledTimes(2); + }); + + it('does not reacquire the lock for a write already inside an aggregate update', async () => { + await runWithGlobalCacheLock(() => setCachedToolsWithinGlobalLock({ mcp: {} })); + + expect(mockRedisClient.set).toHaveBeenCalledTimes(1); + expect(mockRedisClient.eval).toHaveBeenCalledTimes(1); + expect(mockKeyvRedisClient.eval).toHaveBeenCalledWith( + expect.stringContaining("redis.call('GET', KEYS[1])"), + expect.objectContaining({ + keys: [ + `tools:global:write-fence:{${CacheKeys.TOOL_CACHE}:tools:global}`, + `${CacheKeys.TOOL_CACHE}:tools:global`, + ], + }), + ); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('atomically replaces a legacy global catalog while holding its Redis fence', async () => { + mockCache.get.mockResolvedValue({ old_mcp_server: {}, builtin: {} }); + + await updateCachedGlobalTools(() => ({ builtin: {} })); + + expect(mockRedisClient.set).toHaveBeenCalledTimes(1); + expect(mockKeyvRedisClient.eval).toHaveBeenCalledTimes(3); + expect(mockKeyvRedisClient.eval.mock.calls).toEqual( + expect.arrayContaining([ + [ + expect.stringContaining("redis.call('PSETEX', KEYS[2]"), + expect.objectContaining({ + keys: [ + `tools:global:write-fence:{${CacheKeys.TOOL_CACHE}:tools:global}`, + `${CacheKeys.TOOL_CACHE}:tools:global`, + ], + arguments: [ + expect.any(String), + JSON.stringify({ builtin: {} }), + expect.any(String), + expect.any(String), + ], + }), + ], + ]), + ); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('rejects a Redis-backed global write made without distributed lock ownership', async () => { + await expect(setCachedToolsWithinGlobalLock({ unsafe: {} })).rejects.toThrow( + 'Global tool cache write requires lock ownership', + ); + + expect(mockCache.set).not.toHaveBeenCalled(); + expect(mockKeyvRedisClient.eval).not.toHaveBeenCalled(); + }); + + it('rejects a delayed global write after its distributed lease is lost', async () => { + mockKeyvRedisClient.eval.mockResolvedValueOnce(1).mockResolvedValueOnce(0).mockResolvedValue(1); + + await expect( + runWithGlobalCacheLock(() => setCachedToolsWithinGlobalLock({ stale: {} })), + ).rejects.toThrow('Global tool cache lock ownership was lost before write'); + + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('rejects a delayed ownership claim after a newer owner has fenced the slot', async () => { + const operation = jest.fn(); + mockKeyvRedisClient.eval.mockResolvedValueOnce(0); + + await expect(runWithGlobalCacheLock(operation)).rejects.toThrow( + 'Tool cache lock expired or was superseded before ownership could be fenced', + ); + + expect(operation).not.toHaveBeenCalled(); + const [claimScript, claimOptions] = mockKeyvRedisClient.eval.mock.calls[0]; + expect(claimScript).toContain('current ~= ARGV[1]'); + expect(claimScript).toContain("redis.call('TIME')"); + expect(claimOptions.arguments).toEqual([ + mockRedisClient.set.mock.calls[0][1], + expect.any(String), + '1000', + ]); + }); + + it('atomically checks the generation and writes a generation-guarded user catalog', async () => { + await expect( + setCachedToolsIfCurrent( + { current: {} }, + { + userId: 'user-1', + serverName: 'server-1', + configGeneration: 'config-current', + publicationGeneration: 'generation-current', + }, + ), + ).resolves.toBe(true); + + expect(mockRedisClient.set).toHaveBeenCalledWith( + `${CacheKeys.TOOL_CACHE}:tools:mcp-write-lock:user-1:server-1`, + expect.any(String), + 'PX', + 30_000, + 'NX', + ); + expect(mockRedisClient.eval).toHaveBeenCalledTimes(1); + expect(mockKeyvRedisClient.eval).toHaveBeenCalledWith( + expect.stringContaining("redis.call('PSETEX', KEYS[2]"), + expect.objectContaining({ + keys: [ + `${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-generation:{user-1:server-1}`, + `${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:config-current`, + ], + arguments: [ + 'generation-current', + expect.any(String), + expect.any(String), + expect.stringContaining('"publicationGeneration":"generation-current"'), + expect.any(String), + expect.any(String), + ], + }), + ); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('does not write tools when the atomic generation check observes a replacement', async () => { + mockKeyvRedisClient.eval.mockResolvedValue(0); + + await expect( + setCachedToolsIfCurrent( + { stale: {} }, + { + userId: 'user-1', + serverName: 'server-1', + configGeneration: 'config-old', + publicationGeneration: 'generation-old', + }, + ), + ).resolves.toBe(false); + + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('orders app snapshots atomically in the app catalog Redis slot', async () => { + mockCache.get.mockResolvedValue(null); + mockKeyvRedisClient.eval + .mockResolvedValueOnce(1) + .mockResolvedValueOnce(2) + .mockResolvedValueOnce(1) + .mockResolvedValueOnce(0); + + const older = await getNextAppToolsPublicationRevision('server-1', 'config-current'); + const newer = await getNextAppToolsPublicationRevision('server-1', 'config-current'); + await expect( + setCachedAppServerTools('server-1', 'config-current', { current: {} }, newer), + ).resolves.toBe(true); + await expect( + setCachedAppServerTools('server-1', 'config-current', { stale: {} }, older), + ).resolves.toBe(false); + + const [reserveScript, reserveOptions] = mockKeyvRedisClient.eval.mock.calls[0]; + const [writeScript, writeOptions] = mockKeyvRedisClient.eval.mock.calls[2]; + expect(reserveScript).toContain("redis.call('INCR', KEYS[1])"); + expect(writeScript).toContain('tonumber(current) > tonumber(ARGV[1])'); + expect(writeScript).toContain("currentEntry['value']['publicationRevision']"); + expect(calculateSlot(reserveOptions.keys[0])).toBe(calculateSlot(writeOptions.keys[1])); + expect(calculateSlot(writeOptions.keys[0])).toBe(calculateSlot(writeOptions.keys[1])); + expect(writeOptions.keys[0]).toContain('app-committed-revision'); + expect(writeOptions.keys[0]).not.toBe(reserveOptions.keys[0]); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('serializes legacy user catalog migration with Redis-backed writers', async () => { + const legacy = { legacy: {} }; + mockCache.get + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(legacy) + .mockResolvedValueOnce('generation-current') + .mockResolvedValueOnce('generation-current'); + + await expect( + getCachedTools({ + userId: 'user-1', + serverName: 'server-1', + configGeneration: 'config-current', + }), + ).resolves.toBe(legacy); + + expect(mockRedisClient.set).toHaveBeenCalledWith( + `${CacheKeys.TOOL_CACHE}:tools:mcp-write-lock:user-1:server-1`, + expect.any(String), + 'PX', + 30_000, + 'NX', + ); + expect(mockKeyvRedisClient.eval).toHaveBeenCalledWith( + expect.stringContaining("redis.call('EXISTS', KEYS[2])"), + expect.objectContaining({ + keys: [ + `tools:mcp:write-fence:{user-1:server-1}`, + `${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-legacy-fence:{user-1:server-1}`, + `${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-generation:{user-1:server-1}`, + `${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:config-current`, + ], + }), + ); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('rejects first-generation creation after its distributed lock ownership is lost', async () => { + mockCache.get.mockResolvedValue(null); + mockKeyvRedisClient.eval + .mockResolvedValueOnce(1) + .mockResolvedValueOnce(-1) + .mockResolvedValue(1); + + await expect( + getMCPToolsCacheGeneration({ userId: 'user-1', serverName: 'server-1' }), + ).rejects.toThrow('Tool cache lock ownership was lost before generation creation'); + + const createCall = mockKeyvRedisClient.eval.mock.calls.find(([script]) => + script.includes("redis.call('EXISTS', KEYS[2])"), + ); + expect(calculateSlot(createCall[1].keys[0])).toBe(calculateSlot(createCall[1].keys[1])); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('waits through the full abandoned Redis lease before giving up', async () => { + jest.useFakeTimers(); + const startedAt = Date.now(); + mockRedisClient.set.mockImplementation(async () => + Date.now() - startedAt >= 30_000 ? 'OK' : null, + ); + const operation = jest.fn().mockResolvedValue('recovered'); + + const result = runWithGlobalCacheLock(operation); + await jest.advanceTimersByTimeAsync(5_000); + expect(operation).not.toHaveBeenCalled(); + + await jest.advanceTimersByTimeAsync(25_100); + await expect(result).resolves.toBe('recovered'); + expect(operation).toHaveBeenCalledTimes(1); + }); +}); diff --git a/api/server/services/Config/__tests__/getCachedTools.spec.js b/api/server/services/Config/__tests__/getCachedTools.spec.js index 3f85a018f0..6d3947392f 100644 --- a/api/server/services/Config/__tests__/getCachedTools.spec.js +++ b/api/server/services/Config/__tests__/getCachedTools.spec.js @@ -1,4 +1,4 @@ -const { CacheKeys } = require('librechat-data-provider'); +const { CacheKeys, Time } = require('librechat-data-provider'); jest.mock('~/cache/getLogStores'); const getLogStores = require('~/cache/getLogStores'); @@ -9,76 +9,453 @@ getLogStores.mockReturnValue(mockCache); const { ToolCacheKeys, getCachedTools, + updateCachedGlobalTools, setCachedTools, + setCachedToolsIfCurrent, + getMCPToolsCacheGeneration, + renewMCPToolsCacheGeneration, + getCachedAppServerTools, + getNextAppToolsPublicationRevision, + setCachedAppServerTools, + runWithGlobalCacheLock, invalidateCachedTools, } = require('../getCachedTools'); -describe('getCachedTools', () => { +describe('MCP tool cache', () => { beforeEach(() => { jest.clearAllMocks(); getLogStores.mockReturnValue(mockCache); }); - describe('ToolCacheKeys.MCP_SERVER', () => { - it('should generate cache keys that include userId', () => { - const key = ToolCacheKeys.MCP_SERVER('user123', 'github'); - expect(key).toBe('tools:mcp:user123:github'); + it('uses collision-safe configuration-addressed keys', () => { + expect(ToolCacheKeys.MCP_APP_SERVER('server:name', 'config/a')).toBe( + 'tools:mcp:app:server%3Aname:config%2Fa', + ); + expect(ToolCacheKeys.MCP_SERVER('tenant:user', 'server:name', 'config/a')).toBe( + 'tools:mcp:user:{tenant%3Auser:server%3Aname}:config%2Fa', + ); + expect(ToolCacheKeys.MCP_SERVER('tenant:user', 'server:name', 'config/a')).not.toBe( + ToolCacheKeys.MCP_SERVER('tenant', 'user:server:name', 'config/a'), + ); + expect(ToolCacheKeys.MCP_SERVER_GENERATION('tenant:user', 'server:name')).toBe( + 'tools:metadata:mcp:user-generation:{tenant%3Auser:server%3Aname}', + ); + expect(ToolCacheKeys.MCP_SERVER_GENERATION('tenant:user', 'server:name')).not.toBe( + ToolCacheKeys.MCP_SERVER_GENERATION('tenant', 'user:server:name'), + ); + expect(ToolCacheKeys.MCP_SERVER_LEGACY_FENCE('tenant:user', 'server:name')).toBe( + 'tools:metadata:mcp:user-legacy-fence:{tenant%3Auser:server%3Aname}', + ); + }); + + it('keeps the legacy user key available for non-generation callers', () => { + expect(ToolCacheKeys.MCP_SERVER('user123', 'github')).toBe('tools:mcp:user123:github'); + }); + + it('gets and sets static global tools without touching MCP slices', async () => { + const tools = { builtin: { type: 'function' } }; + mockCache.get.mockResolvedValue(tools); + mockCache.set.mockResolvedValue(true); + + await expect(getCachedTools()).resolves.toBe(tools); + await expect(setCachedTools(tools)).resolves.toBe(true); + + expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL); + expect(mockCache.set).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL, tools, expect.any(Number)); + expect(mockCache.delete).not.toHaveBeenCalled(); + }); + + it('updates the global catalog atomically through the catalog store', async () => { + const current = { builtin: { type: 'function' }, old_mcp_server: { type: 'function' } }; + mockCache.get.mockResolvedValue(current); + mockCache.set.mockResolvedValue(true); + + await updateCachedGlobalTools(({ old_mcp_server: _removed, ...staticTools }) => staticTools); + + expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL); + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.GLOBAL, + { builtin: { type: 'function' } }, + Time.TWELVE_HOURS, + ); + }); + + it('recreates the authoritative global catalog after its cache entry expires', async () => { + mockCache.get.mockResolvedValue(null); + mockCache.set.mockResolvedValue(true); + + await updateCachedGlobalTools(() => ({ builtin: { type: 'function' } })); + + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.GLOBAL, + { builtin: { type: 'function' } }, + Time.TWELVE_HOURS, + ); + }); + + it('gets and sets an authoritative app slice, including an empty catalog', async () => { + mockCache.get.mockResolvedValue({}); + mockCache.set.mockResolvedValue(true); + + await expect(getCachedAppServerTools('github', 'config-v2')).resolves.toEqual({}); + await expect(setCachedAppServerTools('github', 'config-v2', {})).resolves.toBe(true); + + const key = ToolCacheKeys.MCP_APP_SERVER('github', 'config-v2'); + expect(mockCache.get).toHaveBeenCalledWith(key); + expect(mockCache.set).toHaveBeenCalledWith( + key, + { version: 1, publicationRevision: '0', tools: {} }, + expect.any(Number), + ); + }); + + it('prevents a slow older app snapshot from replacing a newer revision', async () => { + const key = ToolCacheKeys.MCP_APP_SERVER('github', 'config-v2'); + let cached = null; + mockCache.get.mockImplementation(async (requestedKey) => + requestedKey === key ? cached : null, + ); + mockCache.set.mockImplementation(async (requestedKey, value) => { + if (requestedKey === key) { + cached = value; + } + return true; + }); + + const older = await getNextAppToolsPublicationRevision('github', 'config-v2'); + const newer = await getNextAppToolsPublicationRevision('github', 'config-v2'); + const currentTools = { current: { type: 'function' } }; + const staleTools = { stale: { type: 'function' } }; + + await expect(setCachedAppServerTools('github', 'config-v2', currentTools, newer)).resolves.toBe( + true, + ); + await expect(setCachedAppServerTools('github', 'config-v2', staleTools, older)).resolves.toBe( + false, + ); + await expect(getCachedAppServerTools('github', 'config-v2')).resolves.toEqual(currentTools); + }); + + it('allows a completed snapshot when a later reserved request aborts', async () => { + const key = ToolCacheKeys.MCP_APP_SERVER('github', 'config-v2'); + let cached = null; + mockCache.get.mockImplementation(async (requestedKey) => + requestedKey === key ? cached : null, + ); + mockCache.set.mockImplementation(async (requestedKey, value) => { + if (requestedKey === key) { + cached = value; + } + return true; + }); + + const completed = await getNextAppToolsPublicationRevision('github', 'config-v2'); + await getNextAppToolsPublicationRevision('github', 'config-v2'); + + await expect( + setCachedAppServerTools('github', 'config-v2', { completed: {} }, completed), + ).resolves.toBe(true); + await expect(getCachedAppServerTools('github', 'config-v2')).resolves.toEqual({ + completed: {}, }); }); - describe('TOOL_CACHE namespace usage', () => { - it('getCachedTools should use TOOL_CACHE namespace', async () => { - mockCache.get.mockResolvedValue(null); - await getCachedTools(); - expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE); + it('stores unguarded user tools under the supplied config generation', async () => { + const tools = { search: { type: 'function' } }; + mockCache.set.mockResolvedValue(true); + + await setCachedTools(tools, { + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', }); - it('getCachedTools with MCP server options should use TOOL_CACHE namespace', async () => { - mockCache.get.mockResolvedValue({ tool1: {} }); - await getCachedTools({ userId: 'user1', serverName: 'github' }); - expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE); - expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github')); - }); + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'), + tools, + expect.any(Number), + ); + }); - it('setCachedTools should use TOOL_CACHE namespace', async () => { - mockCache.set.mockResolvedValue(true); - const tools = { tool1: { type: 'function' } }; - await setCachedTools(tools); - expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE); - expect(mockCache.set).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL, tools, expect.any(Number)); - }); + it('writes guarded user tools under both config and connection generations', async () => { + const tools = { current: { type: 'function' } }; + mockCache.get.mockResolvedValue('connection-a'); + mockCache.set.mockResolvedValue(true); - it('setCachedTools with MCP server options should use TOOL_CACHE namespace', async () => { - mockCache.set.mockResolvedValue(true); - const tools = { tool1: { type: 'function' } }; - await setCachedTools(tools, { userId: 'user1', serverName: 'github' }); - expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE); - expect(mockCache.set).toHaveBeenCalledWith( - ToolCacheKeys.MCP_SERVER('user1', 'github'), + await expect( + setCachedToolsIfCurrent(tools, { + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + publicationGeneration: 'connection-a', + }), + ).resolves.toBe(true); + + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'), + { version: 1, publicationGeneration: 'connection-a', tools }, + expect.any(Number), + ); + }); + + it('reads guarded user tools only while their connection generation is current', async () => { + const tools = { current: { type: 'function' } }; + mockCache.get + .mockResolvedValueOnce({ + version: 1, + publicationGeneration: 'connection-a', tools, - expect.any(Number), - ); - }); + }) + .mockResolvedValueOnce('connection-a'); - it('invalidateCachedTools should use TOOL_CACHE namespace', async () => { - mockCache.delete.mockResolvedValue(true); - await invalidateCachedTools({ invalidateGlobal: true }); - expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE); - expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL); - }); + await expect( + getCachedTools({ + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + }), + ).resolves.toEqual(tools); - it('should NOT use CONFIG_STORE namespace', async () => { - mockCache.get.mockResolvedValue(null); - await getCachedTools(); - await getCachedTools({ userId: 'user1', serverName: 'github' }); - mockCache.set.mockResolvedValue(true); - await setCachedTools({ tool1: {} }); - mockCache.delete.mockResolvedValue(true); - await invalidateCachedTools({ invalidateGlobal: true }); + expect(mockCache.get.mock.calls[0][0]).toBe( + ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'), + ); + }); - const allCalls = getLogStores.mock.calls.flat(); - expect(allCalls).not.toContain(CacheKeys.CONFIG_STORE); - expect(allCalls.every((key) => key === CacheKeys.TOOL_CACHE)).toBe(true); - }); + it('copies a legacy user catalog into the config-addressed key on rollout', async () => { + const tools = { legacy: { type: 'function' } }; + mockCache.get + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(tools) + .mockResolvedValueOnce('connection-a') + .mockResolvedValueOnce('connection-a'); + mockCache.set.mockResolvedValue(true); + + await expect( + getCachedTools({ + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + }), + ).resolves.toBe(tools); + + expect(mockCache.get).toHaveBeenNthCalledWith(4, ToolCacheKeys.MCP_SERVER('user1', 'github')); + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'), + { + version: 1, + publicationGeneration: 'connection-a', + tools, + }, + Time.TWELVE_HOURS, + ); + }); + + it('creates a generation fence before migrating a legacy user catalog', async () => { + const tools = { legacy: { type: 'function' } }; + mockCache.get + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(tools) + .mockResolvedValueOnce(null) + .mockImplementationOnce(async () => mockCache.set.mock.calls[0][1]); + mockCache.set.mockResolvedValue(true); + + await expect( + getCachedTools({ + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + }), + ).resolves.toBe(tools); + + const [generationKey, generation] = mockCache.set.mock.calls[0]; + expect(generationKey).toBe(ToolCacheKeys.MCP_SERVER_GENERATION('user1', 'github')); + expect(generation).toEqual(expect.any(String)); + expect(mockCache.set.mock.calls[1][1]).toEqual( + expect.objectContaining({ publicationGeneration: generation, tools }), + ); + }); + + it('rejects a legacy catalog recreated after the scope was fenced', async () => { + mockCache.get + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(true); + + await expect( + getCachedTools({ + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + }), + ).resolves.toBeNull(); + + expect(mockCache.get).not.toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github')); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('preserves a config-addressed catalog published during legacy fallback', async () => { + const current = { current: { type: 'function' } }; + mockCache.get.mockResolvedValueOnce(null).mockResolvedValueOnce(current); + + await expect( + getCachedTools({ + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + }), + ).resolves.toBe(current); + + expect(mockCache.get).toHaveBeenCalledTimes(2); + expect(mockCache.set).not.toHaveBeenCalled(); + }); + + it('hides a guarded entry after its connection generation is replaced', async () => { + mockCache.get + .mockResolvedValueOnce({ + version: 1, + publicationGeneration: 'connection-a', + tools: { stale: {} }, + }) + .mockResolvedValueOnce('connection-b'); + + await expect( + getCachedTools({ + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v1', + }), + ).resolves.toBeNull(); + }); + + it('cannot let a late old-config write replace the current config key', async () => { + mockCache.get.mockResolvedValue('connection-a'); + mockCache.set.mockResolvedValue(true); + + await setCachedToolsIfCurrent( + { current: {} }, + { + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v2', + publicationGeneration: 'connection-a', + }, + ); + await setCachedToolsIfCurrent( + { stale: {} }, + { + userId: 'user1', + serverName: 'github', + configGeneration: 'config-v1', + publicationGeneration: 'connection-a', + }, + ); + + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'), + expect.objectContaining({ tools: { current: {} } }), + expect.any(Number), + ); + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v1'), + expect.objectContaining({ tools: { stale: {} } }), + expect.any(Number), + ); + }); + + it('creates and reuses a durable connection publication generation', async () => { + mockCache.get + .mockResolvedValueOnce(null) + .mockResolvedValueOnce(null) + .mockResolvedValueOnce('existing-generation'); + mockCache.set.mockResolvedValue(true); + + const created = await getMCPToolsCacheGeneration({ userId: 'user1', serverName: 'github' }); + const existing = await getMCPToolsCacheGeneration({ userId: 'user1', serverName: 'github' }); + + expect(created).toEqual(expect.any(String)); + expect(existing).toBe('existing-generation'); + expect(mockCache.set).toHaveBeenCalledWith( + ToolCacheKeys.MCP_SERVER_GENERATION('user1', 'github'), + created, + expect.any(Number), + ); + }); + + it('renews a lease only for its current publication generation', async () => { + mockCache.get.mockResolvedValue('connection-a'); + mockCache.set.mockResolvedValue(true); + + await expect( + renewMCPToolsCacheGeneration({ + userId: 'user1', + serverName: 'github', + publicationGeneration: 'connection-a', + }), + ).resolves.toBe(true); + await expect( + renewMCPToolsCacheGeneration({ + userId: 'user1', + serverName: 'github', + publicationGeneration: 'connection-b', + }), + ).resolves.toBe(false); + }); + + it('rotates the connection generation before deleting the legacy user key', async () => { + mockCache.set.mockResolvedValue(true); + mockCache.delete.mockResolvedValue(true); + + await invalidateCachedTools({ userId: 'user1', serverName: 'github' }); + + expect(mockCache.set).toHaveBeenNthCalledWith( + 1, + ToolCacheKeys.MCP_SERVER_LEGACY_FENCE('user1', 'github'), + true, + expect.any(Number), + ); + expect(mockCache.set.mock.calls[0][2]).toBeGreaterThanOrEqual(Time.ONE_DAY); + expect(mockCache.set).toHaveBeenNthCalledWith( + 2, + ToolCacheKeys.MCP_SERVER_GENERATION('user1', 'github'), + expect.any(String), + expect.any(Number), + ); + expect(mockCache.set.mock.calls[1][2]).toBeGreaterThanOrEqual(Time.ONE_DAY); + expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github')); + expect(mockCache.set.mock.invocationCallOrder[1]).toBeLessThan( + mockCache.delete.mock.invocationCallOrder[0], + ); + }); + + it('invalidates only the static global key for broad config changes', async () => { + mockCache.delete.mockResolvedValue(true); + + await invalidateCachedTools({ invalidateGlobal: true }); + + expect(mockCache.delete).toHaveBeenCalledTimes(1); + expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL); + }); + + it('runs global cache operations directly when the cache is in memory', async () => { + const operation = jest.fn().mockResolvedValue('done'); + await expect(runWithGlobalCacheLock(operation)).resolves.toBe('done'); + expect(operation).toHaveBeenCalledTimes(1); + }); + + it('uses only the TOOL_CACHE namespace', async () => { + mockCache.get.mockResolvedValue(null); + mockCache.set.mockResolvedValue(true); + mockCache.delete.mockResolvedValue(true); + + await getCachedTools(); + await getCachedAppServerTools('github', 'config-v2'); + await setCachedTools({}); + await invalidateCachedTools({ invalidateGlobal: true }); + + expect(getLogStores.mock.calls.flat().every((key) => key === CacheKeys.TOOL_CACHE)).toBe(true); }); }); diff --git a/api/server/services/Config/getCachedTools.js b/api/server/services/Config/getCachedTools.js index 2877234b58..0dc0da0897 100644 --- a/api/server/services/Config/getCachedTools.js +++ b/api/server/services/Config/getCachedTools.js @@ -1,89 +1,20 @@ -const { CacheKeys, Time } = require('librechat-data-provider'); +const { CacheKeys } = require('librechat-data-provider'); +const { + cacheConfig, + ioredisClient, + keyvRedisClient, + mcpConfig, + ToolCacheKeys, + createMCPCatalogStore, +} = require('@librechat/api'); const getLogStores = require('~/cache/getLogStores'); -/** - * Cache key generators for different tool access patterns - */ -const ToolCacheKeys = { - /** Global tools available to all users */ - GLOBAL: 'tools:global', - /** MCP tools cached by user ID and server name */ - MCP_SERVER: (userId, serverName) => `tools:mcp:${userId}:${serverName}`, -}; +const store = createMCPCatalogStore({ + cacheConfig, + ioredisClient, + keyvRedisClient, + userConnectionIdleTimeout: mcpConfig.USER_CONNECTION_IDLE_TIMEOUT, + getCache: () => getLogStores(CacheKeys.TOOL_CACHE), +}); -/** - * Retrieves available tools from cache - * @function getCachedTools - * @param {Object} options - Options for retrieving tools - * @param {string} [options.userId] - User ID for user-specific MCP tools - * @param {string} [options.serverName] - MCP server name to get cached tools for - * @returns {Promise} The available tools object or null if not cached - */ -async function getCachedTools(options = {}) { - const cache = getLogStores(CacheKeys.TOOL_CACHE); - const { userId, serverName } = options; - - // Return MCP server-specific tools if requested - if (serverName && userId) { - return await cache.get(ToolCacheKeys.MCP_SERVER(userId, serverName)); - } - - // Default to global tools - return await cache.get(ToolCacheKeys.GLOBAL); -} - -/** - * Sets available tools in cache - * @function setCachedTools - * @param {Object} tools - The tools object to cache - * @param {Object} options - Options for caching tools - * @param {string} [options.userId] - User ID for user-specific MCP tools - * @param {string} [options.serverName] - MCP server name for server-specific tools - * @param {number} [options.ttl] - Time to live in milliseconds (default: 12 hours) - * @returns {Promise} Whether the operation was successful - */ -async function setCachedTools(tools, options = {}) { - const cache = getLogStores(CacheKeys.TOOL_CACHE); - const { userId, serverName, ttl = Time.TWELVE_HOURS } = options; - - // Cache by MCP server if specified (requires userId) - if (serverName && userId) { - return await cache.set(ToolCacheKeys.MCP_SERVER(userId, serverName), tools, ttl); - } - - // Default to global cache - return await cache.set(ToolCacheKeys.GLOBAL, tools, ttl); -} - -/** - * Invalidates cached tools - * @function invalidateCachedTools - * @param {Object} options - Options for invalidating tools - * @param {string} [options.userId] - User ID for user-specific MCP tools - * @param {string} [options.serverName] - MCP server name to invalidate - * @param {boolean} [options.invalidateGlobal=false] - Whether to invalidate global tools - * @returns {Promise} - */ -async function invalidateCachedTools(options = {}) { - const cache = getLogStores(CacheKeys.TOOL_CACHE); - const { userId, serverName, invalidateGlobal = false } = options; - - const keysToDelete = []; - - if (invalidateGlobal) { - keysToDelete.push(ToolCacheKeys.GLOBAL); - } - - if (serverName && userId) { - keysToDelete.push(ToolCacheKeys.MCP_SERVER(userId, serverName)); - } - - await Promise.all(keysToDelete.map((key) => cache.delete(key))); -} - -module.exports = { - ToolCacheKeys, - getCachedTools, - setCachedTools, - invalidateCachedTools, -}; +module.exports = { ToolCacheKeys, ...store }; diff --git a/api/server/services/Config/mcp.js b/api/server/services/Config/mcp.js index 2bd64cc31b..2a6a66d8fa 100644 --- a/api/server/services/Config/mcp.js +++ b/api/server/services/Config/mcp.js @@ -1,17 +1,43 @@ const { createMCPToolCacheService, MCPServersRegistry } = require('@librechat/api'); -const { getCachedTools, setCachedTools } = require('./getCachedTools'); +const { + getCachedTools, + updateCachedGlobalTools, + setCachedToolsWithinGlobalLock, + getCachedAppServerTools, + setCachedAppServerTools, + setCachedToolsIfCurrent, + getMCPToolsCacheGeneration, + renewMCPToolsCacheGeneration, + getNextAppToolsPublicationRevision, +} = require('./getCachedTools'); -const { mergeAppTools, cacheMCPServerTools, updateMCPServerTools, getMCPServerTools } = - createMCPToolCacheService({ - getCachedTools, - setCachedTools, - getServerConfig: (serverName, userId) => - MCPServersRegistry.getInstance().getServerConfig(serverName, userId), - }); +const { + syncStaticTools, + mergeAppTools, + cacheMCPServerTools, + updateMCPServerTools, + getMCPServerTools, +} = createMCPToolCacheService({ + getCachedTools, + updateCachedGlobalTools, + setCachedTools: setCachedToolsWithinGlobalLock, + setCachedToolsIfCurrent, + getCachedAppServerTools, + setCachedAppServerTools, + getServerConfig: (serverName, userId) => + MCPServersRegistry.getInstance().getServerConfig(serverName, userId), + getAllServerConfigs: () => MCPServersRegistry.getInstance().getAllServerConfigs(), + isAppServerConfig: (serverName, effectiveConfig) => + MCPServersRegistry.getInstance().isAppServerConfig(serverName, effectiveConfig), +}); module.exports = { + syncStaticTools, mergeAppTools, getMCPServerTools, cacheMCPServerTools, updateMCPServerTools, + getMCPToolsCacheGeneration, + renewMCPToolsCacheGeneration, + getNextAppToolsPublicationRevision, }; diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index acc8a41034..6eb0c723a3 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -12,6 +12,7 @@ const { normalizeMCPToolKey, buildServerNameAliases, findShadowedServerNames, + getAssistantToolDefinitions: loadAssistantToolDefinitions, resolveMCPServerContext, normalizeJsonSchema, GenerationJobManager, @@ -27,6 +28,7 @@ const { isUserSourced, checkAccessWithRequestCache, getMissingCustomUserVars, + getUserMCPAuthMap, getServerCustomUserVars, requiresEphemeralUserConnection, requiresOAuthMachinery, @@ -49,12 +51,17 @@ const { getMCPManager, } = require('~/config'); const db = require('~/models'); -const { findToken, createToken, updateToken, deleteTokens } = db; +const { findToken, createToken, updateToken, deleteTokens, findPluginAuthsByKeys } = db; const { getGraphApiToken } = require('./GraphTokenService'); const { exchangeOboToken } = require('./OboTokenService'); const { createOboTrustChecker } = require('./OboPolicyService'); const { reinitMCPServer } = require('./Tools/mcp'); -const { getAppConfig } = require('./Config'); +const { + getAppConfig, + getCachedTools, + getMCPServerTools, + cacheMCPServerTools, +} = require('./Config'); const { getLogStores } = require('~/cache'); const MAX_CACHE_SIZE = 1000; @@ -303,6 +310,57 @@ async function healMcpToolNames({ req, tools, toolDefinitions }) { return healedList; } +/** + * Loads static and MCP function definitions used by assistant create/update writes. MCP catalogs + * are stored per server and effective config, so assistant writers must resolve the referenced + * server slices instead of relying on the static aggregate cache. + * @param {object} params + * @param {ServerRequest} params.req + * @param {Array} [params.tools] + * @returns {Promise} + */ +async function getAssistantToolDefinitions({ req, tools }) { + const registry = getMCPServersRegistry(); + const appConfig = await getAppConfigForRequest(req); + return await loadAssistantToolDefinitions( + { + user: req.user, + tools, + staticTools: (await getCachedTools()) ?? {}, + mcpConfig: appConfig?.mcpConfig ?? {}, + }, + { + ensureConfigServers: (mcpConfig) => registry.ensureConfigServers(mcpConfig), + getAllServerConfigs: (userId, configServers, role) => + registry.getAllServerConfigs(userId, configServers, role), + getMCPServerTools, + getServerToolFunctionsSnapshot: async (userId, serverName, serverConfig) => + (await getMCPManager()?.getServerToolFunctionsSnapshot( + userId, + serverName, + serverConfig, + )) ?? { + tools: null, + }, + recoverServerTools: async (serverName, serverConfig) => { + const userMCPAuthMap = await getUserMCPAuthMap({ + userId: req.user.id, + servers: [serverName], + findPluginAuthsByKeys, + }); + const result = await reinitMCPServer({ + user: req.user, + serverName, + serverConfig, + userMCPAuthMap, + }); + return result?.availableTools ?? null; + }, + cacheMCPServerTools, + }, + ); +} + /** * Resolves the name set MCP collision guards audit against. Prefers the * caller-threaded accessible set; self-fetches only when a configured name @@ -1431,6 +1489,7 @@ module.exports = { resolveMcpServerContext, getAccessibleMcpServerNames, healMcpToolNames, + getAssistantToolDefinitions, resolveCollisionAuditNames, resolveMcpConfigNames, resolveAllMcpConfigs, diff --git a/api/server/services/Tools/mcp.js b/api/server/services/Tools/mcp.js index 24b8bcd5eb..1560694419 100644 --- a/api/server/services/Tools/mcp.js +++ b/api/server/services/Tools/mcp.js @@ -10,7 +10,7 @@ const { findToken, createToken, updateToken, deleteTokens } = require('~/models' const { getGraphApiToken } = require('~/server/services/GraphTokenService'); const { exchangeOboToken } = require('~/server/services/OboTokenService'); const { createOboTrustChecker } = require('~/server/services/OboPolicyService'); -const { updateMCPServerTools } = require('~/server/services/Config'); +const { getMCPToolsCacheGeneration, updateMCPServerTools } = require('~/server/services/Config'); const { getLogStores } = require('~/cache'); const MCP_REINITIALIZE_FAILURE_REASONS = { @@ -65,6 +65,7 @@ async function reinitMCPServer({ let oauthUrl = null; let oauthExpiresAt; let ephemeralServer = false; + let publicationGeneration; try { const registry = getMCPServersRegistry(); @@ -167,6 +168,13 @@ async function reinitMCPServer({ const mcpManager = getMCPManager(); const tokenMethods = { findToken, updateToken, createToken, deleteTokens }; + if (!ephemeralServer) { + publicationGeneration = await getMCPToolsCacheGeneration({ + userId: user.id, + serverName, + }); + } + const oauthStart = _oauthStart ?? (async (authURL, options) => { @@ -259,7 +267,36 @@ async function reinitMCPServer({ } if (connection && !oauthRequired) { - tools = await connection.fetchTools(); + publicationGeneration = + mcpManager.getToolPublicationGeneration(connection) ?? publicationGeneration; + let snapshot; + if (typeof connection.fetchOrderedToolsSnapshot === 'function') { + snapshot = await connection.fetchOrderedToolsSnapshot(); + } else if (typeof connection.fetchToolsSnapshot === 'function') { + snapshot = await connection.fetchToolsSnapshot(); + } else { + snapshot = { tools: await connection.fetchTools(), complete: true }; + } + if (snapshot.complete) { + tools = snapshot.tools; + } else { + logger.warn( + `[MCP Reinitialize] Preserving cached tools for ${serverName} because tools/list returned an incomplete snapshot`, + ); + } + } + + if (tools && !ephemeralServer && publicationGeneration) { + const currentGeneration = await getMCPToolsCacheGeneration({ + userId: user.id, + serverName, + }); + if (currentGeneration !== publicationGeneration) { + logger.warn( + `[MCP Reinitialize] Discarding stale tools for ${serverName} because its publication generation changed during discovery`, + ); + tools = null; + } } if (tools) { @@ -268,7 +305,11 @@ async function reinitMCPServer({ serverName, tools, serverConfig, + ...(publicationGeneration && { publicationGeneration }), }); + if (availableTools == null) { + tools = null; + } } logger.debug( @@ -325,12 +366,9 @@ async function reinitMCPServer({ } finally { if (connection && ephemeralServer && !requestScopedConnections) { try { - await connection.disconnect(); + await connection.dispose(); } catch (error) { - logger.warn( - `[MCP Reinitialize] Failed to disconnect ephemeral server ${serverName}`, - error, - ); + logger.warn(`[MCP Reinitialize] Failed to dispose ephemeral server ${serverName}`, error); } } } diff --git a/api/server/services/Tools/mcp.spec.js b/api/server/services/Tools/mcp.spec.js index d8e788390e..65029bd68b 100644 --- a/api/server/services/Tools/mcp.spec.js +++ b/api/server/services/Tools/mcp.spec.js @@ -4,11 +4,14 @@ const mockGetConnection = jest.fn(); const mockDiscoverServerTools = jest.fn(); const mockGetGraphApiToken = jest.fn(); const mockUpdateMCPServerTools = jest.fn(); +const mockGetMCPToolsCacheGeneration = jest.fn().mockResolvedValue('generation-current'); +const mockGetToolPublicationGeneration = jest.fn().mockReturnValue('generation-current'); jest.mock('~/config', () => ({ getMCPManager: jest.fn(() => ({ getConnection: mockGetConnection, discoverServerTools: mockDiscoverServerTools, + getToolPublicationGeneration: mockGetToolPublicationGeneration, })), getMCPServersRegistry: jest.fn(() => ({ getServerConfig: jest.fn() })), getFlowStateManager: jest.fn(() => ({})), @@ -21,6 +24,7 @@ jest.mock('~/models', () => ({ })); jest.mock('~/server/services/Config', () => ({ updateMCPServerTools: mockUpdateMCPServerTools, + getMCPToolsCacheGeneration: mockGetMCPToolsCacheGeneration, })); jest.mock('~/server/services/GraphTokenService', () => ({ getGraphApiToken: mockGetGraphApiToken, @@ -117,9 +121,71 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => { serverName, tools: [], serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' }, + publicationGeneration: 'generation-current', }); }); + it('preserves cached tools when live recovery returns an incomplete snapshot', async () => { + const fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + tools: [{ name: 'partial', inputSchema: { type: 'object' } }], + complete: false, + }); + mockGetConnection.mockResolvedValue({ + fetchOrderedToolsSnapshot, + }); + + const result = await reinitMCPServer({ + user, + serverName, + serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' }, + }); + + expect(result.tools).toBeNull(); + expect(fetchOrderedToolsSnapshot).toHaveBeenCalledTimes(1); + expect(mockUpdateMCPServerTools).not.toHaveBeenCalled(); + }); + + it('discards a snapshot when another replica rotates its generation during discovery', async () => { + mockGetMCPToolsCacheGeneration + .mockResolvedValueOnce('generation-current') + .mockResolvedValueOnce('generation-replaced'); + mockGetConnection.mockResolvedValue({ + fetchOrderedToolsSnapshot: jest.fn().mockResolvedValue({ + tools: [{ name: 'stale', inputSchema: { type: 'object' } }], + complete: true, + }), + }); + + const result = await reinitMCPServer({ + user, + serverName, + serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' }, + }); + + expect(result.tools).toBeNull(); + expect(result.availableTools).toBeNull(); + expect(mockUpdateMCPServerTools).not.toHaveBeenCalled(); + }); + + it('does not return tools when the guarded publication loses its generation race', async () => { + mockGetConnection.mockResolvedValue({ + fetchOrderedToolsSnapshot: jest.fn().mockResolvedValue({ + tools: [{ name: 'stale', inputSchema: { type: 'object' } }], + complete: true, + }), + }); + mockUpdateMCPServerTools.mockResolvedValue(null); + + const result = await reinitMCPServer({ + user, + serverName, + serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' }, + }); + + expect(result.tools).toBeNull(); + expect(result.availableTools).toBeNull(); + }); + it('passes request body and Graph resolver into connection creation', async () => { mockGetConnection.mockResolvedValue({ fetchTools: jest.fn().mockResolvedValue([]) }); const requestBody = { conversationId: 'conv-123', messageId: 'msg-123' }; @@ -167,8 +233,8 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => { ); }); - it('disconnects ephemeral BODY-scoped connections after loading tools', async () => { - const disconnect = jest.fn().mockResolvedValue(undefined); + it('disposes ephemeral BODY-scoped connections after loading tools', async () => { + const dispose = jest.fn().mockResolvedValue(undefined); const tools = [{ name: 'search', inputSchema: { type: 'object', properties: {} } }]; const serverConfig = { type: 'streamable-http', @@ -176,7 +242,7 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => { source: 'yaml', }; mockGetConnection.mockResolvedValue({ - disconnect, + dispose, fetchTools: jest.fn().mockResolvedValue(tools), }); @@ -188,7 +254,7 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => { userMCPAuthMap: undefined, }); - expect(disconnect).toHaveBeenCalledTimes(1); + expect(dispose).toHaveBeenCalledTimes(1); expect(mockUpdateMCPServerTools).toHaveBeenCalledWith( expect.objectContaining({ tools, @@ -261,9 +327,8 @@ describe('reinitMCPServer — runtime BODY placeholder pre-check (issue #14074)' }); it('connects normally when the request body provides the placeholder fields', async () => { - const disconnect = jest.fn().mockResolvedValue(undefined); mockGetConnection.mockResolvedValue({ - disconnect, + dispose: jest.fn().mockResolvedValue(undefined), fetchTools: jest.fn().mockResolvedValue([]), }); diff --git a/api/server/services/__tests__/MCP.spec.js b/api/server/services/__tests__/MCP.spec.js index 76df1c321d..aa4ee0c080 100644 --- a/api/server/services/__tests__/MCP.spec.js +++ b/api/server/services/__tests__/MCP.spec.js @@ -20,6 +20,7 @@ jest.mock('~/server/services/Config', () => ({ setCachedTools: jest.fn(), getCachedTools: jest.fn(), getMCPServerTools: jest.fn(), + cacheMCPServerTools: jest.fn(), loadCustomConfig: jest.fn(), })); @@ -32,6 +33,7 @@ jest.mock('@librechat/api', () => ({ isMCPDomainAllowed: jest.fn(), GenerationJobManager: jest.fn(), buildOAuthToolCallName: jest.fn((name) => name), + getUserMCPAuthMap: jest.fn(), /** Mirrors the real resolver so these tests still exercise the wrapper's own * plumbing - loading the request config and degrading on failure - rather than * the resolution logic, which is unit-tested in packages/api. Like the real @@ -53,6 +55,7 @@ jest.mock('~/models', () => ({ findToken: jest.fn(), createToken: jest.fn(), updateToken: jest.fn(), + findPluginAuthsByKeys: jest.fn(), })); jest.mock('~/server/services/GraphTokenService', () => ({ getGraphApiToken: jest.fn(), @@ -69,11 +72,18 @@ jest.mock('~/server/services/Tools/mcp', () => ({ const { Constants } = require('librechat-data-provider'); -const { getAppConfig } = require('~/server/services/Config'); +const { + getAppConfig, + getCachedTools, + getMCPServerTools, + cacheMCPServerTools, +} = require('~/server/services/Config'); const { reinitMCPServer } = require('~/server/services/Tools/mcp'); +const { getUserMCPAuthMap } = require('@librechat/api'); const { createMCPTool, healMcpToolNames, + getAssistantToolDefinitions, resolveConfigServers, resolveMcpConfigNames, resolveAllMcpConfigs, @@ -81,6 +91,105 @@ const { resolveCollisionAuditNames, } = require('../MCP'); +describe('getAssistantToolDefinitions', () => { + beforeEach(() => { + jest.clearAllMocks(); + require('~/config').getMCPManager.mockReset(); + }); + + const req = { user: { id: 'u1', role: 'user' } }; + const serverConfig = { type: 'streamable-http', url: 'https://app.example.com/mcp' }; + const toolKey = `search${Constants.mcp_delimiter}app-server`; + const mcpDefinition = { type: 'function', function: { name: toolKey } }; + + it('combines static definitions with referenced configuration-addressed MCP slices', async () => { + getCachedTools.mockResolvedValue({ code_interpreter: { type: 'code_interpreter' } }); + getAppConfig.mockResolvedValue({ mcpConfig: {} }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ 'app-server': serverConfig }); + getMCPServerTools.mockResolvedValue({ [toolKey]: mcpDefinition }); + + const definitions = await getAssistantToolDefinitions({ + req, + tools: ['code_interpreter', toolKey], + }); + + expect(definitions).toEqual({ + code_interpreter: { type: 'code_interpreter' }, + [toolKey]: mcpDefinition, + }); + expect(getMCPServerTools).toHaveBeenCalledWith('u1', 'app-server', serverConfig); + }); + + it('recovers and re-caches a referenced server when its slice is missing', async () => { + getCachedTools.mockResolvedValue({}); + getAppConfig.mockResolvedValue({ mcpConfig: {} }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ 'app-server': serverConfig }); + getMCPServerTools.mockResolvedValue(null); + cacheMCPServerTools.mockResolvedValue(undefined); + const getServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({ + tools: { [toolKey]: mcpDefinition }, + publicationGeneration: 'connection-generation', + }); + require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot }); + + await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).resolves.toEqual({ + [toolKey]: mcpDefinition, + }); + expect(cacheMCPServerTools).toHaveBeenCalledWith({ + userId: 'u1', + serverName: 'app-server', + serverTools: { [toolKey]: mcpDefinition }, + serverConfig, + publicationGeneration: 'connection-generation', + }); + }); + + it('reinitializes a referenced server when its cache and local snapshot are missing', async () => { + getCachedTools.mockResolvedValue({}); + getAppConfig.mockResolvedValue({ mcpConfig: {} }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ 'app-server': serverConfig }); + getMCPServerTools.mockResolvedValue(null); + const getServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({ tools: null }); + require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot }); + const userMCPAuthMap = { 'mcp_app-server': { API_KEY: 'saved' } }; + getUserMCPAuthMap.mockResolvedValue(userMCPAuthMap); + reinitMCPServer.mockResolvedValue({ availableTools: { [toolKey]: mcpDefinition } }); + + await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).resolves.toEqual({ + [toolKey]: mcpDefinition, + }); + expect(reinitMCPServer).toHaveBeenCalledWith({ + user: req.user, + serverName: 'app-server', + serverConfig, + userMCPAuthMap, + }); + expect(getUserMCPAuthMap).toHaveBeenCalledWith({ + userId: 'u1', + servers: ['app-server'], + findPluginAuthsByKeys: expect.any(Function), + }); + }); + + it('propagates config-server resolution failures through the assistant write bridge', async () => { + const resolutionError = new Error('config resolution failed'); + getCachedTools.mockResolvedValue({}); + getAppConfig.mockResolvedValue({ + mcpConfig: { 'app-server': { type: 'streamable-http', url: 'https://example.com/mcp' } }, + }); + mockRegistry.ensureConfigServers.mockRejectedValue(resolutionError); + + await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).rejects.toBe( + resolutionError, + ); + expect(mockRegistry.getAllServerConfigs).not.toHaveBeenCalled(); + expect(getMCPServerTools).not.toHaveBeenCalled(); + }); +}); + describe('resolveConfigServers', () => { beforeEach(() => jest.clearAllMocks()); diff --git a/api/server/services/initializeMCPs.js b/api/server/services/initializeMCPs.js index e3b35a6e86..130bde998a 100644 --- a/api/server/services/initializeMCPs.js +++ b/api/server/services/initializeMCPs.js @@ -1,6 +1,19 @@ const mongoose = require('mongoose'); const { logger } = require('@librechat/data-schemas'); -const { mergeAppTools, getAppConfig } = require('./Config'); +const { + registerShutdownTask, + setMCPToolsChangedHandler, + setMCPToolsChangedGenerationHandler, + setMCPToolsChangedGenerationRenewalHandler, + setMCPToolsChangedRevisionHandler, +} = require('@librechat/api'); +const { syncStaticTools, mergeAppTools, getAppConfig } = require('./Config'); +const { + getMCPToolsCacheGeneration, + renewMCPToolsCacheGeneration, + getNextAppToolsPublicationRevision, + updateMCPServerTools, +} = require('./Config/mcp'); const { createMCPServersRegistry, createMCPManager } = require('~/config'); /** @@ -18,6 +31,36 @@ async function resolveMCPAllowlists(ctx) { }; } +/** + * Refreshes one server's tools after it reported `notifications/tools/list_changed`. + * + * A server that builds tools at runtime is the case this exists for: without it the tool list + * stayed frozen at connection time and only a restart picked up the change (#7117). The list is + * re-fetched from the live connection and written over that server's cache entry, so tools that + * disappeared stop being advertised too. + */ +async function refreshChangedServerTools({ + serverName, + userId, + tools, + serverConfig, + publicationGeneration, + publicationRevision, +}) { + await updateMCPServerTools({ + userId, + serverName, + tools, + serverConfig, + ...(publicationGeneration && { publicationGeneration }), + ...(publicationRevision && { publicationRevision }), + }); + const toolCount = tools.length; + logger.info( + `[MCP][${serverName}] Tool list changed; refreshed ${toolCount} ${toolCount === 1 ? 'tool' : 'tools'}${userId ? ` for user ${userId}` : ''}`, + ); +} + /** * Initialize MCP servers */ @@ -39,16 +82,28 @@ async function initializeMCPs() { try { const mcpManager = await createMCPManager(mcpServers || {}); + setMCPToolsChangedHandler(refreshChangedServerTools); + setMCPToolsChangedGenerationHandler(getMCPToolsCacheGeneration); + setMCPToolsChangedGenerationRenewalHandler(renewMCPToolsCacheGeneration); + setMCPToolsChangedRevisionHandler(({ serverName, configGeneration }) => + getNextAppToolsPublicationRevision(serverName, configGeneration), + ); + registerShutdownTask('MCP app connections', () => mcpManager.disconnectAppServers()); if (mcpServers && Object.keys(mcpServers).length > 0) { const mcpTools = (await mcpManager.getAppToolFunctions()) || {}; - await mergeAppTools(mcpTools); + try { + await mergeAppTools(mcpTools, appConfig.availableTools || {}); + } finally { + await mcpManager.connectAppServers(); + } const serverCount = Object.keys(mcpServers).length; const toolCount = Object.keys(mcpTools).length; logger.info( `[MCP] Initialized with ${serverCount} configured ${serverCount === 1 ? 'server' : 'servers'} and ${toolCount} ${toolCount === 1 ? 'tool' : 'tools'}.`, ); } else { + await syncStaticTools(appConfig.availableTools || {}); logger.debug('[MCP] No servers configured. MCPManager ready for UI-based servers.'); } } catch (error) { @@ -58,3 +113,4 @@ async function initializeMCPs() { } module.exports = initializeMCPs; +module.exports.refreshChangedServerTools = refreshChangedServerTools; diff --git a/api/server/services/initializeMCPs.spec.js b/api/server/services/initializeMCPs.spec.js index fe0766343c..4939cd507c 100644 --- a/api/server/services/initializeMCPs.spec.js +++ b/api/server/services/initializeMCPs.spec.js @@ -27,6 +27,7 @@ jest.mock('@librechat/data-schemas', () => ({ // Mock config functions const mockGetAppConfig = jest.fn(); +const mockSyncStaticTools = jest.fn(); const mockMergeAppTools = jest.fn(); jest.mock('./Config', () => ({ @@ -36,12 +37,17 @@ jest.mock('./Config', () => ({ get mergeAppTools() { return mockMergeAppTools; }, + get syncStaticTools() { + return mockSyncStaticTools; + }, })); // Mock MCP singletons const mockCreateMCPServersRegistry = jest.fn(); const mockCreateMCPManager = jest.fn(); const mockMCPManagerInstance = { + connectAppServers: jest.fn(), + disconnectAppServers: jest.fn(), getAppToolFunctions: jest.fn(), }; @@ -54,6 +60,49 @@ jest.mock('~/config', () => ({ }, })); +const mockSetMCPToolsChangedHandler = jest.fn(); +const mockSetMCPToolsChangedGenerationHandler = jest.fn(); +const mockSetMCPToolsChangedGenerationRenewalHandler = jest.fn(); +const mockSetMCPToolsChangedRevisionHandler = jest.fn(); +const mockRegisterShutdownTask = jest.fn(); +const mockUpdateMCPServerTools = jest.fn(); +const mockGetMCPToolsCacheGeneration = jest.fn(); +const mockRenewMCPToolsCacheGeneration = jest.fn(); +const mockGetNextAppToolsPublicationRevision = jest.fn(); + +jest.mock('@librechat/api', () => ({ + get registerShutdownTask() { + return mockRegisterShutdownTask; + }, + get setMCPToolsChangedHandler() { + return mockSetMCPToolsChangedHandler; + }, + get setMCPToolsChangedGenerationHandler() { + return mockSetMCPToolsChangedGenerationHandler; + }, + get setMCPToolsChangedGenerationRenewalHandler() { + return mockSetMCPToolsChangedGenerationRenewalHandler; + }, + get setMCPToolsChangedRevisionHandler() { + return mockSetMCPToolsChangedRevisionHandler; + }, +})); + +jest.mock('./Config/mcp', () => ({ + get updateMCPServerTools() { + return mockUpdateMCPServerTools; + }, + get getMCPToolsCacheGeneration() { + return mockGetMCPToolsCacheGeneration; + }, + get renewMCPToolsCacheGeneration() { + return mockRenewMCPToolsCacheGeneration; + }, + get getNextAppToolsPublicationRevision() { + return mockGetNextAppToolsPublicationRevision; + }, +})); + const { logger } = require('@librechat/data-schemas'); const initializeMCPs = require('./initializeMCPs'); @@ -65,6 +114,9 @@ describe('initializeMCPs', () => { mockCreateMCPServersRegistry.mockReturnValue(undefined); mockCreateMCPManager.mockResolvedValue(mockMCPManagerInstance); mockMCPManagerInstance.getAppToolFunctions.mockResolvedValue({}); + mockMCPManagerInstance.connectAppServers.mockResolvedValue(undefined); + mockMCPManagerInstance.disconnectAppServers.mockResolvedValue(undefined); + mockSyncStaticTools.mockResolvedValue(undefined); mockMergeAppTools.mockResolvedValue(undefined); }); @@ -183,6 +235,36 @@ describe('initializeMCPs', () => { expect(mockCreateMCPManager).toHaveBeenCalledWith(mcpServers); }); + it('should register app connections for graceful shutdown', async () => { + mockGetAppConfig.mockResolvedValue({ mcpConfig: null }); + + await initializeMCPs(); + + expect(mockRegisterShutdownTask).toHaveBeenCalledWith( + 'MCP app connections', + expect.any(Function), + ); + const shutdown = mockRegisterShutdownTask.mock.calls[0][1]; + await shutdown(); + expect(mockMCPManagerInstance.disconnectAppServers).toHaveBeenCalledTimes(1); + }); + + it('should wire app publication revision allocation into the cache store', async () => { + mockGetAppConfig.mockResolvedValue({ mcpConfig: null }); + mockGetNextAppToolsPublicationRevision.mockResolvedValue('9'); + + await initializeMCPs(); + + const allocateRevision = mockSetMCPToolsChangedRevisionHandler.mock.calls[0][0]; + await expect( + allocateRevision({ serverName: 'dynamic', configGeneration: 'config-generation' }), + ).resolves.toBe('9'); + expect(mockGetNextAppToolsPublicationRevision).toHaveBeenCalledWith( + 'dynamic', + 'config-generation', + ); + }); + it('should throw and log error if MCPManager initialization fails', async () => { const managerError = new Error('Manager initialization failed'); mockCreateMCPManager.mockRejectedValue(managerError); @@ -197,21 +279,24 @@ describe('initializeMCPs', () => { }); describe('Tool merging behavior', () => { - it('should NOT merge tools when no configured servers exist', async () => { + it('should skip app catalog discovery when no configured servers exist', async () => { mockGetAppConfig.mockResolvedValue({ mcpConfig: null, // No configured servers + availableTools: { builtin: { type: 'function' } }, }); await initializeMCPs(); expect(mockMCPManagerInstance.getAppToolFunctions).not.toHaveBeenCalled(); expect(mockMergeAppTools).not.toHaveBeenCalled(); + expect(mockSyncStaticTools).toHaveBeenCalledWith({ builtin: { type: 'function' } }); + expect(mockMCPManagerInstance.connectAppServers).not.toHaveBeenCalled(); expect(logger.debug).toHaveBeenCalledWith( '[MCP] No servers configured. MCPManager ready for UI-based servers.', ); }); - it('should NOT merge tools when mcpConfig is empty object', async () => { + it('should skip app catalog discovery when mcpConfig is empty', async () => { mockGetAppConfig.mockResolvedValue({ mcpConfig: {}, // Empty object }); @@ -220,6 +305,8 @@ describe('initializeMCPs', () => { expect(mockMCPManagerInstance.getAppToolFunctions).not.toHaveBeenCalled(); expect(mockMergeAppTools).not.toHaveBeenCalled(); + expect(mockSyncStaticTools).toHaveBeenCalledWith({}); + expect(mockMCPManagerInstance.connectAppServers).not.toHaveBeenCalled(); expect(logger.debug).toHaveBeenCalledWith( '[MCP] No servers configured. MCPManager ready for UI-based servers.', ); @@ -239,7 +326,11 @@ describe('initializeMCPs', () => { await initializeMCPs(); expect(mockMCPManagerInstance.getAppToolFunctions).toHaveBeenCalledTimes(1); - expect(mockMergeAppTools).toHaveBeenCalledWith(mcpTools); + expect(mockMergeAppTools).toHaveBeenCalledWith(mcpTools, {}); + expect(mockMCPManagerInstance.connectAppServers).toHaveBeenCalledTimes(1); + expect(mockMergeAppTools.mock.invocationCallOrder[0]).toBeLessThan( + mockMCPManagerInstance.connectAppServers.mock.invocationCallOrder[0], + ); expect(logger.info).toHaveBeenCalledWith( '[MCP] Initialized with 1 configured server and 2 tools.', ); @@ -253,11 +344,21 @@ describe('initializeMCPs', () => { await initializeMCPs(); // Should use empty object fallback - expect(mockMergeAppTools).toHaveBeenCalledWith({}); + expect(mockMergeAppTools).toHaveBeenCalledWith({}, {}); expect(logger.info).toHaveBeenCalledWith( '[MCP] Initialized with 1 configured server and 0 tools.', ); }); + + it('should connect app servers when startup cache synchronization fails', async () => { + const mcpServers = { 'test-server': { type: 'sse', url: 'http://localhost:3001' } }; + mockGetAppConfig.mockResolvedValue({ mcpConfig: mcpServers }); + mockMergeAppTools.mockRejectedValueOnce(new Error('cache lock timed out')); + + await expect(initializeMCPs()).rejects.toThrow('cache lock timed out'); + + expect(mockMCPManagerInstance.connectAppServers).toHaveBeenCalledTimes(1); + }); }); describe('Initialization order', () => { @@ -315,3 +416,45 @@ describe('initializeMCPs', () => { }); }); }); + +describe('refreshChangedServerTools', () => { + const { refreshChangedServerTools } = require('./initializeMCPs'); + const event = { + serverName: 'dynamic', + serverConfig: { type: 'streamable-http', url: 'https://mcp.example.com' }, + tools: [{ name: 'tool', inputSchema: { type: 'object' } }], + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('publishes the complete refreshed snapshot in its original cache scope', async () => { + await refreshChangedServerTools({ ...event, userId: 'user-1' }); + + expect(mockUpdateMCPServerTools).toHaveBeenCalledWith({ ...event, userId: 'user-1' }); + expect(logger.info).toHaveBeenCalledWith( + '[MCP][dynamic] Tool list changed; refreshed 1 tool for user user-1', + ); + }); + + it('publishes an empty app-level snapshot so removals take effect', async () => { + await refreshChangedServerTools({ ...event, tools: [] }); + + expect(mockUpdateMCPServerTools).toHaveBeenCalledWith({ ...event, tools: [] }); + }); + + it('is registered as the tools-changed handler during initialization', async () => { + mockGetAppConfig.mockResolvedValue({ mcpConfig: null, mcpSettings: {} }); + + await initializeMCPs(); + + expect(mockSetMCPToolsChangedHandler).toHaveBeenCalledWith(refreshChangedServerTools); + expect(mockSetMCPToolsChangedGenerationHandler).toHaveBeenCalledWith( + mockGetMCPToolsCacheGeneration, + ); + expect(mockSetMCPToolsChangedGenerationRenewalHandler).toHaveBeenCalledWith( + mockRenewMCPToolsCacheGeneration, + ); + }); +}); diff --git a/e2e/config/librechat.e2e.yaml b/e2e/config/librechat.e2e.yaml index b8ebab2851..cf51932fb9 100644 --- a/e2e/config/librechat.e2e.yaml +++ b/e2e/config/librechat.e2e.yaml @@ -23,6 +23,7 @@ mcpSettings: # admin override is honored by inspection/connection. stdio servers skip this check. allowedDomains: - https://allowed.example.com + # __E2E_DYNAMIC_MCP_ALLOWED_DOMAIN__ mcpServers: e2e-memory: @@ -30,6 +31,7 @@ mcpServers: command: node args: - e2e/setup/fake-mcp-server.js + # __E2E_DYNAMIC_MCP_STDIO_ENV__ title: E2E Memory description: Local MCP fixture used by mock end-to-end tests. timeout: 30000 @@ -39,6 +41,7 @@ mcpServers: title: E2E HTTP description: Local HTTP MCP fixture for allowlist-override e2e tests. timeout: 30000 + # __E2E_DYNAMIC_MCP_NETWORK_SERVERS__ endpoints: # Default capabilities plus run_in_background (off by default upstream) so the diff --git a/e2e/playwright.config.mock.ts b/e2e/playwright.config.mock.ts index 9d40af1eed..9b1762c0de 100644 --- a/e2e/playwright.config.mock.ts +++ b/e2e/playwright.config.mock.ts @@ -4,10 +4,23 @@ import path from 'path'; import { getLocalE2EEnv, getE2EBaseURL } from './setup/env'; const rootPath = path.resolve(__dirname, '..'); -const serverPath = path.resolve(rootPath, 'e2e/setup/start-server.js'); +const replicaCount = Number(process.env.E2E_REPLICAS || '1'); +if (replicaCount !== 1 && replicaCount !== 2) { + throw new Error(`E2E_REPLICAS must be 1 or 2, received ${process.env.E2E_REPLICAS}`); +} +const serverPath = path.resolve( + rootPath, + replicaCount === 2 ? 'e2e/setup/start-server-cluster.js' : 'e2e/setup/start-server.js', +); const mcpHttpServerPath = path.resolve(rootPath, 'e2e/setup/fake-mcp-http-server.js'); +const dynamicMcpServerPath = path.resolve(rootPath, 'e2e/setup/fake-mcp-dynamic-network-server.js'); /** Must match the `e2e-http` server URL in e2e/config/librechat.e2e.yaml. */ const MCP_HTTP_PORT = process.env.E2E_MCP_HTTP_PORT || '8765'; +/** Must match the dynamic Streamable HTTP and SSE URLs in the e2e config template. */ +const MCP_DYNAMIC_PORT = process.env.E2E_MCP_DYNAMIC_PORT || '8766'; +const MCP_STATE_PATH = + process.env.E2E_MCP_STATE_PATH || + path.resolve(rootPath, 'e2e/specs/.test-results/mcp-tool-state.json'); const labelServerPath = path.resolve(rootPath, 'e2e/setup/fake-label-server.js'); /** The template's custom-endpoint `baseURL`s hard-code 8889; * `writeRuntimeMockConfig` substitutes any override into the generated copy. */ @@ -17,6 +30,7 @@ const configTemplatePath = path.resolve(rootPath, 'e2e/config/librechat.e2e.yaml const configPath = path.resolve(rootPath, 'e2e/.generated/librechat.e2e.yaml'); const reportPath = path.resolve(rootPath, 'e2e/playwright-report'); const deploymentSkillsPath = path.resolve(rootPath, 'e2e/fixtures/deployment-skills'); +const enableDynamicMcp = process.env.E2E_MCP_LIST_CHANGED === 'true'; const baseURL = getE2EBaseURL(); const chromiumChannel = process.env.E2E_CHROMIUM_CHANNEL || undefined; @@ -43,6 +57,7 @@ const baseEnv = { DEPLOYMENT_SKILLS_DIR: deploymentSkillsPath, /** Loaded in-process by `@librechat/api`'s `createRun` to swap in a fake model. */ LIBRECHAT_TEST_RUN_HOOK: fakeModelHookPath, + ...(enableDynamicMcp ? { E2E_MCP_LIST_CHANGED: 'true', E2E_MCP_STATE_PATH: MCP_STATE_PATH } : {}), ...vanillaOverrides, }; @@ -68,6 +83,34 @@ function writeRuntimeMockConfig() { process.env.E2E_MODEL_SPECS_ENFORCE === 'true' ? template.replace('\n enforce: false\n', '\n enforce: true\n') : template; + const dynamicMcpConfig = enableDynamicMcp + ? { + allowedDomain: '- http://127.0.0.1:8766', + stdioEnv: [ + 'env:', + ' E2E_MCP_LIST_CHANGED: "true"', + ` E2E_MCP_STATE_PATH: ${JSON.stringify(MCP_STATE_PATH)}`, + ].join('\n'), + networkServers: [ + 'e2e-streamable:', + ' type: streamable-http', + ' url: http://127.0.0.1:8766/mcp', + ' title: E2E Streamable HTTP', + ' description: Dynamic real-SDK Streamable HTTP fixture for mock end-to-end tests.', + ' timeout: 30000', + ' e2e-sse:', + ' type: sse', + ' url: http://127.0.0.1:8766/sse', + ' title: E2E SSE', + ' description: Dynamic real-SDK legacy SSE fixture for mock end-to-end tests.', + ' timeout: 30000', + ].join('\n'), + } + : { allowedDomain: '', stdioEnv: '', networkServers: '' }; + config = config + .replace('# __E2E_DYNAMIC_MCP_ALLOWED_DOMAIN__', dynamicMcpConfig.allowedDomain) + .replace('# __E2E_DYNAMIC_MCP_STDIO_ENV__', dynamicMcpConfig.stdioEnv) + .replace('# __E2E_DYNAMIC_MCP_NETWORK_SERVERS__', dynamicMcpConfig.networkServers); /** Keep the generated config in lockstep with the overridable label-server * port: the template hard-codes 8889, so an `E2E_LABEL_PORT` override that * moved only the server and its health check would report ready while @@ -75,8 +118,15 @@ function writeRuntimeMockConfig() { if (LABEL_PORT !== '8889') { config = config.split('127.0.0.1:8889').join(`127.0.0.1:${LABEL_PORT}`); } + if (enableDynamicMcp && MCP_DYNAMIC_PORT !== '8766') { + config = config.split('127.0.0.1:8766').join(`127.0.0.1:${MCP_DYNAMIC_PORT}`); + } fs.mkdirSync(path.dirname(configPath), { recursive: true }); fs.writeFileSync(configPath, config); + if (enableDynamicMcp) { + fs.mkdirSync(path.dirname(MCP_STATE_PATH), { recursive: true }); + fs.writeFileSync(MCP_STATE_PATH, `${JSON.stringify({ revision: 0, tool: null })}\n`); + } } function neutralizeCredentialEnv(env: NodeJS.ProcessEnv, keep: Set) { @@ -147,15 +197,6 @@ export default defineConfig({ }, ], webServer: [ - { - command: `node ${serverPath}`, - cwd: rootPath, - url: baseURL, - stdout: 'pipe', - ignoreHTTPSErrors: true, - timeout: 120_000, - reuseExistingServer: false, - }, { // URL-based MCP fixture for the allowlist-override spec (its health route is GET /). command: `node ${mcpHttpServerPath}`, @@ -166,6 +207,24 @@ export default defineConfig({ timeout: 60_000, reuseExistingServer: false, }, + ...(enableDynamicMcp + ? [ + { + // One real SDK server exposes both current HTTP and legacy SSE transports. + command: `node ${dynamicMcpServerPath}`, + cwd: rootPath, + env: { + ...process.env, + E2E_MCP_DYNAMIC_PORT: MCP_DYNAMIC_PORT, + E2E_MCP_STATE_PATH: MCP_STATE_PATH, + }, + url: `http://127.0.0.1:${MCP_DYNAMIC_PORT}/`, + stdout: 'pipe' as const, + timeout: 60_000, + reuseExistingServer: false, + }, + ] + : []), { // Serves the activity-label model call (the custom endpoints' baseURL). command: `node ${labelServerPath}`, @@ -176,5 +235,16 @@ export default defineConfig({ timeout: 60_000, reuseExistingServer: false, }, + { + // Start one LibreChat process, or a two-process topology behind a test-only proxy, after the + // network fixtures so inspection and persistent connections agree. + command: `node ${serverPath}`, + cwd: rootPath, + url: baseURL, + stdout: 'pipe', + ignoreHTTPSErrors: true, + timeout: 120_000, + reuseExistingServer: false, + }, ], }); diff --git a/e2e/setup/dynamic-mcp-tools.js b/e2e/setup/dynamic-mcp-tools.js new file mode 100644 index 0000000000..cc71cfd121 --- /dev/null +++ b/e2e/setup/dynamic-mcp-tools.js @@ -0,0 +1,130 @@ +const fs = require('node:fs'); +const path = require('node:path'); +const z = require('zod/v4'); + +const DEFAULT_STATE_PATH = path.join('/tmp', 'librechat-e2e-mcp-tool-state.json'); +const POLL_INTERVAL_MS = 50; + +function getStatePath() { + return process.env.E2E_MCP_STATE_PATH || DEFAULT_STATE_PATH; +} + +function emptyState() { + return { revision: 0, tool: null }; +} + +function readState() { + try { + const parsed = JSON.parse(fs.readFileSync(getStatePath(), 'utf8')); + if (typeof parsed.revision !== 'number') { + throw new Error('revision must be a number'); + } + if ( + parsed.tool !== null && + (typeof parsed.tool !== 'object' || typeof parsed.tool.description !== 'string') + ) { + throw new Error('tool must be null or contain a description'); + } + return parsed; + } catch (error) { + if (error?.code === 'ENOENT') { + return emptyState(); + } + throw error; + } +} + +function writeState(state) { + const statePath = getStatePath(); + const temporaryPath = `${statePath}.${process.pid}.tmp`; + fs.mkdirSync(path.dirname(statePath), { recursive: true }); + fs.writeFileSync(temporaryPath, `${JSON.stringify(state)}\n`); + fs.renameSync(temporaryPath, statePath); +} + +function resetState() { + writeState(emptyState()); +} + +function schemaForVersion(schemaVersion) { + if (schemaVersion === 2) { + return { + value: z.string(), + uppercase: z.boolean().optional(), + }; + } + return { value: z.string() }; +} + +function toolCallback({ value, uppercase = false }) { + const text = uppercase ? value.toUpperCase() : value; + return Promise.resolve({ content: [{ type: 'text', text }] }); +} + +/** + * Keeps one SDK McpServer's live registry synchronized with the shared e2e state file. + * registerTool/update/remove intentionally exercise the SDK's real list-changed notifications. + */ +function watchDynamicTool(server) { + let lastRevision = -1; + let registeredTool; + + const sync = () => { + const state = readState(); + if (state.revision === lastRevision) { + return; + } + lastRevision = state.revision; + + if (state.tool == null) { + registeredTool?.remove(); + registeredTool = undefined; + console.error( + `[dynamic-mcp-tools] applied revision ${state.revision}: removed runtime_probe`, + ); + return; + } + + const paramsSchema = schemaForVersion(state.tool.schemaVersion); + if (registeredTool) { + registeredTool.update({ + description: state.tool.description, + paramsSchema, + callback: toolCallback, + }); + console.error( + `[dynamic-mcp-tools] applied revision ${state.revision}: updated runtime_probe`, + ); + return; + } + + registeredTool = server.registerTool( + 'runtime_probe', + { + description: state.tool.description, + inputSchema: paramsSchema, + }, + toolCallback, + ); + console.error(`[dynamic-mcp-tools] applied revision ${state.revision}: added runtime_probe`); + }; + + sync(); + const timer = setInterval(() => { + try { + sync(); + } catch (error) { + console.error('[dynamic-mcp-tools] failed to synchronize tool state', error); + } + }, POLL_INTERVAL_MS); + + return () => clearInterval(timer); +} + +module.exports = { + emptyState, + getStatePath, + resetState, + watchDynamicTool, + writeState, +}; diff --git a/e2e/setup/env.ts b/e2e/setup/env.ts index da4fb3fb52..95d8217320 100644 --- a/e2e/setup/env.ts +++ b/e2e/setup/env.ts @@ -9,6 +9,9 @@ const GENERATED_CREDS_IV = crypto.randomBytes(16).toString('hex'); const GENERATED_JWT_SECRET = crypto.randomBytes(32).toString('hex'); const GENERATED_JWT_REFRESH_SECRET = crypto.randomBytes(32).toString('hex'); const DEFAULT_REDIS_URI = 'redis://127.0.0.1:6379/15'; +const DEFAULT_REDIS_CLUSTER_URI = [7001, 7002, 7003] + .map((port) => `redis://127.0.0.1:${port}`) + .join(','); const DEFAULT_REDIS_KEY_PREFIX = 'LibreChatE2E'; const PASSTHROUGH_ENV_KEYS = [ 'APPDATA', @@ -80,6 +83,7 @@ function getStreamStoreEnv(): Record { E2E_REQUIRE_REDIS_STREAMS: 'false', USE_REDIS: 'false', USE_REDIS_STREAMS: 'false', + USE_REDIS_CLUSTER: 'false', REDIS_KEY_PREFIX: '', REDIS_KEY_PREFIX_VAR: '', }; @@ -89,11 +93,23 @@ function getStreamStoreEnv(): Record { E2E_REQUIRE_REDIS_STREAMS: 'true', USE_REDIS: 'true', USE_REDIS_STREAMS: 'true', + USE_REDIS_CLUSTER: 'false', REDIS_URI: process.env.REDIS_URI ?? DEFAULT_REDIS_URI, REDIS_KEY_PREFIX: process.env.E2E_REDIS_KEY_PREFIX ?? DEFAULT_REDIS_KEY_PREFIX, REDIS_KEY_PREFIX_VAR: '', }; } + if (streamStore === 'redis-cluster') { + return { + E2E_REQUIRE_REDIS_STREAMS: 'true', + USE_REDIS: 'true', + USE_REDIS_STREAMS: 'true', + USE_REDIS_CLUSTER: 'true', + REDIS_URI: process.env.REDIS_URI ?? DEFAULT_REDIS_CLUSTER_URI, + REDIS_KEY_PREFIX: process.env.E2E_REDIS_KEY_PREFIX ?? DEFAULT_REDIS_KEY_PREFIX, + REDIS_KEY_PREFIX_VAR: '', + }; + } throw new Error(`Unsupported E2E_STREAM_STORE "${streamStore}"`); } diff --git a/e2e/setup/fake-mcp-dynamic-network-server.js b/e2e/setup/fake-mcp-dynamic-network-server.js new file mode 100644 index 0000000000..381fa91c09 --- /dev/null +++ b/e2e/setup/fake-mcp-dynamic-network-server.js @@ -0,0 +1,125 @@ +#!/usr/bin/env node + +const http = require('node:http'); +const { randomUUID } = require('node:crypto'); +const { McpServer } = require('@modelcontextprotocol/sdk/server/mcp.js'); +const { SSEServerTransport } = require('@modelcontextprotocol/sdk/server/sse.js'); +const { + StreamableHTTPServerTransport, +} = require('@modelcontextprotocol/sdk/server/streamableHttp.js'); +const { watchDynamicTool } = require('./dynamic-mcp-tools'); + +const PORT = Number.parseInt(process.env.E2E_MCP_DYNAMIC_PORT || '8766', 10); +const HOST = '127.0.0.1'; + +function createMcpServer(name, transportLabel) { + const server = new McpServer({ name, version: '1.0.0' }); + server.registerTool( + 'transport_probe', + { + description: `Confirms that the real ${transportLabel} MCP transport is connected.`, + inputSchema: {}, + }, + async () => ({ content: [{ type: 'text', text: `${transportLabel} connected` }] }), + ); + const stopWatching = watchDynamicTool(server); + return { server, stopWatching }; +} + +/** @type {Map, server: InstanceType, stopWatching: () => void }>} */ +const streamableSessions = new Map(); +/** @type {Map, server: InstanceType, stopWatching: () => void }>} */ +const sseSessions = new Map(); + +async function handleStreamableRequest(req, res) { + const sessionId = req.headers['mcp-session-id']; + let session = typeof sessionId === 'string' ? streamableSessions.get(sessionId) : undefined; + + if (!session) { + const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID() }); + const mcp = createMcpServer('e2e-streamable', 'Streamable HTTP'); + session = { transport, ...mcp }; + await mcp.server.connect(transport); + } + + await session.transport.handleRequest(req, res); + + const connectedSessionId = session.transport.sessionId; + if (connectedSessionId && !streamableSessions.has(connectedSessionId)) { + streamableSessions.set(connectedSessionId, session); + session.transport.onclose = () => { + streamableSessions.delete(connectedSessionId); + session.stopWatching(); + }; + } +} + +async function handleSSEConnect(res) { + const transport = new SSEServerTransport('/messages', res); + const mcp = createMcpServer('e2e-sse', 'SSE'); + const session = { transport, ...mcp }; + sseSessions.set(transport.sessionId, session); + transport.onclose = () => { + sseSessions.delete(transport.sessionId); + session.stopWatching(); + }; + await mcp.server.connect(transport); +} + +const httpServer = http.createServer(async (req, res) => { + try { + const url = new URL(req.url, `http://${req.headers.host}`); + if (req.method === 'GET' && url.pathname === '/') { + res.writeHead(200, { 'Content-Type': 'text/plain' }); + res.end('ok'); + return; + } + if (url.pathname === '/mcp') { + await handleStreamableRequest(req, res); + return; + } + if (req.method === 'GET' && url.pathname === '/sse') { + await handleSSEConnect(res); + return; + } + if (req.method === 'POST' && url.pathname === '/messages') { + const sessionId = url.searchParams.get('sessionId') || ''; + const session = sseSessions.get(sessionId); + if (!session) { + res.writeHead(404); + res.end(); + return; + } + await session.transport.handlePostMessage(req, res); + return; + } + res.writeHead(404); + res.end(); + } catch (error) { + console.error('[fake-mcp-dynamic-network-server] request failed', error); + if (!res.headersSent) { + res.writeHead(500); + } + res.end(); + } +}); + +async function shutdown() { + const sessions = [...streamableSessions.values(), ...sseSessions.values()]; + streamableSessions.clear(); + sseSessions.clear(); + await Promise.all( + sessions.map(async ({ server, stopWatching }) => { + stopWatching(); + await server.close().catch(() => undefined); + }), + ); + httpServer.close(() => process.exit(0)); +} + +process.once('SIGINT', shutdown); +process.once('SIGTERM', shutdown); + +httpServer.listen(PORT, HOST, () => { + console.log(`[e2e] dynamic MCP server listening on http://${HOST}:${PORT}`); +}); diff --git a/e2e/setup/fake-mcp-server.js b/e2e/setup/fake-mcp-server.js index 03c2602147..77454aa5f7 100644 --- a/e2e/setup/fake-mcp-server.js +++ b/e2e/setup/fake-mcp-server.js @@ -5,6 +5,7 @@ const path = require('node:path'); const { McpServer } = require('@modelcontextprotocol/sdk/server/mcp.js'); const { StdioServerTransport } = require('@modelcontextprotocol/sdk/server/stdio.js'); const z = require('zod/v4'); +const { watchDynamicTool } = require('./dynamic-mcp-tools'); const APPROVAL_AUDIT_DIR = path.join('/tmp', 'librechat-e2e-approval-audit'); @@ -99,6 +100,19 @@ server.registerTool( }, ); +if (process.env.E2E_MCP_LIST_CHANGED === 'true') { + server.registerTool( + 'transport_probe', + { + description: 'Confirms that the real stdio MCP transport is connected.', + inputSchema: {}, + }, + async () => ({ content: [{ type: 'text', text: 'stdio connected' }] }), + ); + + watchDynamicTool(server); +} + async function main() { await server.connect(new StdioServerTransport()); } diff --git a/e2e/setup/start-server-cluster.js b/e2e/setup/start-server-cluster.js new file mode 100644 index 0000000000..4cf5bf9ece --- /dev/null +++ b/e2e/setup/start-server-cluster.js @@ -0,0 +1,165 @@ +const { spawn } = require('child_process'); +const fs = require('fs'); +const http = require('http'); +const path = require('path'); +const { MongoMemoryServer } = require('mongodb-memory-server'); + +const DEFAULT_BASE_URL = 'http://localhost:3080'; +const DEFAULT_RUNTIME_ENV_PATH = path.resolve(__dirname, '../specs/.test-results/runtime-env.json'); +const REPLICA_STARTUP_TIMEOUT_MS = 120_000; +const serverPath = path.resolve(__dirname, 'start-server.js'); + +let shuttingDown = false; +let mongoServer; +let proxyServer; +const children = []; + +function getTopology() { + const baseURL = new URL(process.env.E2E_BASE_URL || DEFAULT_BASE_URL); + if (baseURL.protocol !== 'http:') { + throw new Error(`[e2e] Replica proxy requires an http base URL, received ${baseURL.protocol}`); + } + const basePort = Number(baseURL.port || 80); + if (!Number.isInteger(basePort) || basePort < 1 || basePort > 65533) { + throw new Error(`[e2e] Invalid replica base port: ${baseURL.port}`); + } + return { + baseURL, + replicaPorts: [basePort + 1, basePort + 2], + }; +} + +function writeRuntimeEnv(mongoUri) { + const runtimeEnvPath = process.env.E2E_RUNTIME_ENV_PATH || DEFAULT_RUNTIME_ENV_PATH; + fs.mkdirSync(path.dirname(runtimeEnvPath), { recursive: true }); + fs.writeFileSync(runtimeEnvPath, JSON.stringify({ MONGO_URI: mongoUri }, null, 2)); +} + +function startReplica(port, index, mongoUri) { + const child = spawn(process.execPath, [serverPath], { + cwd: path.resolve(__dirname, '../..'), + env: { + ...process.env, + E2E_REPLICA_INDEX: String(index), + E2E_USE_MEMORY_MONGO: 'false', + HOST: process.env.E2E_HOST || '127.0.0.1', + MONGO_URI: mongoUri, + PORT: String(port), + }, + stdio: 'inherit', + }); + children.push(child); + child.once('exit', (code, signal) => { + if (!shuttingDown) { + console.error( + `[e2e] LibreChat replica ${index} exited unexpectedly (${signal || `code ${code}`})`, + ); + void shutdown(code || 1); + } + }); + return child; +} + +async function waitForReplica(port) { + const deadline = Date.now() + REPLICA_STARTUP_TIMEOUT_MS; + while (Date.now() < deadline) { + const isReady = await new Promise((resolve) => { + const request = http.get(`http://127.0.0.1:${port}/readyz`, (response) => { + response.resume(); + resolve(response.statusCode === 200); + }); + request.setTimeout(500, () => request.destroy()); + request.once('error', () => resolve(false)); + }); + if (isReady) { + return; + } + await new Promise((resolve) => setTimeout(resolve, 100)); + } + throw new Error(`[e2e] LibreChat replica on port ${port} did not become ready`); +} + +function startProxy(baseURL, targetPort) { + proxyServer = http.createServer((request, response) => { + const upstream = http.request( + { + hostname: '127.0.0.1', + port: targetPort, + path: request.url, + method: request.method, + headers: request.headers, + }, + (upstreamResponse) => { + response.writeHead(upstreamResponse.statusCode || 502, upstreamResponse.headers); + upstreamResponse.pipe(response); + }, + ); + upstream.once('error', (error) => { + if (!response.headersSent) { + response.writeHead(502, { 'content-type': 'text/plain' }); + } + response.end(`Replica unavailable: ${error.message}`); + }); + request.pipe(upstream); + }); + proxyServer.listen(Number(baseURL.port || 80), baseURL.hostname, () => { + console.log( + `[e2e] Replica proxy listening at ${baseURL.origin}; primary target is ${targetPort}`, + ); + }); +} + +async function shutdown(exitCode = 0) { + if (shuttingDown) { + return; + } + shuttingDown = true; + if (proxyServer) { + proxyServer.close(); + } + for (const child of children) { + child.kill('SIGTERM'); + } + await Promise.all( + children.map( + (child) => + new Promise((resolve) => { + if (child.exitCode != null || child.signalCode != null) { + resolve(); + return; + } + child.once('exit', resolve); + }), + ), + ); + if (mongoServer) { + await mongoServer.stop(); + } + process.exit(exitCode); +} + +async function startCluster() { + const { baseURL, replicaPorts } = getTopology(); + mongoServer = await MongoMemoryServer.create({ + instance: { + dbName: 'LibreChat-e2e', + ip: '127.0.0.1', + }, + }); + const mongoUri = new URL('LibreChat-e2e', mongoServer.getUri()).toString(); + writeRuntimeEnv(mongoUri); + console.log(`[e2e] Started shared memory MongoDB at ${mongoUri}`); + startReplica(replicaPorts[0], 1, mongoUri); + await waitForReplica(replicaPorts[0]); + startReplica(replicaPorts[1], 2, mongoUri); + await waitForReplica(replicaPorts[1]); + startProxy(baseURL, replicaPorts[0]); +} + +process.once('SIGINT', () => void shutdown(130)); +process.once('SIGTERM', () => void shutdown(143)); + +startCluster().catch((error) => { + console.error('[e2e] Failed to start LibreChat replicas:', error); + void shutdown(1); +}); diff --git a/e2e/specs/mock/mcp-fixture-isolation.spec.ts b/e2e/specs/mock/mcp-fixture-isolation.spec.ts new file mode 100644 index 0000000000..47a8f3857c --- /dev/null +++ b/e2e/specs/mock/mcp-fixture-isolation.spec.ts @@ -0,0 +1,17 @@ +import fs from 'node:fs'; +import path from 'node:path'; +import { expect, test } from '@playwright/test'; + +test('keeps dynamic MCP fixtures out of the general mock suite', () => { + expect(process.env.E2E_MCP_LIST_CHANGED).not.toBe('true'); + + const config = fs.readFileSync( + path.resolve(__dirname, '../../.generated/librechat.e2e.yaml'), + 'utf8', + ); + expect(config).not.toContain('e2e-streamable:'); + expect(config).not.toContain('e2e-sse:'); + expect(config).not.toContain('E2E_MCP_LIST_CHANGED:'); + expect(config).not.toContain('E2E_MCP_STATE_PATH:'); + expect(config).not.toContain('127.0.0.1:8766'); +}); diff --git a/e2e/specs/mock/mcp-tool-list-changed.spec.ts b/e2e/specs/mock/mcp-tool-list-changed.spec.ts new file mode 100644 index 0000000000..8c2f97a2a7 --- /dev/null +++ b/e2e/specs/mock/mcp-tool-list-changed.spec.ts @@ -0,0 +1,161 @@ +import fs from 'node:fs'; +import path from 'node:path'; +import { expect, test } from '@playwright/test'; +import type { Page } from '@playwright/test'; +import { getAccessToken } from './helpers'; + +const STATE_PATH = + process.env.E2E_MCP_STATE_PATH || path.resolve(__dirname, '../.test-results/mcp-tool-state.json'); +const SERVERS = ['e2e-memory', 'e2e-streamable', 'e2e-sse'] as const; +const TRANSPORT_PROBE = 'transport_probe'; +const DYNAMIC_TOOL = 'runtime_probe'; +const REPLICA_COUNT = Number(process.env.E2E_REPLICAS || '1'); + +type MCPTool = { + name: string; + pluginKey: string; + description?: string; +}; + +type MCPToolsResponse = { + servers?: Record; +}; + +function writeToolState( + revision: number, + tool: { description: string; schemaVersion: number } | null, +) { + const temporaryPath = `${STATE_PATH}.${process.pid}.tmp`; + fs.mkdirSync(path.dirname(STATE_PATH), { recursive: true }); + fs.writeFileSync(temporaryPath, `${JSON.stringify({ revision, tool })}\n`); + fs.renameSync(temporaryPath, STATE_PATH); +} + +function getCatalogURLs() { + if (REPLICA_COUNT === 1) { + return ['/api/mcp/tools']; + } + const baseURL = new URL(process.env.E2E_BASE_URL || 'http://localhost:3080'); + const basePort = Number(baseURL.port || 80); + return [1, 2].map((offset) => { + const replicaURL = new URL('/api/mcp/tools', baseURL); + replicaURL.port = String(basePort + offset); + return replicaURL.toString(); + }); +} + +async function getTools( + page: Page, + accessToken: string, + catalogURL: string, +): Promise { + const response = await page.request.get(catalogURL, { + headers: { Authorization: `Bearer ${accessToken}` }, + }); + expect(response.ok(), await response.text()).toBe(true); + return response.json() as Promise; +} + +async function expectCatalog( + page: Page, + accessToken: string, + expected: { present: boolean; description?: string; toolName: string }, +) { + await expect + .poll( + async () => { + const catalogs = await Promise.all( + getCatalogURLs().map((catalogURL) => getTools(page, accessToken, catalogURL)), + ); + return catalogs.flatMap((catalog, replicaIndex) => + SERVERS.map((serverName) => { + const pluginKey = `${expected.toolName}_mcp_${serverName}`; + const tool = catalog.servers?.[serverName]?.tools?.find( + (candidate) => candidate.pluginKey === pluginKey, + ); + if (expected.description === undefined) { + return { replica: replicaIndex + 1, serverName, present: tool != null }; + } + return tool + ? { replica: replicaIndex + 1, serverName, description: tool.description } + : { replica: replicaIndex + 1, serverName }; + }), + ); + }, + { + message: `${expected.toolName} should be ${expected.present ? 'present' : 'absent'} on every MCP transport`, + timeout: 30_000, + intervals: [100, 250, 500], + }, + ) + .toEqual( + getCatalogURLs().flatMap((_, replicaIndex) => + SERVERS.map((serverName) => { + if (expected.description === undefined) { + return { replica: replicaIndex + 1, serverName, present: expected.present }; + } + if (expected.present) { + return { + replica: replicaIndex + 1, + serverName, + description: expected.description, + }; + } + return { replica: replicaIndex + 1, serverName }; + }), + ), + ); +} + +test.describe('MCP tools/list_changed transports', () => { + test.skip( + process.env.E2E_MCP_LIST_CHANGED !== 'true', + 'runs only in the dedicated dynamic MCP topology matrix', + ); + + test('refreshes the public tool catalog over stdio, Streamable HTTP, and SSE', async ({ + page, + }) => { + test.setTimeout(120_000); + await page.goto('/c/new'); + const accessToken = await getAccessToken(page); + const revisionBase = Date.now(); + + try { + writeToolState(revisionBase, null); + await expectCatalog(page, accessToken, { + present: true, + toolName: TRANSPORT_PROBE, + description: undefined, + }); + + writeToolState(revisionBase + 1, { + description: 'Dynamic MCP tool version one', + schemaVersion: 1, + }); + await expectCatalog(page, accessToken, { + present: true, + toolName: DYNAMIC_TOOL, + description: 'Dynamic MCP tool version one', + }); + + writeToolState(revisionBase + 2, { + description: 'Dynamic MCP tool version two', + schemaVersion: 2, + }); + await expectCatalog(page, accessToken, { + present: true, + toolName: DYNAMIC_TOOL, + description: 'Dynamic MCP tool version two', + }); + + writeToolState(revisionBase + 3, null); + await expectCatalog(page, accessToken, { + present: false, + toolName: DYNAMIC_TOOL, + }); + } finally { + writeToolState(revisionBase + 4, null); + } + }); +}); diff --git a/packages/api/src/agents/__tests__/load.spec.ts b/packages/api/src/agents/__tests__/load.spec.ts index 737642e6ac..83aa9cda3c 100644 --- a/packages/api/src/agents/__tests__/load.spec.ts +++ b/packages/api/src/agents/__tests__/load.spec.ts @@ -162,11 +162,37 @@ describe('loadAgent', () => { ); expect(mockGetMCPServerTools).toHaveBeenCalledTimes(1); - expect(mockGetMCPServerTools).toHaveBeenCalledWith('user123', 'server1'); + expect(mockGetMCPServerTools).toHaveBeenCalledWith('user123', 'server1', undefined); expect(result?.tools).toContain(`${Constants.mcp_all}${Constants.mcp_delimiter}body-scoped`); expect(result?.tools).toContain('tool1_mcp_server1'); }); + test('addresses cached tools with a non-ephemeral request overlay', async () => { + const { EPHEMERAL_AGENT_ID } = Constants; + const overlayConfig = { + type: 'streamable-http' as const, + url: 'https://overlay.example.com/mcp', + }; + mockGetMCPServerTools.mockResolvedValue({ overlay_tool_mcp_overlay: {} }); + + const result = await loadAgent( + { + req: { + user: { id: 'user123' }, + config: { mcpConfig: { overlay: overlayConfig } } as unknown as AppConfig, + body: { ephemeralAgent: { mcp: ['overlay'] } }, + }, + agent_id: EPHEMERAL_AGENT_ID as string, + endpoint: 'openai', + model_parameters: { model: 'gpt-4' } as unknown as AgentModelParameters, + }, + deps, + ); + + expect(mockGetMCPServerTools).toHaveBeenCalledWith('user123', 'overlay', overlayConfig); + expect(result?.tools).toContain('overlay_tool_mcp_overlay'); + }); + test('should return null for non-existent agent', async () => { const mockReq = { user: { id: 'user123' } }; const result = await loadAgent( @@ -671,6 +697,37 @@ describe('loadAgent', () => { expect(result?.subagents).toBeUndefined(); }); + test('addresses added-agent cached tools with the effective config overlay', async () => { + const overlayConfig = { + type: 'streamable-http' as const, + url: 'https://overlay.example.com/mcp', + }; + mockGetMCPServerTools.mockResolvedValue({ overlay_tool_mcp_overlay: {} }); + + const result = await loadAddedAgent( + { + req: { + user: { id: 'user123' }, + config: { + config: {}, + fileStrategy: FileSources.local, + imageOutputType: 'png', + mcpConfig: { overlay: overlayConfig }, + }, + }, + conversation: { + endpoint: 'openai', + model: 'gpt-4', + ephemeralAgent: { mcp: ['overlay'] }, + } as unknown as TConversation, + }, + deps, + ); + + expect(mockGetMCPServerTools).toHaveBeenCalledWith('user123', 'overlay', overlayConfig); + expect(result?.tools).toContain('overlay_tool_mcp_overlay'); + }); + test('should enable full skill scope for added ephemeral model spec with skills true', async () => { const result = await loadAddedAgent( { diff --git a/packages/api/src/agents/added.ts b/packages/api/src/agents/added.ts index b2b80a437a..0d23204500 100644 --- a/packages/api/src/agents/added.ts +++ b/packages/api/src/agents/added.ts @@ -9,11 +9,12 @@ import { } from 'librechat-data-provider'; import type { Agent, AgentToolOptions, TConversation, TModelSpec } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; +import type { ParsedServerConfig } from '~/mcp/types'; +import { requiresEphemeralUserConnection, validateMCPServerConfig } from '~/mcp/utils'; import { ASK_USER_QUESTION_TOOL_NAME } from '~/agents/hitl/askUserQuestionTool'; import { synthesizeBackgroundToolOptions } from '~/agents/background'; import { mergeSynthesizedToolOptions } from '~/agents/selection'; import { synthesizeIntentToolOptions } from '~/agents/intent'; -import { requiresEphemeralUserConnection } from '~/mcp/utils'; import { getCustomEndpointConfig } from '~/app/config'; const { mcp_all, mcp_delimiter } = Constants; @@ -53,6 +54,7 @@ export interface LoadAddedAgentDeps { getMCPServerTools: ( userId: string, serverName: string, + serverConfig?: ParsedServerConfig, ) => Promise | null>; } @@ -218,13 +220,14 @@ export async function loadAddedAgent( if (addedServers.has(mcpServer)) { continue; } - /** Request-tier overlays are invisible to the cache service's registry - * resolver — overlay-scoped servers expand fresh via `mcp_all` instead */ - const overlayConfig = appConfig?.mcpConfig?.[mcpServer]; + /** Address durable catalogs by the effective request overlay; request-scoped + * overlays still expand fresh through `mcp_all`. */ + const rawOverlayConfig = appConfig?.mcpConfig?.[mcpServer]; + const overlayConfig = rawOverlayConfig ? validateMCPServerConfig(rawOverlayConfig) : undefined; const serverTools = overlayConfig && requiresEphemeralUserConnection(overlayConfig) ? null - : await deps.getMCPServerTools(userId, mcpServer); + : await deps.getMCPServerTools(userId, mcpServer, overlayConfig); if (!serverTools) { tools.push(`${mcp_all}${mcp_delimiter}${mcpServer}`); addedServers.add(mcpServer); diff --git a/packages/api/src/agents/load.ts b/packages/api/src/agents/load.ts index b4a73a8f54..a160319f62 100644 --- a/packages/api/src/agents/load.ts +++ b/packages/api/src/agents/load.ts @@ -14,11 +14,12 @@ import type { Agent, } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; +import type { ParsedServerConfig } from '~/mcp/types'; +import { requiresEphemeralUserConnection, validateMCPServerConfig } from '~/mcp/utils'; import { ASK_USER_QUESTION_TOOL_NAME } from '~/agents/hitl/askUserQuestionTool'; import { synthesizeBackgroundToolOptions } from '~/agents/background'; import { mergeSynthesizedToolOptions } from '~/agents/selection'; import { synthesizeIntentToolOptions } from '~/agents/intent'; -import { requiresEphemeralUserConnection } from '~/mcp/utils'; import { getCustomEndpointConfig } from '~/app/config'; const { mcp_all, mcp_delimiter } = Constants; @@ -29,6 +30,7 @@ export interface LoadAgentDeps { getMCPServerTools: ( userId: string, serverName: string, + serverConfig?: ParsedServerConfig, ) => Promise | null>; } @@ -94,13 +96,16 @@ export async function loadEphemeralAgent( if (addedServers.has(mcpServer)) { continue; } - /** Request-tier overlays are invisible to the cache service's registry - * resolver — overlay-scoped servers expand fresh via `mcp_all` instead */ - const overlayConfig = req.config?.mcpConfig?.[mcpServer]; + /** Address durable catalogs by the effective request overlay; request-scoped + * overlays still expand fresh through `mcp_all`. */ + const rawOverlayConfig = req.config?.mcpConfig?.[mcpServer]; + const overlayConfig = rawOverlayConfig + ? validateMCPServerConfig(rawOverlayConfig) + : undefined; const serverTools = overlayConfig && requiresEphemeralUserConnection(overlayConfig) ? null - : await deps.getMCPServerTools(userId, mcpServer); + : await deps.getMCPServerTools(userId, mcpServer, overlayConfig); if (!serverTools) { tools.push(`${mcp_all}${mcp_delimiter}${mcpServer}`); addedServers.add(mcpServer); diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index ec61eee882..fad99cd525 100644 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -15,12 +15,15 @@ export * from './mcp/authority'; export * from './mcp/registry/MCPServersRegistry'; export * from './mcp/MCPManager'; export * from './mcp/connection'; +export * from './mcp/toolsChanged'; export * from './mcp/oauth'; export * from './mcp/auth'; export * from './mcp/zod'; export * from './mcp/errors'; export * from './mcp/cache'; export * from './mcp/tools'; +export * from './mcp/catalog/store'; +export * from './mcp/assistants'; export * from './mcp/request'; /* Utilities */ export * from './mcp/utils'; diff --git a/packages/api/src/mcp/ConnectionsRepository.ts b/packages/api/src/mcp/ConnectionsRepository.ts index 7dd2695467..63e9396192 100644 --- a/packages/api/src/mcp/ConnectionsRepository.ts +++ b/packages/api/src/mcp/ConnectionsRepository.ts @@ -1,12 +1,22 @@ import { logger } from '@librechat/data-schemas'; import type * as t from './types'; +import { + cancelMCPToolsChanged, + getMCPAppToolsPublicationGeneration, + notifyMCPToolsChanged, +} from './toolsChanged'; import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry'; -import { isUserSourced, requiresUserScopedConnection } from './utils'; import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory'; +import { canUseAppConnection, isUserSourced } from './utils'; import { MCPConnection } from './connection'; const CONNECT_CONCURRENCY = 3; +interface ConnectionLoadOptions { + continueOnError?: boolean; + refreshTools?: boolean; +} + /** * Manages MCP connections with lazy loading and reconnection. * Maintains a pool of connections and handles connection lifecycle management. @@ -20,6 +30,8 @@ export class ConnectionsRepository { protected connections: Map = new Map(); protected oauthOpts: t.OAuthConnectionOptions | undefined; private readonly ownerId: string | undefined; + private readonly connectionOperations = new Map>(); + private shuttingDown = false; constructor(ownerId?: string, oauthOpts?: t.OAuthConnectionOptions) { this.ownerId = ownerId; @@ -31,6 +43,23 @@ export class ConnectionsRepository { return this.connections.size; } + /** Serializes connection lifecycle transitions for one server without blocking other servers. */ + private runConnectionOperation(serverName: string, operation: () => Promise): Promise { + const previous = this.connectionOperations.get(serverName) ?? Promise.resolve(); + const result = previous.then(operation, operation); + const tail = result.then( + () => undefined, + () => undefined, + ); + this.connectionOperations.set(serverName, tail); + void tail.then(() => { + if (this.connectionOperations.get(serverName) === tail) { + this.connectionOperations.delete(serverName); + } + }); + return result; + } + /** Checks whether this repository can connect to a specific server */ async has(serverName: string): Promise { const config = await MCPServersRegistry.getInstance().getServerConfig(serverName, this.ownerId); @@ -43,7 +72,23 @@ export class ConnectionsRepository { } /** Gets or creates a connection for the specified server with lazy loading */ - async get(serverName: string): Promise { + async get( + serverName: string, + options: ConnectionLoadOptions = {}, + ): Promise { + if (this.shuttingDown) { + return null; + } + return this.runConnectionOperation(serverName, () => this.loadConnection(serverName, options)); + } + + private async loadConnection( + serverName: string, + options: ConnectionLoadOptions, + ): Promise { + if (this.shuttingDown) { + return null; + } const serverConfig = await MCPServersRegistry.getInstance().getServerConfig( serverName, this.ownerId, @@ -51,9 +96,7 @@ export class ConnectionsRepository { const existingConnection = this.connections.get(serverName); if (!serverConfig || !this.isAllowedToConnectToServer(serverConfig)) { - if (existingConnection) { - await existingConnection.disconnect(); - } + await this.disconnectConnection(serverName); return null; } if (existingConnection) { @@ -68,18 +111,19 @@ export class ConnectionsRepository { ); // Disconnect stale connection - await existingConnection.disconnect(); - this.connections.delete(serverName); + await this.disconnectConnection(serverName); // Fall through to create new connection } else if (await existingConnection.isConnected()) { return existingConnection; } else { - await this.disconnect(serverName); + await this.disconnectConnection(serverName); } } const registry = MCPServersRegistry.getInstance(); const { allowedDomains, allowedAddresses, useSSRFProtection } = await registry.resolveAllowlists({ userId: this.ownerId }); + const publicationGeneration = + this.ownerId === undefined ? getMCPAppToolsPublicationGeneration(serverConfig) : undefined; const connection = await MCPConnectionFactory.create( { serverName, @@ -92,19 +136,81 @@ export class ConnectionsRepository { this.oauthOpts, ); + if (this.shuttingDown) { + await connection.dispose(); + await cancelMCPToolsChanged({ userId: this.ownerId, serverName }); + return null; + } + + let toolsChangedGeneration = 0; + let latestToolsChangedPublication = Promise.resolve(); + + /* Both scopes get the same treatment: this repository is per-owner, so ownerId already says + * whose tool cache a change belongs to (undefined = the app-level, shared one). */ + connection.on('toolsChanged', (tools: t.MCPTool[], publicationRevision?: string) => { + toolsChangedGeneration++; + latestToolsChangedPublication = notifyMCPToolsChanged({ + tools, + serverName, + serverConfig, + userId: this.ownerId, + publicationGeneration, + publicationRevision, + }); + void latestToolsChangedPublication; + }); + this.connections.set(serverName, connection); + if (this.ownerId === undefined && options.refreshTools !== false) { + if (connection.client.getServerCapabilities()?.tools == null) { + await notifyMCPToolsChanged({ + tools: [], + serverName, + serverConfig, + publicationGeneration, + }); + return connection; + } + const initialGeneration = toolsChangedGeneration; + const snapshot = await connection.fetchToolsSnapshot(); + if (snapshot.complete) { + if (toolsChangedGeneration !== initialGeneration) { + await latestToolsChangedPublication; + } else { + await notifyMCPToolsChanged({ + tools: snapshot.tools, + serverName, + serverConfig, + publicationGeneration, + }); + } + } else { + await connection.refreshToolList(); + } + } return connection; } /** Gets or creates connections for multiple servers concurrently */ - async getMany(serverNames: string[]): Promise> { + async getMany( + serverNames: string[], + options: ConnectionLoadOptions = {}, + ): Promise> { const results: [string, MCPConnection | null][] = []; for (let i = 0; i < serverNames.length; i += CONNECT_CONCURRENCY) { const batch = serverNames.slice(i, i + CONNECT_CONCURRENCY); const batchResults = await Promise.all( - batch.map( - async (name): Promise<[string, MCPConnection | null]> => [name, await this.get(name)], - ), + batch.map(async (name): Promise<[string, MCPConnection | null]> => { + try { + return [name, await this.get(name, options)]; + } catch (error) { + if (!options.continueOnError) { + throw error; + } + logger.warn(`${this.prefix(name)} Failed to establish connection`, error); + return [name, null]; + } + }), ); results.push(...batchResults); } @@ -117,27 +223,45 @@ export class ConnectionsRepository { } /** Gets or creates connections for all configured servers in this repository's scope */ - async getAll(): Promise> { + async getAll(options: ConnectionLoadOptions = {}): Promise> { //TODO in the future we should use a scoped config getter (APPLevel, UserLevel, Private) //for now the absent config will not throw error const allConfigs = await MCPServersRegistry.getInstance().getAllServerConfigs(this.ownerId); - return this.getMany(Object.keys(allConfigs)); + return this.getMany(Object.keys(allConfigs), options); } /** Disconnects and removes a specific server connection from the pool */ async disconnect(serverName: string): Promise { + return this.runConnectionOperation(serverName, () => this.disconnectConnection(serverName)); + } + + private async disconnectConnection(serverName: string): Promise { const connection = this.connections.get(serverName); - if (!connection) return Promise.resolve(); + if (!connection) { + await cancelMCPToolsChanged({ userId: this.ownerId, serverName }); + return; + } this.connections.delete(serverName); - return connection.disconnect().catch((err) => { - logger.error(`${this.prefix(serverName)} Error disconnecting`, err); - }); + try { + connection.removeAllListeners?.('toolsChanged'); + await connection.dispose(); + } catch (err) { + logger.error(`${this.prefix(serverName)} Error disposing`, err); + } finally { + await cancelMCPToolsChanged({ userId: this.ownerId, serverName }); + } } /** Disconnects all active connections and returns array of disconnect promises */ disconnectAll(): Promise[] { + this.shuttingDown = true; + return [this.drainAndDisconnectAll()]; + } + + private async drainAndDisconnectAll(): Promise { + await Promise.allSettled(Array.from(this.connectionOperations.values())); const serverNames = Array.from(this.connections.keys()); - return serverNames.map((serverName) => this.disconnect(serverName)); + await Promise.all(serverNames.map((serverName) => this.disconnect(serverName))); } // Returns formatted log prefix for server messages @@ -156,10 +280,7 @@ export class ConnectionsRepository { if (config.inspectionFailed) { return false; } - if ( - this.ownerId === undefined && - (config.startup === false || requiresUserScopedConnection(config)) - ) { + if (this.ownerId === undefined && !canUseAppConnection(config)) { return false; } return true; diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index 34d42f55a7..977978d9d3 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -204,9 +204,14 @@ export class MCPConnectionFactory { ); if (await connection.isConnected()) { - const tools = await connection.fetchTools(); + const snapshot = await connection.fetchOrderedToolsSnapshot(); connection.removeListener('oauthRequired', oauthHandler); - return { tools, connection, oauthRequired: false, oauthUrl: null }; + return { + tools: snapshot.complete ? snapshot.tools : null, + connection, + oauthRequired: false, + oauthUrl: null, + }; } } catch { MCPConnection.decrementCycleCount(this.serverName); @@ -223,7 +228,7 @@ export class MCPConnectionFactory { `${this.logPrefix} [Discovery] Successfully discovered ${tools.length} tools without auth`, ); try { - await connection.disconnect(); + await connection.dispose(); } catch { // Ignore cleanup errors } @@ -238,7 +243,7 @@ export class MCPConnectionFactory { connection.removeListener('oauthRequired', oauthHandler); try { - await connection.disconnect(); + await connection.dispose(); } catch { // Ignore cleanup errors } @@ -272,16 +277,16 @@ export class MCPConnectionFactory { await withTimeout(unauthConnection.connect(), connectTimeout, `Unauth connection timeout`); if (await unauthConnection.isConnected()) { - const tools = await unauthConnection.fetchTools(); - await unauthConnection.disconnect(); - return tools; + const snapshot = await unauthConnection.fetchOrderedToolsSnapshot(); + await unauthConnection.dispose(); + return snapshot.complete ? snapshot.tools : null; } } catch { logger.debug(`${this.logPrefix} [Discovery] Unauthenticated connection attempt failed`); } try { - await unauthConnection.disconnect(); + await unauthConnection.dispose(); } catch { // Ignore cleanup errors } diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index 6d39acfd68..bb43719843 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -12,12 +12,14 @@ import type { RequestBody } from '~/types'; import type * as t from './types'; import { getMissingRuntimeBodyPlaceholderFields, + canUseAppConnection, isOAuthServer, isUserSourced, requiresEphemeralUserConnection, requiresOAuthMachinery, requiresUserScopedConnection, } from './utils'; +import { getMCPAppToolsPublicationGeneration, getMCPToolsChangedGeneration } from './toolsChanged'; import { MCPServersInitializer } from './registry/MCPServersInitializer'; import { OboTokenResolutionError, resolveOboToken } from '~/mcp/oauth'; import { MCPServerInspector } from './registry/MCPServerInspector'; @@ -127,8 +129,12 @@ export class MCPManager extends UserConnectionManager { try { const existingAppConnection = await this.appConnections?.get(serverName); if (existingAppConnection && (await existingAppConnection.isConnected())) { - const tools = await existingAppConnection.fetchTools(); - return { tools, oauthRequired: false, oauthUrl: null }; + const snapshot = await existingAppConnection.fetchOrderedToolsSnapshot(); + return { + tools: snapshot.complete ? snapshot.tools : null, + oauthRequired: false, + oauthUrl: null, + }; } } catch { logger.debug(`${logPrefix} [Discovery] App connection not available, trying discovery mode`); @@ -186,9 +192,9 @@ export class MCPManager extends UserConnectionManager { ): Promise => { if (result.connection) { try { - await result.connection.disconnect(); + await result.connection.dispose(); } catch (error) { - logger.warn(`${logPrefix} [Discovery] Failed to disconnect discovery connection`, error); + logger.warn(`${logPrefix} [Discovery] Failed to dispose discovery connection`, error); } } return { @@ -237,43 +243,129 @@ export class MCPManager extends UserConnectionManager { const toolFunctions: t.LCAvailableTools = {}; const configs = await MCPServersRegistry.getInstance().getAllServerConfigs(); for (const config of Object.values(configs)) { - if (config.toolFunctions != null) { + if (canUseAppConnection(config) && config.toolFunctions != null) { Object.assign(toolFunctions, config.toolFunctions); } } return toolFunctions; } - /** Returns all available tool functions from all connections available to user */ - public async getServerToolFunctions( + /** Opens eligible app-shared sessions after the inspected startup catalog has been cached. */ + public async connectAppServers(): Promise { + try { + const configs = await MCPServersRegistry.getInstance().getAllServerConfigs(); + const serverNames = Object.entries(configs) + .filter(([, config]) => canUseAppConnection(config)) + .map(([serverName]) => serverName); + const connections = await this.appConnections?.getMany(serverNames, { + continueOnError: true, + refreshTools: false, + }); + if (!connections) { + return; + } + await Promise.all( + Array.from(connections.values(), (connection) => connection.refreshToolList()), + ); + } catch (error) { + logger.warn('[MCP] Failed to establish one or more app connections after inspection', error); + } + } + + /** Closes app-shared MCP sessions during graceful process shutdown. */ + public async disconnectAppServers(): Promise { + await Promise.all(this.appConnections?.disconnectAll() ?? []); + } + + /** Returns tool functions with the generation bound to their originating user connection. */ + public async getServerToolFunctionsSnapshot( userId: string, serverName: string, - ): Promise { + serverConfig?: t.ParsedServerConfig, + ): Promise<{ + tools: t.LCAvailableTools | null; + publicationGeneration?: string; + }> { try { - //try get the appConnection (if the config is not in the app level anymore any existing connection will disconnect and get will return null) - const existingAppConnection = await this.appConnections?.get(serverName); - if (existingAppConnection) { - return MCPServerInspector.getToolFunctions(serverName, existingAppConnection); + const registry = MCPServersRegistry.getInstance(); + const effectiveConfig = serverConfig ?? (await registry.getServerConfig(serverName, userId)); + const useAppConnection = + effectiveConfig != null && + canUseAppConnection(effectiveConfig) && + (await registry.isAppServerConfig(serverName, effectiveConfig)); + const existingAppConnection = useAppConnection + ? await this.appConnections?.get(serverName) + : null; + if (existingAppConnection != null) { + return { + tools: await MCPServerInspector.getToolFunctions(serverName, existingAppConnection), + }; } const userConnections = this.getUserConnections(userId); if (!userConnections || userConnections.size === 0) { - return null; + return { tools: null }; } if (!userConnections.has(serverName)) { - return null; + return { tools: null }; } - return MCPServerInspector.getToolFunctions(serverName, userConnections.get(serverName)!); + const connection = userConnections.get(serverName)!; + 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 }; + } + 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}`, error, ); - return null; + return { tools: null }; } } + /** Returns all available tool functions from all connections available to user. */ + public async getServerToolFunctions( + userId: string, + serverName: string, + ): Promise { + return (await this.getServerToolFunctionsSnapshot(userId, serverName)).tools; + } + /** * Get instructions for MCP servers * @param serverNames Optional array of server names. If not provided or empty, returns all servers. @@ -529,7 +621,7 @@ Please follow these instructions when using tools from the respective MCP server const hasPersistentUserConnections = !!userId && (this.userConnections.get(userId)?.size ?? 0) > 0; if (!ephemeralConnection && hasPersistentUserConnections) { - this.updateUserLastActivity(userId); + await this.updateUserLastActivity(userId); } this.checkIdleConnections(); return formatToolContent(result as t.MCPToolCallResponse, provider); diff --git a/packages/api/src/mcp/UserConnectionManager.ts b/packages/api/src/mcp/UserConnectionManager.ts index 12902b2c85..7eadbf8252 100644 --- a/packages/api/src/mcp/UserConnectionManager.ts +++ b/packages/api/src/mcp/UserConnectionManager.ts @@ -3,6 +3,13 @@ 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, + getMCPAppToolsPublicationGeneration, + getMCPToolsChangedGeneration, + notifyMCPToolsChanged, + renewMCPToolsChangedGeneration, +} from '~/mcp/toolsChanged'; import { getMissingRuntimeBodyPlaceholderFields, hasRuntimeUrlPlaceholders, @@ -38,6 +45,8 @@ type PendingConnection = { oauth: PendingOAuthState; }; +type ConnectionCreationGuard = { cancelled: boolean }; + /** * Abstract base class for managing user-specific MCP connections with lifecycle management. * Only meant to be extended by MCPManager. @@ -54,14 +63,149 @@ export abstract class UserConnectionManager { protected userLastActivity: Map = new Map(); /** In-flight connection promises keyed by `userId:serverName` — coalesces concurrent attempts */ protected pendingConnections: Map = new Map(); + /** All durable creations, including forced replacements, visible to mutation teardown. */ + private readonly activeConnectionCreations: Map> = new Map(); + /** Serializes explicit durable replacements without coalescing their callers. */ + private readonly forceNewConnectionQueues: Map> = new Map(); + /** Fences durable connections whose credentials were invalidated on another replica. */ + protected readonly toolPublicationGenerations: WeakMap = new WeakMap(); + /** Binds a durable connection to the stored config that created it, independently of Redis. */ + protected readonly toolConfigGenerations: WeakMap = new WeakMap(); + /** Limits Redis lease refreshes while ensuring active connections cannot outlive their lease. */ + private readonly toolPublicationLeaseRefreshes: WeakMap = new WeakMap(); + /** Coalesces concurrent activity updates for the same durable connection. */ + private readonly toolPublicationLeaseRenewals: WeakMap> = + new WeakMap(); - /** Updates the last activity timestamp for a user */ - protected updateUserLastActivity(userId: string): void { + /** Records connections whose distributed publication authority was replaced during renewal. */ + private readonly lostToolPublicationLeases: WeakSet = new WeakSet(); + + /** Returns the cache-publication generation captured for a durable connection. */ + public getToolPublicationGeneration(connection: MCPConnection): string | undefined { + return this.toolPublicationGenerations.get(connection); + } + + /** Returns the config identity captured when a durable connection was created. */ + public getToolConfigGeneration(connection: MCPConnection): string | undefined { + return this.toolConfigGenerations.get(connection); + } + + private runWithForceNewConnectionQueue(key: string, operation: () => Promise): Promise { + const previous = this.forceNewConnectionQueues.get(key) ?? Promise.resolve(); + const result = previous.then(operation, operation); + const tail = result.then( + () => undefined, + () => undefined, + ); + this.forceNewConnectionQueues.set(key, tail); + void tail.then(() => { + if (this.forceNewConnectionQueues.get(key) === tail) { + this.forceNewConnectionQueues.delete(key); + } + }); + return result; + } + + private registerConnectionCreation(key: string, guard: ConnectionCreationGuard): void { + const guards = this.activeConnectionCreations.get(key) ?? new Set(); + guards.add(guard); + this.activeConnectionCreations.set(key, guards); + } + + private unregisterConnectionCreation(key: string, guard: ConnectionCreationGuard): void { + const guards = this.activeConnectionCreations.get(key); + guards?.delete(guard); + if (guards?.size === 0) { + this.activeConnectionCreations.delete(key); + } + } + + private cancelConnectionCreations(key: string, preserved?: ConnectionCreationGuard): void { + for (const guard of this.activeConnectionCreations.get(key) ?? []) { + if (guard !== preserved) { + guard.cancelled = true; + } + } + } + + private async renewUserToolPublicationLeases(userId: string, now: number): Promise { + const userConnections = this.userConnections.get(userId); + if (!userConnections) { + return; + } + const configuredIdleTimeout = Number(mcpConfig.USER_CONNECTION_IDLE_TIMEOUT); + const refreshInterval = + Number.isFinite(configuredIdleTimeout) && configuredIdleTimeout > 0 + ? Math.max(1_000, configuredIdleTimeout / 2) + : 15 * 60 * 1000; + const renewals: Promise[] = []; + for (const [serverName, connection] of userConnections) { + const publicationGeneration = this.toolPublicationGenerations.get(connection); + if (!publicationGeneration) { + continue; + } + const lastRefresh = this.toolPublicationLeaseRefreshes.get(connection) ?? 0; + if (now - lastRefresh < refreshInterval) { + continue; + } + const pendingRenewal = this.toolPublicationLeaseRenewals.get(connection); + if (pendingRenewal) { + renewals.push(pendingRenewal); + continue; + } + const renewal = renewMCPToolsChangedGeneration({ + userId, + serverName, + publicationGeneration, + }) + .then((renewed) => { + if (renewed === true) { + this.toolPublicationLeaseRefreshes.set(connection, now); + } else if (renewed === false) { + this.lostToolPublicationLeases.add(connection); + logger.info( + `[MCP][User: ${userId}][${serverName}] Publication lease is no longer current`, + ); + } + }) + .catch((error) => { + logger.warn( + `[MCP][User: ${userId}][${serverName}] Failed to renew tool publication lease`, + error, + ); + }); + this.toolPublicationLeaseRenewals.set(connection, renewal); + void renewal.finally(() => { + if (this.toolPublicationLeaseRenewals.get(connection) === renewal) { + this.toolPublicationLeaseRenewals.delete(connection); + } + }); + renewals.push(renewal); + } + await Promise.all(renewals); + } + + /** Updates activity and keeps every durable connection retained by that user leased. */ + protected async updateUserLastActivity(userId: string): Promise { const now = Date.now(); this.userLastActivity.set(userId, now); logger.debug( `[MCP][User: ${userId}] Updated last activity timestamp: ${new Date(now).toISOString()}`, ); + await this.renewUserToolPublicationLeases(userId, now); + } + + private async assertToolPublicationLeaseCurrent( + connection: MCPConnection, + userId: string, + serverName: string, + creationGuard?: ConnectionCreationGuard, + ): Promise { + if (!this.lostToolPublicationLeases.has(connection)) { + return; + } + await this.disconnectUserConnection(userId, serverName, creationGuard); + throw new Error(`[MCP][User: ${userId}][${serverName}] Publication lease is no longer current`); } /** Gets or creates a connection for a specific user, coalescing concurrent attempts */ @@ -165,25 +309,36 @@ export abstract class UserConnectionManager { } const pendingOAuth = this.createPendingOAuthState(opts.oauthStart); - const connectionPromise = this.createUserConnectionInternal( - { - ...opts, - forceNew: forceNewConnection, - ephemeralConnection, - serverConfig: config, - oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth), - }, - userId, - clearCooldown, - ); + const creationGuard: ConnectionCreationGuard = { cancelled: false }; + this.registerConnectionCreation(lockKey, creationGuard); + const createConnection = () => + this.createUserConnectionInternal( + { + ...opts, + forceNew: forceNewConnection, + ephemeralConnection, + serverConfig: config, + oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth), + }, + userId, + clearCooldown, + creationGuard, + ); + const connectionPromise = !ephemeralConnection + ? this.runWithForceNewConnectionQueue(lockKey, createConnection) + : createConnection(); if (!forceNewConnection) { - this.pendingConnections.set(lockKey, { promise: connectionPromise, oauth: pendingOAuth }); + this.pendingConnections.set(lockKey, { + promise: connectionPromise, + oauth: pendingOAuth, + }); } try { return await connectionPromise; } finally { + this.unregisterConnectionCreation(lockKey, creationGuard); if ( !forceNewConnection && this.pendingConnections.get(lockKey)?.promise === connectionPromise @@ -379,7 +534,13 @@ export abstract class UserConnectionManager { }: t.UserMCPConnectionOptions, userId: string, clearCooldown: boolean, + creationGuard?: ConnectionCreationGuard, ): Promise { + if (creationGuard?.cancelled) { + throw new Error( + `[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`, + ); + } if (await this.appConnections!.has(serverName)) { throw new McpError( ErrorCode.InvalidRequest, @@ -391,20 +552,67 @@ export abstract class UserConnectionManager { providedConfig ?? (await MCPServersRegistry.getInstance().getServerConfig(serverName, userId)); + /** Capture before resolving credentials/creating the connection. If another replica rotates + * the generation while creation is in flight, this connection's publications are fenced. */ + const publicationGeneration = ephemeralConnection + ? undefined + : await getMCPToolsChangedGeneration({ userId, serverName }); + const userServerMap = this.userConnections.get(userId); - let connection = forceNew ? undefined : userServerMap?.get(serverName); + let connection = userServerMap?.get(serverName); + if (forceNew && connection && !ephemeralConnection) { + logger.info( + `[MCP][User: ${userId}][${serverName}] Disposing existing connection before forced replacement`, + ); + await this.disconnectUserConnection(userId, serverName, creationGuard); + connection = undefined; + } else if (forceNew) { + connection = undefined; + } if (clearCooldown) { MCPConnection.clearCooldown(serverName); } const now = Date.now(); + const existingPublicationGeneration = connection + ? this.toolPublicationGenerations.get(connection) + : undefined; + const configGeneration = config ? getMCPAppToolsPublicationGeneration(config) : undefined; + const existingConfigGeneration = connection + ? this.toolConfigGenerations.get(connection) + : undefined; + if ( + connection && + configGeneration && + existingConfigGeneration && + configGeneration !== existingConfigGeneration + ) { + logger.info( + `[MCP][User: ${userId}][${serverName}] Config identity changed, disconnecting stale connection`, + ); + await this.disconnectUserConnection(userId, serverName, creationGuard); + connection = undefined; + } + if ( + connection && + publicationGeneration && + existingPublicationGeneration && + publicationGeneration !== existingPublicationGeneration + ) { + logger.info( + `[MCP][User: ${userId}][${serverName}] Cache generation changed, disconnecting stale connection`, + ); + await this.disconnectUserConnection(userId, serverName, creationGuard); + connection = undefined; + } + // Check if user is idle const lastActivity = this.userLastActivity.get(userId); if (lastActivity && now - lastActivity > mcpConfig.USER_CONNECTION_IDLE_TIMEOUT) { logger.info(`[MCP][User: ${userId}] User idle for too long. Disconnecting all connections.`); // Disconnect all user connections try { - await this.disconnectUserConnections(userId); + await this.disconnectUserConnections(userId, creationGuard); } catch (err) { logger.error(`[MCP][User: ${userId}] Error disconnecting idle connections:`, err); } @@ -416,18 +624,23 @@ export abstract class UserConnectionManager { `[MCP][User: ${userId}][${serverName}] Config was updated, disconnecting stale connection`, ); } - await this.disconnectUserConnection(userId, serverName); + await this.disconnectUserConnection(userId, serverName, creationGuard); connection = undefined; } else if (await connection.isConnected()) { logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing active connection`); - this.updateUserLastActivity(userId); + 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 { - // Connection exists but is not connected, attempt to remove potentially stale entry logger.warn( `[MCP][User: ${userId}][${serverName}] Found existing but disconnected connection object. Cleaning up.`, ); - this.removeUserConnection(userId, serverName); // Clean up maps + await this.disconnectUserConnection(userId, serverName, creationGuard); connection = undefined; } } @@ -512,11 +725,41 @@ export abstract class UserConnectionManager { connection = await MCPConnectionFactory.create(basic, connectionOptions); + if (creationGuard?.cancelled) { + throw new Error( + `[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`, + ); + } + + if (publicationGeneration) { + this.toolPublicationGenerations.set(connection, publicationGeneration); + } + if (configGeneration) { + this.toolConfigGenerations.set(connection, configGeneration); + } + + connection.on('toolsChanged', (tools: t.MCPTool[], publicationRevision?: string) => { + void notifyMCPToolsChanged({ + tools, + userId, + serverName, + serverConfig: config, + ...(publicationGeneration && { publicationGeneration }), + ...(publicationRevision && { publicationRevision }), + }); + }); + if (!(await connection?.isConnected())) { throw new Error('Failed to establish connection after initialization attempt.'); } if (!ephemeralConnection) { + await connection.refreshToolList(); + if (creationGuard?.cancelled) { + throw new Error( + `[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`, + ); + } if (!this.userConnections.has(userId)) { this.userConnections.set(userId, new Map()); } @@ -525,20 +768,29 @@ export abstract class UserConnectionManager { logger.info(`[MCP][User: ${userId}][${serverName}] Connection successfully established`); if (!ephemeralConnection) { - this.updateUserLastActivity(userId); + 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; } catch (error) { logger.error(`[MCP][User: ${userId}][${serverName}] Failed to establish connection`, error); // Ensure partial connection state is cleaned up if initialization fails - await connection?.disconnect().catch((disconnectError) => { + connection?.removeAllListeners?.('toolsChanged'); + await connection?.dispose().catch((disconnectError) => { logger.error( `[MCP][User: ${userId}][${serverName}] Error during cleanup after failed connection`, disconnectError, ); }); // Ensure cleanup even if connection attempt fails - this.removeUserConnection(userId, serverName); + if (connection && this.userConnections.get(userId)?.get(serverName) === connection) { + this.removeUserConnection(userId, serverName); + } throw error; // Re-throw the error to the caller } } @@ -711,19 +963,48 @@ export abstract class UserConnectionManager { } /** Disconnects and removes a specific user connection */ - public async disconnectUserConnection(userId: string, serverName: string): Promise { - this.pendingConnections.delete(`${userId}:${serverName}`); + public async disconnectUserConnection( + userId: string, + serverName: string, + preservedCreation?: ConnectionCreationGuard, + ): Promise { + const pendingKey = `${userId}:${serverName}`; + const pending = this.pendingConnections.get(pendingKey); + this.cancelConnectionCreations(pendingKey, preservedCreation); + if (pending && preservedCreation == null) { + this.pendingConnections.delete(pendingKey); + } const userMap = this.userConnections.get(userId); const connection = userMap?.get(serverName); - if (connection) { - logger.info(`[MCP][User: ${userId}][${serverName}] Disconnecting...`); - await connection.disconnect(); - this.removeUserConnection(userId, serverName); + try { + if (connection) { + logger.info(`[MCP][User: ${userId}][${serverName}] Disconnecting...`); + connection.removeAllListeners?.('toolsChanged'); + this.removeUserConnection(userId, serverName); + await connection.dispose(); + } + } finally { + await cancelMCPToolsChanged({ userId, serverName }); } } /** Disconnects and removes all connections for a specific user */ - public async disconnectUserConnections(userId: string): Promise { + public async disconnectUserConnections( + userId: string, + preservedCreation?: ConnectionCreationGuard, + ): Promise { + for (const key of this.activeConnectionCreations.keys()) { + if (key.startsWith(`${userId}:`)) { + this.cancelConnectionCreations(key, preservedCreation); + } + } + if (preservedCreation == null) { + for (const key of this.pendingConnections.keys()) { + if (key.startsWith(`${userId}:`)) { + this.pendingConnections.delete(key); + } + } + } const userMap = this.userConnections.get(userId); const disconnectPromises: Promise[] = []; if (userMap) { @@ -731,7 +1012,7 @@ export abstract class UserConnectionManager { const userServers = Array.from(userMap.keys()); for (const serverName of userServers) { disconnectPromises.push( - this.disconnectUserConnection(userId, serverName).catch((error) => { + this.disconnectUserConnection(userId, serverName, preservedCreation).catch((error) => { logger.error( `[MCP][User: ${userId}][${serverName}] Error during disconnection:`, error, @@ -740,12 +1021,6 @@ export abstract class UserConnectionManager { ); } await Promise.allSettled(disconnectPromises); - // Clean up any pending connection promises for this user - for (const key of this.pendingConnections.keys()) { - if (key.startsWith(`${userId}:`)) { - this.pendingConnections.delete(key); - } - } logger.info(`[MCP][User: ${userId}] All connections processed for disconnection.`); } /** diff --git a/packages/api/src/mcp/__tests__/ConnectionsRepository.test.ts b/packages/api/src/mcp/__tests__/ConnectionsRepository.test.ts index 73b88587de..5105df396e 100644 --- a/packages/api/src/mcp/__tests__/ConnectionsRepository.test.ts +++ b/packages/api/src/mcp/__tests__/ConnectionsRepository.test.ts @@ -1,5 +1,6 @@ import { logger } from '@librechat/data-schemas'; import type * as t from '~/mcp/types'; +import { getMCPAppToolsPublicationGeneration, setMCPToolsChangedHandler } from '~/mcp/toolsChanged'; import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory'; import { MCPConnection } from '~/mcp/connection'; @@ -7,8 +8,10 @@ import { MCPConnection } from '~/mcp/connection'; // Mock external dependencies jest.mock('@librechat/data-schemas', () => ({ logger: { + debug: jest.fn(), error: jest.fn(), info: jest.fn(), + warn: jest.fn(), }, })); @@ -54,6 +57,7 @@ describe('ConnectionsRepository', () => { let mockConnection: jest.Mocked; beforeEach(() => { + setMCPToolsChangedHandler(null); mockServerConfigs = { server1: { url: 'http://localhost:3001', type: 'sse' }, server2: { command: 'test-command', args: ['--test'], type: 'stdio' }, @@ -72,10 +76,19 @@ describe('ConnectionsRepository', () => { ) as jest.Mock; mockConnection = { + client: { + getServerCapabilities: jest.fn().mockReturnValue({ tools: {} }), + }, isConnected: jest.fn().mockResolvedValue(true), disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), + refreshToolList: jest.fn().mockResolvedValue(undefined), createdAt: Date.now(), isStale: jest.fn().mockReturnValue(false), + /* A real connection is an EventEmitter and the repository subscribes to it. */ + on: jest.fn(), + removeAllListeners: jest.fn(), } as unknown as jest.Mocked; (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); @@ -87,6 +100,7 @@ describe('ConnectionsRepository', () => { }); afterEach(() => { + setMCPToolsChangedHandler(null); jest.clearAllMocks(); }); @@ -127,19 +141,126 @@ describe('ConnectionsRepository', () => { undefined, ); expect(repository['connections'].get('server1')).toBe(mockConnection); + expect(mockConnection.fetchToolsSnapshot).toHaveBeenCalledTimes(1); + expect(mockConnection.refreshToolList).not.toHaveBeenCalled(); + }); + + it('serializes concurrent creation so only one connection is retained', async () => { + const [first, second] = await Promise.all([ + repository.get('server1'), + repository.get('server1'), + ]); + + expect(first).toBe(mockConnection); + expect(second).toBe(mockConnection); + expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(1); + expect(mockConnection.on).toHaveBeenCalledTimes(1); + }); + + it('awaits initial app tool publication before returning a new connection', async () => { + let releasePublication: (() => void) | undefined; + const publication = new Promise((resolve) => { + releasePublication = resolve; + }); + const handler = jest.fn(() => publication); + setMCPToolsChangedHandler(handler); + + let loaded = false; + const load = repository.get('server1').then((connection) => { + loaded = true; + return connection; + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ + serverName: 'server1', + tools: [], + publicationGeneration: getMCPAppToolsPublicationGeneration(mockServerConfigs.server1), + }), + ); + expect(loaded).toBe(false); + + releasePublication?.(); + await expect(load).resolves.toBe(mockConnection); + }); + + it('can defer the initial app tool refresh for startup synchronization', async () => { + await repository.get('server1', { refreshTools: false }); + + expect(mockConnection.fetchToolsSnapshot).not.toHaveBeenCalled(); + expect(mockConnection.refreshToolList).not.toHaveBeenCalled(); + }); + + it('queues the connection retry path when the initial app snapshot is incomplete', async () => { + mockConnection.fetchToolsSnapshot.mockResolvedValue({ tools: [], complete: false }); + + await repository.get('server1'); + + expect(mockConnection.refreshToolList).toHaveBeenCalledTimes(1); + }); + + it('does not publish an initial snapshot superseded by a list-changed refresh', async () => { + let toolsChanged: ((tools: t.MCPTool[]) => void) | undefined; + mockConnection.on.mockImplementation((event, listener) => { + if (event === 'toolsChanged') { + toolsChanged = listener as (tools: t.MCPTool[]) => void; + } + return mockConnection; + }); + let releaseInitialSnapshot: + | ((snapshot: { tools: t.MCPTool[]; complete: boolean }) => void) + | undefined; + mockConnection.fetchToolsSnapshot.mockReturnValue( + new Promise((resolve) => { + releaseInitialSnapshot = resolve; + }), + ); + const handler = jest.fn().mockResolvedValue(undefined); + setMCPToolsChangedHandler(handler); + + const load = repository.get('server1'); + await new Promise((resolve) => setImmediate(resolve)); + toolsChanged?.([{ name: 'current', inputSchema: { type: 'object' } } as t.MCPTool]); + releaseInitialSnapshot?.({ + tools: [{ name: 'stale', inputSchema: { type: 'object' } } as t.MCPTool], + complete: true, + }); + await load; + + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ tools: [expect.objectContaining({ name: 'current' })] }), + ); + expect(handler).not.toHaveBeenCalledWith( + expect.objectContaining({ tools: [expect.objectContaining({ name: 'stale' })] }), + ); + }); + + it('publishes an empty snapshot without listing tools when the server lacks the capability', async () => { + const handler = jest.fn().mockResolvedValue(undefined); + setMCPToolsChangedHandler(handler); + (mockConnection.client.getServerCapabilities as jest.Mock).mockReturnValue({ resources: {} }); + + await repository.get('server1'); + + expect(mockConnection.fetchToolsSnapshot).not.toHaveBeenCalled(); + expect(handler).toHaveBeenCalledWith( + expect.objectContaining({ serverName: 'server1', tools: [] }), + ); }); it('should create new connection if existing connection is not connected', async () => { const oldConnection = { isConnected: jest.fn().mockResolvedValue(false), disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as jest.Mocked; repository['connections'].set('server1', oldConnection); const result = await repository.get('server1'); expect(result).toBe(mockConnection); - expect(oldConnection.disconnect).toHaveBeenCalled(); + expect(oldConnection.dispose).toHaveBeenCalled(); expect(MCPConnectionFactory.create).toHaveBeenCalledWith( { serverName: 'server1', @@ -160,6 +281,7 @@ describe('ConnectionsRepository', () => { const staleConnection = { isConnected: jest.fn().mockResolvedValue(true), disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), createdAt: connectionCreatedAt, isStale: jest.fn().mockReturnValue(true), } as unknown as jest.Mocked; @@ -179,7 +301,7 @@ describe('ConnectionsRepository', () => { expect(staleConnection.isStale).toHaveBeenCalledWith(configCachedAt); // Verify old connection was disconnected - expect(staleConnection.disconnect).toHaveBeenCalled(); + expect(staleConnection.dispose).toHaveBeenCalled(); // Verify new connection was created expect(MCPConnectionFactory.create).toHaveBeenCalledWith( @@ -209,6 +331,7 @@ describe('ConnectionsRepository', () => { const freshConnection = { isConnected: jest.fn().mockResolvedValue(true), disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), createdAt: connectionCreatedAt, isStale: jest.fn().mockReturnValue(false), } as unknown as jest.Mocked; @@ -228,7 +351,7 @@ describe('ConnectionsRepository', () => { expect(freshConnection.isStale).toHaveBeenCalledWith(configCachedAt); // Verify connection was not disconnected - expect(freshConnection.disconnect).not.toHaveBeenCalled(); + expect(freshConnection.dispose).not.toHaveBeenCalled(); // Verify no new connection was created expect(MCPConnectionFactory.create).not.toHaveBeenCalled(); @@ -239,6 +362,17 @@ describe('ConnectionsRepository', () => { // Verify repository still has the same connection expect(repository['connections'].get('server1')).toBe(freshConnection); }); + + it('uses the repository cleanup path when a loaded server config is removed', async () => { + repository['connections'].set('server1', mockConnection); + delete mockServerConfigs.server1; + + await expect(repository.get('server1')).resolves.toBeNull(); + + expect(mockConnection.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(mockConnection.dispose).toHaveBeenCalled(); + expect(repository['connections'].has('server1')).toBe(false); + }); //todo revist later when async getAll(): in packages/api/src/mcp/ConnectionsRepository.ts is refactored it.skip('should throw error for non-existent server configuration', async () => { await expect(repository.get('nonexistent')).rejects.toThrow( @@ -295,6 +429,27 @@ describe('ConnectionsRepository', () => { expect(result.get('server2')).toBe(mockConnection); expect(result.get('server3')).toBe(mockConnection); }); + + it('continues loading later servers when one connection fails', async () => { + (MCPConnectionFactory.create as jest.Mock).mockImplementation( + ({ serverName }: { serverName: string }) => { + if (serverName === 'server1') { + return Promise.reject(new Error('server unavailable')); + } + return Promise.resolve(mockConnection); + }, + ); + + const result = await repository.getAll({ continueOnError: true }); + + expect(result.has('server1')).toBe(false); + expect(result.get('server2')).toBe(mockConnection); + expect(result.get('server3')).toBe(mockConnection); + expect(mockLogger.warn).toHaveBeenCalledWith( + '[MCP][server1] Failed to establish connection', + expect.any(Error), + ); + }); }); describe('disconnect', () => { @@ -303,36 +458,39 @@ describe('ConnectionsRepository', () => { await repository.disconnect('server1'); - expect(mockConnection.disconnect).toHaveBeenCalled(); + expect(mockConnection.dispose).toHaveBeenCalled(); expect(repository['connections'].has('server1')).toBe(false); }); it('should handle disconnect error gracefully', async () => { const disconnectError = new Error('Disconnect failed'); - mockConnection.disconnect.mockRejectedValue(disconnectError); + mockConnection.dispose.mockRejectedValue(disconnectError); repository['connections'].set('server1', mockConnection); await repository.disconnect('server1'); - expect(mockConnection.disconnect).toHaveBeenCalled(); + expect(mockConnection.dispose).toHaveBeenCalled(); expect(repository['connections'].has('server1')).toBe(false); expect(mockLogger.error).toHaveBeenCalledWith( - '[MCP][server1] Error disconnecting', + '[MCP][server1] Error disposing', disconnectError, ); }); }); describe('disconnectAll', () => { - it('should disconnect all active connections', () => { + it('should disconnect all active connections', async () => { const mockConnection1 = { disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as jest.Mocked; const mockConnection2 = { disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as jest.Mocked; const mockConnection3 = { disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as jest.Mocked; repository['connections'].set('server1', mockConnection1); @@ -341,8 +499,32 @@ describe('ConnectionsRepository', () => { const promises = repository.disconnectAll(); - expect(promises).toHaveLength(3); + expect(promises).toHaveLength(1); expect(Array.isArray(promises)).toBe(true); + await Promise.all(promises); + expect(mockConnection1.dispose).toHaveBeenCalledTimes(1); + expect(mockConnection2.dispose).toHaveBeenCalledTimes(1); + expect(mockConnection3.dispose).toHaveBeenCalledTimes(1); + }); + + it('drains and disposes a connection whose creation finishes during shutdown', async () => { + let releaseCreation: ((connection: MCPConnection) => void) | undefined; + (MCPConnectionFactory.create as jest.Mock).mockReturnValue( + new Promise((resolve) => { + releaseCreation = resolve; + }), + ); + + const load = repository.get('server1'); + await new Promise((resolve) => setImmediate(resolve)); + const shutdown = Promise.all(repository.disconnectAll()); + releaseCreation?.(mockConnection); + + await expect(load).resolves.toBeNull(); + await shutdown; + expect(mockConnection.dispose).toHaveBeenCalledTimes(1); + await expect(repository.get('server1')).resolves.toBeNull(); + expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(1); }); }); @@ -375,6 +557,17 @@ describe('ConnectionsRepository', () => { expect(await repository.has('defaultServer')).toBe(true); }); + it('should NOT allow app connections to public user-managed servers', async () => { + mockServerConfigs.publicServer = { + type: 'streamable-http', + url: 'https://public.example.com/mcp', + source: 'user', + requiresOAuth: false, + }; + + expect(await repository.has('publicServer')).toBe(false); + }); + it('should NOT allow connection to OAuth servers', async () => { mockServerConfigs.oauthServer = { type: 'streamable-http', @@ -475,7 +668,7 @@ describe('ConnectionsRepository', () => { const allowed = await repository.has('changingServer'); expect(allowed).toBe(false); - expect(mockConnection.disconnect).toHaveBeenCalled(); + expect(mockConnection.dispose).toHaveBeenCalled(); }); }); @@ -496,6 +689,17 @@ describe('ConnectionsRepository', () => { expect(await repository.has('regularServer')).toBe(true); }); + it('should lazily allow user connections to public user-managed servers', async () => { + mockServerConfigs.publicServer = { + type: 'streamable-http', + url: 'https://public.example.com/mcp', + source: 'user', + requiresOAuth: false, + }; + + expect(await repository.has('publicServer')).toBe(true); + }); + it('should allow connection to OAuth servers', async () => { mockServerConfigs.oauthServer = { type: 'streamable-http', diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index b4700b8bc1..9b28a0f8e9 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -3758,7 +3758,9 @@ describe('MCPConnectionFactory', () => { mockConnectionInstance.connect.mockResolvedValue(undefined); mockConnectionInstance.isConnected.mockResolvedValue(true); - mockConnectionInstance.fetchTools = jest.fn().mockResolvedValue(mockTools); + mockConnectionInstance.fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: mockTools, complete: true }); const result = await MCPConnectionFactory.discoverTools(basicOptions); @@ -3768,6 +3770,25 @@ describe('MCPConnectionFactory', () => { expect(result.connection).toBe(mockConnectionInstance); }); + it('does not expose an incomplete discovery snapshot as authoritative', async () => { + const basicOptions = { + serverName: 'test-server', + serverConfig: mockServerConfig, + }; + + mockConnectionInstance.connect.mockResolvedValue(undefined); + mockConnectionInstance.isConnected.mockResolvedValue(true); + mockConnectionInstance.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + tools: [mockTools[0]], + complete: false, + }); + + const result = await MCPConnectionFactory.discoverTools(basicOptions); + + expect(result.tools).toBeNull(); + expect(result.connection).toBe(mockConnectionInstance); + }); + it('should forward user context to processMCPEnv for non-OAuth discovery', async () => { const serverConfig: t.MCPOptions = { type: 'streamable-http', @@ -3787,7 +3808,9 @@ describe('MCPConnectionFactory', () => { mockConnectionInstance.connect.mockResolvedValue(undefined); mockConnectionInstance.isConnected.mockResolvedValue(true); - mockConnectionInstance.fetchTools = jest.fn().mockResolvedValue(mockTools); + mockConnectionInstance.fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: mockTools, complete: true }); const result = await MCPConnectionFactory.discoverTools(basicOptions, userContext); @@ -4340,7 +4363,9 @@ describe('MCPConnectionFactory', () => { mockFlowManager.createFlowWithHandler.mockResolvedValue(null); mockConnectionInstance.connect.mockResolvedValue(undefined); mockConnectionInstance.isConnected.mockResolvedValue(true); - mockConnectionInstance.fetchTools = jest.fn().mockResolvedValue(mockTools); + mockConnectionInstance.fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: mockTools, complete: true }); const result = await MCPConnectionFactory.discoverTools( { serverName: 'bigquery', serverConfig }, diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts index 212f4dc498..58b3dc892a 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts @@ -54,7 +54,7 @@ function createConnectionWithListTools(listTools: jest.Mock): MCPConnection { serverConfig: { type: 'streamable-http', url: 'http://localhost/mcp' }, useSSRFProtection: false, }); - conn.client = { listTools } as unknown as MCPConnection['client']; + conn.client.listTools = listTools; return conn; } @@ -81,6 +81,16 @@ describe('MCPConnection.fetchTools pagination', () => { mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000; }); + it('does not queue tool-list retries when the server has no tools capability', async () => { + const listTools = jest.fn(); + const conn = createConnectionWithListTools(listTools); + jest.spyOn(conn.client, 'getServerCapabilities').mockReturnValue({ resources: {} }); + + await conn.refreshToolList(); + + expect(listTools).not.toHaveBeenCalled(); + }); + it('returns the tools from a single page and makes one request when there is no nextCursor', async () => { const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('a'), makeTool('b')] }); const conn = createConnectionWithListTools(listTools); @@ -93,6 +103,33 @@ describe('MCPConnection.fetchTools pagination', () => { expect(mockLogger.warn).not.toHaveBeenCalled(); }); + it('returns the notification snapshot when a request snapshot races list_changed', async () => { + let releaseStale: ((value: { tools: ReturnType[] }) => void) | undefined; + const stale = new Promise<{ tools: ReturnType[] }>((resolve) => { + releaseStale = resolve; + }); + const listTools = jest + .fn() + .mockReturnValueOnce(stale) + .mockResolvedValueOnce({ tools: [makeTool('current')] }); + const conn = createConnectionWithListTools(listTools); + Reflect.set(conn, 'connectionState', 'connected'); + jest.spyOn(conn.client, 'getServerCapabilities').mockReturnValue({ + tools: { listChanged: true }, + }); + + const requested = conn.fetchOrderedToolsSnapshot(); + await Promise.resolve(); + const notified = conn.refreshToolList(); + await notified; + releaseStale?.({ tools: [makeTool('stale')] }); + + await expect(requested).resolves.toEqual({ + tools: [makeTool('current')], + complete: true, + }); + }); + it('follows nextCursor across pages, concatenating every tool and passing the cursor back', async () => { const listTools = jest.fn(async (params?: { cursor?: string }) => { switch (params?.cursor) { @@ -183,7 +220,7 @@ describe('MCPConnection.fetchTools pagination', () => { expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('size budget')); }); - it('stops at the elapsed-time budget before requesting another page', async () => { + it('marks a time-truncated snapshot incomplete before requesting another page', async () => { mcpConfig.TOOLS_LIST_TIMEOUT_MS = 1; const listTools = jest.fn(async () => ({ tools: [makeTool('a')], nextCursor: 'c1' })); const conn = createConnectionWithListTools(listTools); @@ -193,9 +230,10 @@ describe('MCPConnection.fetchTools pagination', () => { .mockReturnValueOnce(1000) .mockReturnValueOnce(1001); - const tools = await conn.fetchTools(); + const snapshot = await conn.fetchToolsSnapshot(); - expect(tools.map((t) => t.name)).toEqual(['a']); + expect(snapshot.tools.map((t) => t.name)).toEqual(['a']); + expect(snapshot.complete).toBe(false); expect(listTools).toHaveBeenCalledTimes(1); expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('time budget')); dateNow.mockRestore(); @@ -228,11 +266,12 @@ describe('MCPConnection.fetchTools pagination', () => { const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('x')], nextCursor: 'same' }); const conn = createConnectionWithListTools(listTools); - const tools = await conn.fetchTools(); + const snapshot = await conn.fetchToolsSnapshot(); expect(listTools).toHaveBeenCalledTimes(2); // The second page's tools are collected before the repeated cursor is detected, hence two copies. - expect(tools.map((t) => t.name)).toEqual(['x', 'x']); + expect(snapshot.tools.map((tool) => tool.name)).toEqual(['x', 'x']); + expect(snapshot.complete).toBe(false); expect(mockLogger.warn).toHaveBeenCalledWith( expect.stringContaining('repeated tools/list cursor'), ); diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index 5b88acf1a9..cafbe0091d 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -8,6 +8,7 @@ import { MCPServerInspector } from '~/mcp/registry/MCPServerInspector'; import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory'; import { isMCPDomainAllowed } from '~/auth/domain'; +import * as toolsChanged from '~/mcp/toolsChanged'; import { MCPConnection } from '~/mcp/connection'; import { MCPManager } from '~/mcp/MCPManager'; import * as graphUtils from '~/utils/graph'; @@ -47,6 +48,7 @@ const mockGetAllowedDomains = jest.fn().mockReturnValue(null); const mockGetAllowedAddresses = jest.fn().mockReturnValue(null); const mockRegistryInstance = { getServerConfig: jest.fn(), + isAppServerConfig: jest.fn(), getAllServerConfigs: jest.fn(), getOAuthServers: jest.fn(), shouldEnableSSRFProtection: mockShouldEnableSSRFProtection, @@ -97,6 +99,13 @@ describe('MCPManager', () => { // Set up default mock implementations (MCPServersInitializer.initialize as jest.Mock).mockResolvedValue(undefined); (mockRegistryInstance.getAllServerConfigs as jest.Mock).mockResolvedValue({}); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'stdio', + command: 'test', + args: [], + source: 'yaml', + }); + (mockRegistryInstance.isAppServerConfig as jest.Mock).mockResolvedValue(true); (mockRegistryInstance.shouldEnableSSRFProtection as jest.Mock).mockReturnValue(false); (mockRegistryInstance.getAllowedDomains as jest.Mock).mockReturnValue(null); (mockRegistryInstance.getAllowedAddresses as jest.Mock).mockReturnValue(null); @@ -114,6 +123,8 @@ describe('MCPManager', () => { const mock = { has: jest.fn().mockResolvedValue(false), get: jest.fn().mockResolvedValue({} as unknown as MCPConnection), + getAll: jest.fn().mockResolvedValue(new Map()), + getMany: jest.fn().mockResolvedValue(new Map()), ...appConnectionsConfig, }; return ( @@ -144,6 +155,17 @@ describe('MCPManager', () => { expect(result).toEqual({}); }); + it('does not open live connections while collecting the startup snapshot', async () => { + const getAll = jest.fn().mockResolvedValue(new Map()); + mockAppConnections({ getAll }); + (mockRegistryInstance.getAllServerConfigs as jest.Mock).mockResolvedValue({}); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getAppToolFunctions(); + + expect(getAll).not.toHaveBeenCalled(); + }); + it('should collect tool functions from multiple servers', async () => { const toolFunctions1 = { tool1_mcp_server1: { @@ -228,6 +250,81 @@ describe('MCPManager', () => { expect(result).toEqual(toolFunctions1); }); + + it('excludes public user-managed servers from the app catalog', async () => { + const operatorKey = 'operator_tool_mcp_operator'; + const publicKey = 'public_tool_mcp_public'; + (mockRegistryInstance.getAllServerConfigs as jest.Mock).mockResolvedValue({ + operator: { + type: 'stdio', + command: 'operator', + source: 'yaml', + toolFunctions: { + [operatorKey]: { type: 'function', function: { name: operatorKey } }, + }, + }, + public: { + type: 'streamable-http', + url: 'https://public.example.com/mcp', + source: 'user', + toolFunctions: { + [publicKey]: { type: 'function', function: { name: publicKey } }, + }, + }, + }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + + await expect(manager.getAppToolFunctions()).resolves.toEqual({ + [operatorKey]: { type: 'function', function: { name: operatorKey } }, + }); + }); + }); + + describe('connectAppServers', () => { + it('opens only operator app connections and refreshes their current catalogs', async () => { + const connection = new MCPConnection({ + serverName: 'dynamic', + serverConfig: { type: 'streamable-http', url: 'http://localhost/mcp' }, + useSSRFProtection: false, + }); + const refreshToolList = jest.spyOn(connection, 'refreshToolList').mockResolvedValue(); + const getMany = jest.fn().mockResolvedValue(new Map([['dynamic', connection]])); + mockAppConnections({ getMany }); + (mockRegistryInstance.getAllServerConfigs as jest.Mock).mockResolvedValue({ + dynamic: { + type: 'streamable-http', + url: 'http://localhost/mcp', + source: 'yaml', + }, + public: { + type: 'streamable-http', + url: 'https://public.example.com/mcp', + source: 'user', + }, + }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.connectAppServers(); + + expect(getMany).toHaveBeenCalledWith(['dynamic'], { + continueOnError: true, + refreshTools: false, + }); + expect(refreshToolList).toHaveBeenCalledTimes(1); + }); + }); + + describe('disconnectAppServers', () => { + it('waits for every loaded app connection to disconnect', async () => { + const disconnectAll = jest.fn().mockReturnValue([Promise.resolve(), Promise.resolve()]); + mockAppConnections({ disconnectAll }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.disconnectAppServers(); + + expect(disconnectAll).toHaveBeenCalledTimes(1); + }); }); describe('formatInstructionsForContext', () => { @@ -443,6 +540,45 @@ describe('MCPManager', () => { expect.any(Error), ); }); + + it('uses a user connection when an effective overlay shadows an app server', async () => { + const overlayConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://tenant.example.com/mcp', + source: 'yaml', + startup: false, + }; + const overlayConnection = {} as MCPConnection; + const expectedTools: t.LCAvailableTools = { + overlay_mcp_test_server: { + type: 'function', + function: { + name: 'overlay_mcp_test_server', + description: 'Overlay tool', + parameters: { type: 'object' }, + }, + }, + }; + const appGet = jest.fn().mockResolvedValue({} as MCPConnection); + mockAppConnections({ get: appGet }); + (mockRegistryInstance.isAppServerConfig as jest.Mock).mockResolvedValue(false); + (MCPServerInspector.getToolFunctions as jest.Mock).mockResolvedValue(expectedTools); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const internals = manager as unknown as { + userConnections: Map>; + }; + internals.userConnections.set(userId, new Map([[serverName, overlayConnection]])); + + await expect( + manager.getServerToolFunctionsSnapshot(userId, serverName, overlayConfig), + ).resolves.toEqual({ tools: expectedTools, publicationGeneration: undefined }); + expect(appGet).not.toHaveBeenCalled(); + expect(MCPServerInspector.getToolFunctions).toHaveBeenCalledWith( + serverName, + overlayConnection, + ); + }); }); describe('callTool - Activity Tracking', () => { @@ -540,6 +676,9 @@ describe('MCPManager', () => { const mockConnection = { isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), + disconnect: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), setRequestHeaders: jest.fn(), timeout: 30000, client: { @@ -697,8 +836,9 @@ describe('MCPManager', () => { >[0]['flowManager'], }); - /** One pass from user-connection runtime resolution, one from callTool — none from the handler attach */ - expect(mockProcessMCPEnv).toHaveBeenCalledTimes(2); + /** Runtime resolution, config-identity binding, and callTool each process once — none from + * attaching the request handler. */ + expect(mockProcessMCPEnv).toHaveBeenCalledTimes(3); expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledWith( expect.objectContaining({ serverConfig: processedServerConfig, @@ -1352,7 +1492,10 @@ describe('MCPManager', () => { const mockConnection = { isConnected: jest.fn().mockResolvedValue(true), fetchTools: jest.fn().mockResolvedValue(mockTools), + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: mockTools, complete: true }), + fetchOrderedToolsSnapshot: jest.fn().mockResolvedValue({ tools: mockTools, complete: true }), disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as MCPConnection; beforeEach(() => { @@ -1376,6 +1519,7 @@ describe('MCPManager', () => { it('should use MCPConnectionFactory.discoverTools when no app connection available', async () => { const discoveryConnection = { disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as MCPConnection; mockAppConnections({ get: jest.fn().mockResolvedValue(null), @@ -1400,7 +1544,7 @@ describe('MCPManager', () => { expect(result.tools).toEqual(mockTools); expect(result.oauthRequired).toBe(false); expect(MCPConnectionFactory.discoverTools).toHaveBeenCalled(); - expect(discoveryConnection.disconnect).toHaveBeenCalledTimes(1); + expect(discoveryConnection.dispose).toHaveBeenCalledTimes(1); }); it('should forward runtime context to discoverTools in the non-OAuth path', async () => { @@ -1590,7 +1734,9 @@ describe('MCPManager', () => { const mockConnection = { isConnected: jest.fn().mockResolvedValue(true), isStale: jest.fn().mockReturnValue(false), - disconnect: jest.fn(), + disconnect: jest.fn().mockResolvedValue(undefined), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), } as unknown as MCPConnection; it('should pass useOAuth for servers with configured oauth and no requiresOAuth value', async () => { @@ -1651,6 +1797,561 @@ describe('MCPManager', () => { ); }); + it('routes tool-list snapshots from user connections to the user-scoped publisher', async () => { + const serverConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/{{LIBRECHAT_BODY_CONVERSATIONID}}', + source: 'yaml', + requiresOAuth: false, + }; + const tools: t.MCPTool[] = [{ name: 'dynamic', inputSchema: { type: 'object' } }]; + const notifySpy = jest + .spyOn(toolsChanged, 'notifyMCPToolsChanged') + .mockResolvedValue(undefined); + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + mockProcessMCPEnv.mockImplementation(({ options }) => ({ + ...options, + url: 'https://mcp.example.com/conversation-1', + })); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ + serverName, + user: mockUser, + requestBody: { conversationId: 'conversation-1' }, + }); + + const listener = (mockConnection.on as jest.Mock).mock.calls.find( + ([eventName]) => eventName === 'toolsChanged', + )?.[1]; + expect(listener).toEqual(expect.any(Function)); + listener(tools); + + expect(notifySpy).toHaveBeenCalledWith({ + tools, + userId, + serverName, + serverConfig, + }); + } finally { + notifySpy.mockRestore(); + } + }); + + it('refreshes a newly recreated durable user connection after subscribing', async () => { + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + + expect(mockConnection.on).toHaveBeenCalledWith('toolsChanged', expect.any(Function)); + expect(mockConnection.refreshToolList).toHaveBeenCalledTimes(1); + }); + + it.each([false, true])( + 'cancels and disposes a pending connection during mutation teardown (forceNew=%s)', + async (forceNew) => { + const pendingConnection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + let resolveConnection: ((connection: MCPConnection) => void) | undefined; + const factoryResult = new Promise((resolve) => { + resolveConnection = resolve; + }); + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'streamable-http', + url: 'https://mcp.example.com/old', + source: 'user', + dbId: 'server-1', + }); + (MCPConnectionFactory.create as jest.Mock).mockReturnValue(factoryResult); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const creation = manager.getUserConnection({ serverName, user: mockUser, forceNew }); + while ((MCPConnectionFactory.create as jest.Mock).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + await manager.disconnectUserConnection(userId, serverName); + resolveConnection?.(pendingConnection); + + await expect(creation).rejects.toThrow('Connection creation was cancelled during teardown'); + expect(pendingConnection.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(pendingConnection.dispose).toHaveBeenCalledTimes(1); + expect(manager.getUserConnections(userId)?.has(serverName) ?? false).toBe(false); + }, + ); + + it('disposes the tracked durable connection before a forced replacement', async () => { + const createConnection = () => + ({ + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + }) as unknown as MCPConnection; + const previousConnection = createConnection(); + const replacementConnection = createConnection(); + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }); + (MCPConnectionFactory.create as jest.Mock) + .mockResolvedValueOnce(previousConnection) + .mockResolvedValueOnce(replacementConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const result = await manager.getUserConnection({ + serverName, + user: mockUser, + forceNew: true, + }); + + expect(previousConnection.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(previousConnection.dispose).toHaveBeenCalledTimes(1); + expect((previousConnection.dispose as jest.Mock).mock.invocationCallOrder[0]).toBeLessThan( + (MCPConnectionFactory.create as jest.Mock).mock.invocationCallOrder[1], + ); + expect(result).toBe(replacementConnection); + expect(manager.getUserConnections(userId)?.get(serverName)).toBe(replacementConnection); + }); + + it('serializes an ordinary load with multiple queued forced replacements', async () => { + const createConnection = () => + ({ + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + }) as unknown as MCPConnection; + const initial = createConnection(); + const firstReplacement = createConnection(); + const secondReplacement = createConnection(); + let resolveFirst: ((connection: MCPConnection) => void) | undefined; + let resolveSecond: ((connection: MCPConnection) => void) | undefined; + const firstFactoryResult = new Promise((resolve) => { + resolveFirst = resolve; + }); + const secondFactoryResult = new Promise((resolve) => { + resolveSecond = resolve; + }); + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }); + (MCPConnectionFactory.create as jest.Mock) + .mockResolvedValueOnce(initial) + .mockReturnValueOnce(firstFactoryResult) + .mockReturnValueOnce(secondFactoryResult); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const firstForced = manager.getUserConnection({ serverName, user: mockUser, forceNew: true }); + while ((MCPConnectionFactory.create as jest.Mock).mock.calls.length < 2) { + await new Promise((resolve) => setImmediate(resolve)); + } + const ordinary = manager.getUserConnection({ serverName, user: mockUser }); + const secondForced = manager.getUserConnection({ + serverName, + user: mockUser, + forceNew: true, + }); + + resolveFirst?.(firstReplacement); + await expect(firstForced).resolves.toBe(firstReplacement); + await expect(ordinary).resolves.toBe(firstReplacement); + while ((MCPConnectionFactory.create as jest.Mock).mock.calls.length < 3) { + await new Promise((resolve) => setImmediate(resolve)); + } + resolveSecond?.(secondReplacement); + + await expect(secondForced).resolves.toBe(secondReplacement); + expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(3); + expect(firstReplacement.dispose).toHaveBeenCalledTimes(1); + expect(manager.getUserConnections(userId)?.get(serverName)).toBe(secondReplacement); + }); + + it('disposes a retained disconnected connection before recreating it', async () => { + const createConnection = (connected: boolean) => + ({ + isConnected: jest.fn().mockResolvedValue(connected), + isStale: jest.fn().mockReturnValue(false), + disconnect: jest.fn().mockResolvedValue(undefined), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + }) as unknown as MCPConnection; + const disconnectedConnection = createConnection(false); + const replacementConnection = createConnection(true); + (disconnectedConnection.isConnected as jest.Mock) + .mockResolvedValueOnce(true) + .mockResolvedValue(false); + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({ + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }); + (MCPConnectionFactory.create as jest.Mock) + .mockResolvedValueOnce(disconnectedConnection) + .mockResolvedValueOnce(replacementConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const result = await manager.getUserConnection({ serverName, user: mockUser }); + + expect(disconnectedConnection.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(disconnectedConnection.dispose).toHaveBeenCalledTimes(1); + expect( + (disconnectedConnection.dispose as jest.Mock).mock.invocationCallOrder[0], + ).toBeLessThan((MCPConnectionFactory.create as jest.Mock).mock.invocationCallOrder[1]); + expect(result).toBe(replacementConnection); + expect(manager.getUserConnections(userId)?.get(serverName)).toBe(replacementConnection); + }); + + it('binds durable tool publications to the generation captured before connection creation', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValue('generation-a'); + const notifySpy = jest + .spyOn(toolsChanged, 'notifyMCPToolsChanged') + .mockResolvedValue(undefined); + const serverConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }; + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + get: jest.fn().mockResolvedValue(null), + }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); + const expectedToolFunctions: t.LCAvailableTools = { + [`current_mcp_${serverName}`]: { + type: 'function', + function: { + name: `current_mcp_${serverName}`, + description: '', + parameters: { type: 'object' }, + }, + }, + }; + (MCPServerInspector.getToolFunctions as jest.Mock).mockResolvedValue(expectedToolFunctions); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const listener = (mockConnection.on as jest.Mock).mock.calls.find( + ([eventName]) => eventName === 'toolsChanged', + )?.[1]; + const tools: t.MCPTool[] = [{ name: 'current', inputSchema: { type: 'object' } }]; + listener(tools); + const snapshot = await manager.getServerToolFunctionsSnapshot(userId, serverName); + + expect(generationSpy).toHaveBeenCalledWith({ userId, serverName }); + expect(snapshot).toEqual({ + tools: expectedToolFunctions, + publicationGeneration: 'generation-a', + }); + expect(notifySpy).toHaveBeenCalledWith({ + tools, + userId, + serverName, + serverConfig, + publicationGeneration: 'generation-a', + }); + } finally { + generationSpy.mockRestore(); + notifySpy.mockRestore(); + } + }); + + it('rejects a live snapshot after another replica rotates its publication generation', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValueOnce('generation-a') + .mockResolvedValue('generation-b'); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const serverConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'user', + dbId: 'server-1', + }; + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(connection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + + await expect( + manager.getServerToolFunctionsSnapshot(userId, serverName, serverConfig), + ).resolves.toEqual({ tools: null }); + expect(MCPServerInspector.getToolFunctions).not.toHaveBeenCalled(); + expect(connection.dispose).toHaveBeenCalledTimes(1); + } finally { + generationSpy.mockRestore(); + } + }); + + it('rejects an old-config connection even when Redis generation rotation failed', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValue('generation-a'); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const oldConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://old.example.com/mcp', + source: 'user', + dbId: 'server-1', + }; + const committedConfig: t.ParsedServerConfig = { + ...oldConfig, + url: 'https://new.example.com/mcp', + }; + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(oldConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(connection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser, serverConfig: oldConfig }); + + await expect( + manager.getServerToolFunctionsSnapshot(userId, serverName, committedConfig), + ).resolves.toEqual({ tools: null }); + expect(MCPServerInspector.getToolFunctions).not.toHaveBeenCalled(); + expect(connection.dispose).toHaveBeenCalledTimes(1); + } finally { + generationSpy.mockRestore(); + } + }); + + it('rejects a connection for a deleted config even when Redis generation rotation failed', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValue('generation-a'); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const oldConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://old.example.com/mcp', + source: 'user', + dbId: 'server-1', + }; + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(oldConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(connection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser, serverConfig: oldConfig }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(null); + + await expect(manager.getServerToolFunctionsSnapshot(userId, serverName)).resolves.toEqual({ + tools: null, + }); + expect(MCPServerInspector.getToolFunctions).not.toHaveBeenCalled(); + expect(connection.dispose).toHaveBeenCalledTimes(1); + } finally { + generationSpy.mockRestore(); + } + }); + + it('renews the publication lease when an active durable connection is reused', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValue('generation-a'); + const renewalSpy = jest + .spyOn(toolsChanged, 'renewMCPToolsChangedGeneration') + .mockResolvedValue(true); + const serverConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }; + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(mockConnection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const internals = manager as unknown as { + toolPublicationLeaseRefreshes: WeakMap; + }; + renewalSpy.mockClear(); + internals.toolPublicationLeaseRefreshes.set(mockConnection, 0); + + await manager.getUserConnection({ serverName, user: mockUser }); + + expect(renewalSpy).toHaveBeenCalledWith({ + userId, + serverName, + publicationGeneration: 'generation-a', + }); + expect(renewalSpy).toHaveBeenCalledTimes(1); + } finally { + generationSpy.mockRestore(); + renewalSpy.mockRestore(); + } + }); + + it('disposes and rejects a reused connection after lease renewal loses ownership', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValue('generation-a'); + const renewalSpy = jest + .spyOn(toolsChanged, 'renewMCPToolsChangedGeneration') + .mockResolvedValue(true); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const serverConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }; + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(connection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const internals = manager as unknown as { + toolPublicationLeaseRefreshes: WeakMap; + }; + internals.toolPublicationLeaseRefreshes.set(connection, 0); + renewalSpy.mockResolvedValue(false); + + await expect(manager.getUserConnection({ serverName, user: mockUser })).rejects.toThrow( + 'Publication lease is no longer current', + ); + + expect(connection.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(connection.dispose).toHaveBeenCalled(); + expect(manager.getUserConnections(userId)?.has(serverName) ?? false).toBe(false); + } finally { + generationSpy.mockRestore(); + renewalSpy.mockRestore(); + } + }); + + it('does not return a reused connection cancelled while its lease renewal is pending', async () => { + const generationSpy = jest + .spyOn(toolsChanged, 'getMCPToolsChangedGeneration') + .mockResolvedValue('generation-a'); + const renewalSpy = jest + .spyOn(toolsChanged, 'renewMCPToolsChangedGeneration') + .mockResolvedValue(true); + let resolveRenewal: ((renewed: boolean) => void) | undefined; + const pendingRenewal = new Promise((resolve) => { + resolveRenewal = resolve; + }); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + isStale: jest.fn().mockReturnValue(false), + refreshToolList: jest.fn().mockResolvedValue(undefined), + on: jest.fn(), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const serverConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', + startup: false, + }; + mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(connection); + + try { + const manager = await MCPManager.createInstance(newMCPServersConfig()); + await manager.getUserConnection({ serverName, user: mockUser }); + const internals = manager as unknown as { + toolPublicationLeaseRefreshes: WeakMap; + }; + internals.toolPublicationLeaseRefreshes.set(connection, 0); + renewalSpy.mockClear(); + renewalSpy.mockReturnValue(pendingRenewal); + + const reuse = manager.getUserConnection({ serverName, user: mockUser }); + while (renewalSpy.mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + await manager.disconnectUserConnection(userId, serverName); + resolveRenewal?.(true); + + await expect(reuse).rejects.toThrow('Connection creation was cancelled during teardown'); + expect(connection.dispose).toHaveBeenCalled(); + } finally { + generationSpy.mockRestore(); + renewalSpy.mockRestore(); + } + }); + it('should detect OAuth after resolving trusted runtime URL placeholders', async () => { const runtimeUrlConfig: t.ParsedServerConfig = { type: 'streamable-http', @@ -1869,9 +2570,11 @@ describe('MCPManager', () => { }; const firstConnection = { isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), } as unknown as MCPConnection; const secondConnection = { isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), } as unknown as MCPConnection; mockAppConnections({ @@ -1914,6 +2617,7 @@ describe('MCPManager', () => { }; const requestScopedConnection = { isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), } as unknown as MCPConnection; const requestScopedConnections: t.RequestScopedMCPConnectionStore = { connections: new Map(), diff --git a/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts b/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts index 169eb2b423..99fc8b2e13 100644 --- a/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts +++ b/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts @@ -67,7 +67,9 @@ describe('MCP OAuth Race Condition Fixes', () => { const manager = new TestManager(); const mockConnection = { + on: jest.fn(), isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), disconnect: jest.fn().mockResolvedValue(undefined), isStale: jest.fn().mockReturnValue(false), }; @@ -144,7 +146,9 @@ describe('MCP OAuth Race Condition Fixes', () => { const manager = new TestManager(); const mockConnection = { + on: jest.fn(), isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), disconnect: jest.fn().mockResolvedValue(undefined), isStale: jest.fn().mockReturnValue(false), }; @@ -229,7 +233,9 @@ describe('MCP OAuth Race Condition Fixes', () => { const manager = new TestManager(); const mockConnection = { + on: jest.fn(), isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), disconnect: jest.fn().mockResolvedValue(undefined), isStale: jest.fn().mockReturnValue(false), }; @@ -335,8 +341,12 @@ describe('MCP OAuth Race Condition Fixes', () => { let callCount = 0; const makeConnection = () => ({ + on: jest.fn(), isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), disconnect: jest.fn().mockResolvedValue(undefined), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), isStale: jest.fn().mockReturnValue(false), }); @@ -400,11 +410,172 @@ describe('MCP OAuth Race Condition Fixes', () => { expect(callCount).toBe(2); expect(conn1).not.toBe(conn2); + expect(conn1.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(conn1.dispose).toHaveBeenCalledTimes(1); } finally { createSpy.mockRestore(); registrySpy.mockRestore(); } }); + + it('waits for an ordinary pending connection before forcing its replacement', async () => { + const { UserConnectionManager } = await import('~/mcp/UserConnectionManager'); + + class TestManager extends UserConnectionManager {} + + const manager = new TestManager(); + const makeConnection = () => ({ + on: jest.fn(), + isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), + disconnect: jest.fn().mockResolvedValue(undefined), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + isStale: jest.fn().mockReturnValue(false), + }); + const firstConnection = makeConnection(); + const secondConnection = makeConnection(); + manager.appConnections = { has: jest.fn().mockResolvedValue(false) } as never; + + const registrySpy = jest + .spyOn( + // eslint-disable-next-line @typescript-eslint/no-require-imports + require('~/mcp/registry/MCPServersRegistry').MCPServersRegistry, + 'getInstance', + ) + .mockReturnValue({ + getServerConfig: jest.fn().mockResolvedValue({ + type: 'streamable-http', + url: 'http://localhost:9999/', + }), + resolveAllowlists: jest.fn().mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + }), + }); + const { MCPConnectionFactory } = await import('~/mcp/MCPConnectionFactory'); + let releaseFirstCreation: () => void = () => undefined; + const firstCreation = new Promise((resolve) => { + releaseFirstCreation = resolve; + }); + const createSpy = jest + .spyOn(MCPConnectionFactory, 'create') + .mockImplementationOnce(async () => { + await firstCreation; + return firstConnection as never; + }) + .mockResolvedValueOnce(secondConnection as never); + const user = { id: 'ordinary-and-force-new-user' }; + + try { + const ordinary = manager.getUserConnection({ + serverName: 'test-server', + user: user as never, + }); + for (let attempt = 0; attempt < 20 && createSpy.mock.calls.length === 0; attempt++) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(createSpy).toHaveBeenCalledTimes(1); + + const forced = manager.getUserConnection({ + serverName: 'test-server', + forceNew: true, + user: user as never, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(createSpy).toHaveBeenCalledTimes(1); + + releaseFirstCreation(); + const [ordinaryConnection, forcedConnection] = await Promise.all([ordinary, forced]); + + expect(ordinaryConnection).toBe(firstConnection); + expect(forcedConnection).toBe(secondConnection); + expect(createSpy).toHaveBeenCalledTimes(2); + expect(firstConnection.removeAllListeners).toHaveBeenCalledWith('toolsChanged'); + expect(firstConnection.dispose).toHaveBeenCalledTimes(1); + } finally { + releaseFirstCreation(); + createSpy.mockRestore(); + registrySpy.mockRestore(); + } + }); + + it('waits for an in-flight forced replacement before creating an ordinary connection', async () => { + const { UserConnectionManager } = await import('~/mcp/UserConnectionManager'); + + class TestManager extends UserConnectionManager {} + + const manager = new TestManager(); + const replacementConnection = { + on: jest.fn(), + isConnected: jest.fn().mockResolvedValue(true), + refreshToolList: jest.fn().mockResolvedValue(undefined), + disconnect: jest.fn().mockResolvedValue(undefined), + removeAllListeners: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), + isStale: jest.fn().mockReturnValue(false), + }; + manager.appConnections = { has: jest.fn().mockResolvedValue(false) } as never; + + const registrySpy = jest + .spyOn( + // eslint-disable-next-line @typescript-eslint/no-require-imports + require('~/mcp/registry/MCPServersRegistry').MCPServersRegistry, + 'getInstance', + ) + .mockReturnValue({ + getServerConfig: jest.fn().mockResolvedValue({ + type: 'streamable-http', + url: 'http://localhost:9999/', + }), + resolveAllowlists: jest.fn().mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + }), + }); + const { MCPConnectionFactory } = await import('~/mcp/MCPConnectionFactory'); + let releaseReplacement: () => void = () => undefined; + const replacementStarted = new Promise((resolve) => { + releaseReplacement = resolve; + }); + const createSpy = jest.spyOn(MCPConnectionFactory, 'create').mockImplementation(async () => { + await replacementStarted; + return replacementConnection as never; + }); + const user = { id: 'force-new-before-ordinary-user' }; + + try { + const forced = manager.getUserConnection({ + serverName: 'test-server', + forceNew: true, + user: user as never, + }); + for (let attempt = 0; attempt < 20 && createSpy.mock.calls.length === 0; attempt++) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(createSpy).toHaveBeenCalledTimes(1); + + const ordinary = manager.getUserConnection({ + serverName: 'test-server', + user: user as never, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(createSpy).toHaveBeenCalledTimes(1); + + releaseReplacement(); + const [forcedConnection, ordinaryConnection] = await Promise.all([forced, ordinary]); + + expect(forcedConnection).toBe(replacementConnection); + expect(ordinaryConnection).toBe(replacementConnection); + expect(createSpy).toHaveBeenCalledTimes(1); + } finally { + releaseReplacement(); + createSpy.mockRestore(); + registrySpy.mockRestore(); + } + }); }); describe('Fix 2: PENDING flow is reused, not deleted', () => { diff --git a/packages/api/src/mcp/__tests__/toolListChanged.integration.test.ts b/packages/api/src/mcp/__tests__/toolListChanged.integration.test.ts new file mode 100644 index 0000000000..141c8e242d --- /dev/null +++ b/packages/api/src/mcp/__tests__/toolListChanged.integration.test.ts @@ -0,0 +1,362 @@ +/** Real-SDK integration coverage for `notifications/tools/list_changed` (#7117). */ +import { CacheKeys } from 'librechat-data-provider'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; +import { + ErrorCode, + ListToolsRequestSchema, + McpError, + type Tool, +} from '@modelcontextprotocol/sdk/types.js'; +import type { LCAvailableTools } from '../types'; +import { + getMCPAppToolsPublicationGeneration, + setMCPToolsChangedRevisionHandler, +} from '../toolsChanged'; +import { createMCPCatalogStore } from '../catalog/store'; +import { MCPConnection } from '../connection'; + +jest.setTimeout(10_000); + +async function waitFor(condition: () => boolean, timeoutMs = 3_000): Promise { + const deadline = Date.now() + timeoutMs; + while (!condition()) { + if (Date.now() >= deadline) { + throw new Error('Timed out waiting for the tool-list refresh'); + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } +} + +interface Barrier { + entered: Promise; + release: () => void; +} + +interface TestHarness { + connection: MCPConnection; + getListCalls: () => number; + setTools: (tools: Tool[]) => void; + failNextList: () => void; + blockNextList: () => Barrier; + notifyChanged: () => Promise; + close: () => Promise; +} + +const tool = (name: string, description = name): Tool => ({ + name, + description, + inputSchema: { type: 'object', properties: {} }, +}); + +async function createHarness(initialTools: Tool[]): Promise { + let tools = initialTools; + let listCalls = 0; + let shouldFailNextList = false; + let nextBarrier: + | { + entered: Promise; + resolveEntered: () => void; + released: Promise; + release: () => void; + } + | undefined; + + const server = new Server( + { name: 'dynamic-tool-server', version: '1.0.0' }, + { capabilities: { tools: { listChanged: true } } }, + ); + server.setRequestHandler(ListToolsRequestSchema, async () => { + listCalls++; + if (shouldFailNextList) { + shouldFailNextList = false; + throw new McpError(ErrorCode.InternalError, 'temporary tools/list failure'); + } + + const snapshot = tools.map((entry) => ({ ...entry })); + const barrier = nextBarrier; + nextBarrier = undefined; + if (barrier) { + barrier.resolveEntered(); + await barrier.released; + } + return { tools: snapshot }; + }); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await server.connect(serverTransport); + + const connection = new MCPConnection({ + serverName: 'dynamic', + serverConfig: { type: 'streamable-http', url: 'https://mcp.example.com' }, + }); + await connection.client.connect(clientTransport); + connection.emit('connectionChange', 'connected'); + + return { + connection, + getListCalls: () => listCalls, + setTools: (nextTools) => { + tools = nextTools; + }, + failNextList: () => { + shouldFailNextList = true; + }, + blockNextList: () => { + let resolveEntered: (() => void) | undefined; + let release: (() => void) | undefined; + const entered = new Promise((resolve) => { + resolveEntered = resolve; + }); + const released = new Promise((resolve) => { + release = resolve; + }); + nextBarrier = { + entered, + resolveEntered: () => resolveEntered?.(), + released, + release: () => release?.(), + }; + return { entered, release: () => release?.() }; + }, + notifyChanged: () => server.sendToolListChanged(), + close: async () => { + connection.removeAllListeners(); + await connection.client.close().catch(() => undefined); + await server.close().catch(() => undefined); + }, + }; +} + +describe('tools/list_changed', () => { + let harness: TestHarness | undefined; + + afterEach(async () => { + setMCPToolsChangedRevisionHandler(null); + await harness?.close(); + harness = undefined; + }); + + it('publishes additions, removals, and schema changes from refreshed snapshots', async () => { + harness = await createHarness([tool('initial', 'version one')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + + harness.setTools([tool('initial', 'version one'), tool('added')]); + await harness.notifyChanged(); + await waitFor(() => snapshots.length === 1); + expect(snapshots[0].map(({ name }) => name)).toEqual(['initial', 'added']); + + harness.setTools([ + { + ...tool('initial', 'version two'), + inputSchema: { type: 'object', properties: { query: { type: 'string' } } }, + }, + ]); + await harness.notifyChanged(); + await waitFor(() => snapshots.length === 2); + + expect(snapshots[1]).toEqual([ + expect.objectContaining({ + name: 'initial', + description: 'version two', + inputSchema: expect.objectContaining({ + properties: { query: { type: 'string' } }, + }), + }), + ]); + }); + + it('does not lose a notification that arrives while a refresh is in flight', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + + const barrier = harness.blockNextList(); + harness.setTools([tool('initial'), tool('stale')]); + await harness.notifyChanged(); + await barrier.entered; + + harness.setTools([tool('initial'), tool('latest')]); + await harness.notifyChanged(); + barrier.release(); + await waitFor(() => snapshots.length === 2); + + expect(snapshots[0].map(({ name }) => name)).toEqual(['initial', 'stale']); + expect(snapshots[1].map(({ name }) => name)).toEqual(['initial', 'latest']); + }); + + it('keeps a newer cross-replica snapshot when an older request finishes last', async () => { + const olderReplica = await createHarness([tool('initial')]); + const newerReplica = await createHarness([tool('initial')]); + const cache = new Map(); + const store = createMCPCatalogStore({ + cacheConfig: { FORCED_IN_MEMORY_CACHE_NAMESPACES: [CacheKeys.TOOL_CACHE] }, + getCache: () => ({ + get: async (key) => cache.get(key), + set: async (key, value) => { + cache.set(key, value); + return true; + }, + delete: async (key) => cache.delete(key), + }), + }); + const configGeneration = getMCPAppToolsPublicationGeneration({ + type: 'streamable-http', + url: 'https://mcp.example.com', + }); + setMCPToolsChangedRevisionHandler(({ serverName, configGeneration }) => + store.getNextAppToolsPublicationRevision(serverName, configGeneration), + ); + const publications: Promise[] = []; + const publish = (tools: Tool[], publicationRevision?: string) => { + const catalog: LCAvailableTools = Object.fromEntries( + tools.map(({ name }) => [name, { type: 'function', function: { name } }]), + ); + publications.push( + store.setCachedAppServerTools('dynamic', configGeneration, catalog, publicationRevision), + ); + }; + olderReplica.connection.on('toolsChanged', publish); + newerReplica.connection.on('toolsChanged', publish); + + try { + const barrier = olderReplica.blockNextList(); + olderReplica.setTools([tool('stale')]); + await olderReplica.notifyChanged(); + await barrier.entered; + + newerReplica.setTools([tool('current')]); + await newerReplica.notifyChanged(); + await waitFor(() => publications.length === 1); + await Promise.all(publications); + + barrier.release(); + await waitFor(() => publications.length === 2); + await Promise.all(publications); + + await expect(store.getCachedAppServerTools('dynamic', configGeneration)).resolves.toEqual({ + current: { type: 'function', function: { name: 'current' } }, + }); + } finally { + await Promise.all([olderReplica.close(), newerReplica.close()]); + } + }); + + it('does not publish an in-flight snapshot after the connection is disconnected', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + + const barrier = harness.blockNextList(); + harness.setTools([tool('initial'), tool('stale-after-disconnect')]); + await harness.notifyChanged(); + await barrier.entered; + await harness.connection.disconnect(); + barrier.release(); + await new Promise((resolve) => setTimeout(resolve, 25)); + + expect(snapshots).toEqual([]); + }); + + it('discards an old transport snapshot when disconnect and reconnect race the refresh', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + + const barrier = harness.blockNextList(); + harness.setTools([tool('initial'), tool('stale-transport')]); + await harness.notifyChanged(); + await barrier.entered; + + harness.connection.emit('connectionChange', 'disconnected'); + harness.setTools([tool('initial'), tool('latest-transport')]); + harness.connection.emit('connectionChange', 'connected'); + barrier.release(); + + await waitFor(() => snapshots.length === 1); + expect(snapshots[0].map(({ name }) => name)).toEqual(['initial', 'latest-transport']); + }); + + it('refreshes after reconnect when the server changed tools while notifications were unavailable', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + + harness.connection.emit('connectionChange', 'disconnected'); + harness.setTools([tool('initial'), tool('added-while-disconnected')]); + harness.connection.emit('connectionChange', 'connected'); + + await waitFor(() => snapshots.length === 1); + expect(snapshots[0].map(({ name }) => name)).toEqual(['initial', 'added-while-disconnected']); + }); + + it('publishes an empty snapshot when a reconnect no longer advertises tools', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + jest.spyOn(harness.connection.client, 'getServerCapabilities').mockReturnValue({}); + + harness.connection.emit('connectionChange', 'disconnected'); + harness.connection.emit('connectionChange', 'connected'); + + await waitFor(() => snapshots.length === 1); + expect(snapshots[0]).toEqual([]); + }); + + it('clears a recreated connection when its current server no longer advertises tools', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + jest.spyOn(harness.connection.client, 'getServerCapabilities').mockReturnValue({}); + + await harness.connection.refreshToolList(); + + expect(snapshots).toEqual([[]]); + }); + + it('retains the last good catalog and retries after a transient list failure', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Tool[][] = []; + harness.connection.on('toolsChanged', (tools: Tool[]) => snapshots.push(tools)); + await harness.connection.fetchTools(); + const callsBeforeNotification = harness.getListCalls(); + + harness.setTools([tool('initial'), tool('recovered')]); + harness.failNextList(); + await harness.notifyChanged(); + await waitFor(() => harness!.getListCalls() > callsBeforeNotification); + expect(snapshots).toEqual([]); + + await waitFor(() => snapshots.length === 1); + expect(snapshots[0].map(({ name }) => name)).toEqual(['initial', 'recovered']); + }); + + it('retries when shared publication ordering is temporarily unavailable', async () => { + harness = await createHarness([tool('initial')]); + const snapshots: Array<{ tools: Tool[]; revision?: string }> = []; + const allocateRevision = jest + .fn, [{ serverName: string; configGeneration: string }]>() + .mockRejectedValueOnce(new Error('Redis unavailable')) + .mockResolvedValue('1'); + setMCPToolsChangedRevisionHandler(allocateRevision); + harness.connection.on('toolsChanged', (tools: Tool[], revision?: string) => + snapshots.push({ tools, revision }), + ); + + harness.setTools([tool('recovered')]); + await harness.notifyChanged(); + expect(snapshots).toEqual([]); + + await waitFor(() => snapshots.length === 1); + expect(snapshots[0]).toEqual({ tools: [tool('recovered')], revision: '1' }); + expect(allocateRevision).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/api/src/mcp/__tests__/utils.test.ts b/packages/api/src/mcp/__tests__/utils.test.ts index b0ff7bbf89..120278f683 100644 --- a/packages/api/src/mcp/__tests__/utils.test.ts +++ b/packages/api/src/mcp/__tests__/utils.test.ts @@ -18,6 +18,7 @@ import { getRuntimeBodyPlaceholderFields, getMissingRuntimeBodyPlaceholderFields, isUserSourced, + validateMCPServerConfig, requiresEphemeralUserConnection, } from '~/mcp/utils'; @@ -849,6 +850,29 @@ describe('getMissingRuntimeBodyPlaceholderFields', () => { }); }); +describe('validateMCPServerConfig', () => { + it('preserves server-managed metadata on a valid effective config', () => { + const config = { + type: 'streamable-http' as const, + url: 'https://example.com/mcp', + source: 'config' as const, + dbId: 'server-123', + }; + + expect(validateMCPServerConfig(config)).toBe(config); + expect(validateMCPServerConfig(config)).toMatchObject({ + source: 'config', + dbId: 'server-123', + }); + }); + + it('rejects an incomplete effective config', () => { + expect(() => validateMCPServerConfig({ type: 'streamable-http' })).toThrow( + 'Invalid effective MCP server configuration', + ); + }); +}); + describe('requiresEphemeralUserConnection', () => { it('returns true when BODY placeholders affect oauth_headers', () => { expect( diff --git a/packages/api/src/mcp/assistants.spec.ts b/packages/api/src/mcp/assistants.spec.ts new file mode 100644 index 0000000000..a326911404 --- /dev/null +++ b/packages/api/src/mcp/assistants.spec.ts @@ -0,0 +1,199 @@ +import { Constants } from 'librechat-data-provider'; +import type { LCAvailableTools, ParsedServerConfig } from './types'; +import type { AssistantToolDefinitionsDeps } from './assistants'; +import { getAssistantToolDefinitions } from './assistants'; + +const serverConfig: ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + source: 'yaml', +}; +const toolKey = `search${Constants.mcp_delimiter}app-server`; +const catalog: LCAvailableTools = { + [toolKey]: { + type: 'function', + ['function']: { + name: toolKey, + description: '', + parameters: { type: 'object', properties: {} }, + }, + }, +}; + +function createDeps( + overrides: Partial = {}, +): AssistantToolDefinitionsDeps { + return { + ensureConfigServers: jest.fn().mockResolvedValue({}), + getAllServerConfigs: jest.fn().mockResolvedValue({ 'app-server': serverConfig }), + getMCPServerTools: jest.fn().mockResolvedValue(catalog), + getServerToolFunctionsSnapshot: jest.fn().mockResolvedValue({ tools: null }), + recoverServerTools: jest.fn().mockResolvedValue(null), + cacheMCPServerTools: jest.fn().mockResolvedValue(undefined), + ...overrides, + }; +} + +describe('getAssistantToolDefinitions', () => { + const params = { + user: { id: 'user-1', role: 'user' }, + tools: ['code_interpreter', toolKey], + staticTools: { + code_interpreter: { type: 'function' as const, ['function']: { name: 'code_interpreter' } }, + }, + mcpConfig: {}, + }; + + it('combines static definitions with referenced configuration-addressed MCP catalogs', async () => { + const deps = createDeps(); + + await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual({ + ...params.staticTools, + ...catalog, + }); + expect(deps.getMCPServerTools).toHaveBeenCalledWith('user-1', 'app-server', serverConfig); + }); + + it('reconnects a user server when neither cache nor local snapshot has a catalog', async () => { + const recoveredCatalog = { ...catalog }; + const recoverServerTools = jest.fn().mockResolvedValue(recoveredCatalog); + const deps = createDeps({ + getMCPServerTools: jest.fn().mockResolvedValue(null), + getServerToolFunctionsSnapshot: jest.fn().mockResolvedValue({ tools: null }), + recoverServerTools, + }); + + await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual({ + ...params.staticTools, + ...recoveredCatalog, + }); + expect(recoverServerTools).toHaveBeenCalledWith('app-server', serverConfig); + }); + + it('limits concurrent cold catalog recovery across referenced servers', async () => { + const serverNames = ['one', 'two', 'three', 'four', 'five']; + const configs = Object.fromEntries( + serverNames.map((name) => [ + name, + { ...serverConfig, url: `https://${name}.example.com/mcp` }, + ]), + ); + let active = 0; + let peak = 0; + const recoverServerTools = jest.fn(async (name: string) => { + active++; + peak = Math.max(peak, active); + await new Promise((resolve) => setImmediate(resolve)); + active--; + const key = `search${Constants.mcp_delimiter}${name}`; + return { [key]: { type: 'function' as const, ['function']: { name: key } } }; + }); + const deps = createDeps({ + getAllServerConfigs: jest.fn().mockResolvedValue(configs), + getMCPServerTools: jest.fn().mockResolvedValue(null), + getServerToolFunctionsSnapshot: jest.fn().mockResolvedValue({ tools: null }), + recoverServerTools, + }); + + await getAssistantToolDefinitions( + { + ...params, + tools: serverNames.map((name) => `search${Constants.mcp_delimiter}${name}`), + }, + deps, + ); + + expect(recoverServerTools).toHaveBeenCalledTimes(serverNames.length); + expect(peak).toBe(3); + }); + + it('re-caches an authoritative empty local snapshot after a cache miss', async () => { + const cacheMCPServerTools = jest.fn().mockResolvedValue(undefined); + const deps = createDeps({ + getMCPServerTools: jest.fn().mockResolvedValue(null), + getServerToolFunctionsSnapshot: jest.fn().mockResolvedValue({ + tools: {}, + publicationGeneration: 'generation-1', + }), + cacheMCPServerTools, + }); + + await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual(params.staticTools); + expect(cacheMCPServerTools).toHaveBeenCalledWith({ + userId: 'user-1', + serverName: 'app-server', + serverTools: {}, + serverConfig, + publicationGeneration: 'generation-1', + }); + }); + + it('propagates config-server resolution failures instead of dropping selected tools', async () => { + const resolutionError = new Error('config resolution failed'); + const deps = createDeps({ + ensureConfigServers: jest.fn().mockRejectedValue(resolutionError), + }); + + await expect(getAssistantToolDefinitions(params, deps)).rejects.toBe(resolutionError); + expect(deps.getAllServerConfigs).not.toHaveBeenCalled(); + expect(deps.getMCPServerTools).not.toHaveBeenCalled(); + }); + + it('rejects a selected config server omitted by partial initialization', async () => { + const deps = createDeps({ + ensureConfigServers: jest.fn().mockResolvedValue({}), + getAllServerConfigs: jest.fn().mockResolvedValue({}), + }); + + await expect( + getAssistantToolDefinitions({ ...params, mcpConfig: { 'app-server': serverConfig } }, deps), + ).rejects.toThrow('MCP server configuration unavailable for assistant server "app-server"'); + expect(deps.getMCPServerTools).not.toHaveBeenCalled(); + }); + + it('does not route a missing normalized config server to a colliding available server', async () => { + const collidingToolKey = `search${Constants.mcp_delimiter}foo`; + const deps = createDeps({ + ensureConfigServers: jest.fn().mockResolvedValue({}), + getAllServerConfigs: jest.fn().mockResolvedValue({ foo: serverConfig }), + }); + + await expect( + getAssistantToolDefinitions( + { + ...params, + tools: [collidingToolKey], + mcpConfig: { 'foo!': { ...serverConfig, url: 'https://other.example.com/mcp' } }, + }, + deps, + ), + ).rejects.toThrow('MCP server configuration unavailable for assistant server "foo!"'); + expect(deps.getMCPServerTools).not.toHaveBeenCalled(); + }); + + it('does not classify a known static tool containing the MCP delimiter as an MCP reference', async () => { + const staticToolKey = `get${Constants.mcp_delimiter}status`; + const staticTools: LCAvailableTools = { + [staticToolKey]: { + type: 'function', + ['function']: { name: staticToolKey }, + }, + }; + const deps = createDeps({ + getAllServerConfigs: jest.fn().mockResolvedValue({ status: serverConfig }), + }); + + await expect( + getAssistantToolDefinitions( + { + ...params, + tools: [staticToolKey], + staticTools, + }, + deps, + ), + ).resolves.toBe(staticTools); + expect(deps.ensureConfigServers).not.toHaveBeenCalled(); + expect(deps.getMCPServerTools).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/api/src/mcp/assistants.ts b/packages/api/src/mcp/assistants.ts new file mode 100644 index 0000000000..25c123d0fc --- /dev/null +++ b/packages/api/src/mcp/assistants.ts @@ -0,0 +1,183 @@ +import { logger } from '@librechat/data-schemas'; +import { + Constants, + buildServerNameAliases, + normalizeServerName, + splitMCPToolKey, +} from 'librechat-data-provider'; +import type { MCPOptions } from 'librechat-data-provider'; +import type { LCAvailableTools, ParsedServerConfig } from '~/mcp/types'; +import { createConcurrencyLimiter } from '~/utils/promise'; +import { findShadowedServerNames } from '~/mcp/utils'; + +const RECOVERY_CONCURRENCY = 3; + +export interface AssistantMCPUser { + id?: string; + role?: string; +} + +export type AssistantToolReference = string | { type?: string; function?: { name?: string } }; + +export interface AssistantToolDefinitionsParams { + user: AssistantMCPUser | undefined; + tools?: readonly AssistantToolReference[]; + staticTools: LCAvailableTools; + mcpConfig: Record; +} + +export interface AssistantToolCatalogSnapshot { + tools: LCAvailableTools | null; + publicationGeneration?: string; +} + +export interface AssistantToolDefinitionsDeps { + ensureConfigServers: ( + mcpConfig: Record, + ) => Promise>; + getAllServerConfigs: ( + userId: string, + configServers: Record, + role?: string, + ) => Promise>; + getMCPServerTools: ( + userId: string, + serverName: string, + serverConfig: ParsedServerConfig, + ) => Promise; + getServerToolFunctionsSnapshot: ( + userId: string, + serverName: string, + serverConfig: ParsedServerConfig, + ) => Promise; + recoverServerTools: ( + serverName: string, + serverConfig: ParsedServerConfig, + ) => Promise; + cacheMCPServerTools: (params: { + userId: string; + serverName: string; + serverTools: LCAvailableTools; + serverConfig: ParsedServerConfig; + publicationGeneration?: string; + }) => Promise; +} + +function isMCPToolReference(tool: AssistantToolReference): tool is string { + return typeof tool === 'string' && tool.includes(Constants.mcp_delimiter); +} + +async function resolveAssistantMcpConfigs( + userId: string, + role: string | undefined, + mcpConfig: Record, + deps: AssistantToolDefinitionsDeps, +): Promise> { + const configServers = await deps.ensureConfigServers(mcpConfig); + return deps.getAllServerConfigs(userId, configServers, role); +} + +function selectReferencedServers( + toolNames: readonly string[], + configs: Record, + configuredServerNames: readonly string[] = [], +): Set { + const serverNames = [...new Set([...Object.keys(configs), ...configuredServerNames])]; + const aliases = buildServerNameAliases(serverNames); + const knownNames = [...new Set([...serverNames, ...aliases.keys()])]; + const shadowed = findShadowedServerNames(serverNames); + + return toolNames.reduce((selected, toolName) => { + const [, parsedServerName] = splitMCPToolKey(toolName, knownNames); + const unavailableConfiguredServer = configuredServerNames.find( + (serverName) => + !Object.prototype.hasOwnProperty.call(configs, serverName) && + (serverName === parsedServerName || normalizeServerName(serverName) === parsedServerName), + ); + if (unavailableConfiguredServer) { + throw new Error( + `MCP server configuration unavailable for assistant server "${unavailableConfiguredServer}"`, + ); + } + const serverName = + parsedServerName != null && Object.prototype.hasOwnProperty.call(configs, parsedServerName) + ? parsedServerName + : aliases.get(parsedServerName ?? ''); + if (serverName && !Object.prototype.hasOwnProperty.call(configs, serverName)) { + throw new Error(`MCP server configuration unavailable for assistant server "${serverName}"`); + } + if (serverName && !shadowed.has(serverName)) { + selected.add(serverName); + } + return selected; + }, new Set()); +} + +async function loadServerCatalog( + userId: string, + serverName: string, + serverConfig: ParsedServerConfig, + deps: AssistantToolDefinitionsDeps, + recover: (task: () => Promise) => Promise, +): Promise { + const cached = await deps.getMCPServerTools(userId, serverName, serverConfig); + if (cached != null) { + return cached; + } + + const snapshot = await deps.getServerToolFunctionsSnapshot(userId, serverName, serverConfig); + if (snapshot.tools != null) { + void deps + .cacheMCPServerTools({ + userId, + serverName, + serverTools: snapshot.tools, + serverConfig, + publicationGeneration: snapshot.publicationGeneration, + }) + .catch((error) => + logger.error( + `[assistant tool definitions] Failed to cache tools for ${serverName}:`, + error, + ), + ); + return snapshot.tools; + } + + const recovered = await recover(() => deps.recoverServerTools(serverName, serverConfig)); + if (recovered != null) { + return recovered; + } + throw new Error(`MCP tool definitions unavailable for assistant server "${serverName}"`); +} + +/** Loads the static catalog with the configuration-addressed MCP slices referenced by an assistant. */ +export async function getAssistantToolDefinitions( + params: AssistantToolDefinitionsParams, + deps: AssistantToolDefinitionsDeps, +): Promise { + const mcpToolNames = + params.tools?.filter( + (tool): tool is string => + isMCPToolReference(tool) && !Object.prototype.hasOwnProperty.call(params.staticTools, tool), + ) ?? []; + const userId = params.user?.id; + if (mcpToolNames.length === 0 || !userId) { + return params.staticTools; + } + + const configs = await resolveAssistantMcpConfigs( + userId, + params.user?.role, + params.mcpConfig, + deps, + ); + const recover = createConcurrencyLimiter(RECOVERY_CONCURRENCY); + const serverCatalogs = await Promise.all( + Array.from( + selectReferencedServers(mcpToolNames, configs, Object.keys(params.mcpConfig)), + (serverName) => loadServerCatalog(userId, serverName, configs[serverName], deps, recover), + ), + ); + return Object.assign({}, params.staticTools, ...serverCatalogs); +} diff --git a/packages/api/src/mcp/catalog/store.ts b/packages/api/src/mcp/catalog/store.ts new file mode 100644 index 0000000000..0ce949bbd9 --- /dev/null +++ b/packages/api/src/mcp/catalog/store.ts @@ -0,0 +1,886 @@ +import { randomUUID } from 'crypto'; +import { logger } from '@librechat/data-schemas'; +import { CacheKeys, Time } from 'librechat-data-provider'; +import type { LCAvailableTools } from '../types'; + +const GLOBAL_LOCK_TTL_MS = 30_000; +const LOCK_FENCE_SAFETY_MS = 1_000; +const LOCK_RETRY_MS = 25; +const CACHE_ENTRY_VERSION = 1; +const GLOBAL_LOCK_KEY = `${CacheKeys.TOOL_CACHE}:tools:global:write-lock`; + +const RELEASE_LOCK_SCRIPT = ` +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +`; + +const CLAIM_OWNERSHIP_FENCE_SCRIPT = ` +local current = redis.call('GET', KEYS[1]) +if current and current ~= ARGV[1] then return 0 end +local redisTime = redis.call('TIME') +local now = (tonumber(redisTime[1]) * 1000) + math.floor(tonumber(redisTime[2]) / 1000) +local ttl = tonumber(ARGV[2]) - now - tonumber(ARGV[3]) +if ttl <= 0 then return 0 end +redis.call('PSETEX', KEYS[1], ttl, ARGV[1]) +return 1 +`; + +const CREATE_GENERATION_IF_OWNER_SCRIPT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end +if redis.call('EXISTS', KEYS[2]) == 1 then return 0 end +redis.call('PSETEX', KEYS[2], ARGV[4], cjson.encode({value=ARGV[2], expires=tonumber(ARGV[3])})) +return 1 +`; + +const MIGRATE_USER_TOOLS_IF_OWNER_SCRIPT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end +if redis.call('EXISTS', KEYS[2]) == 1 or redis.call('EXISTS', KEYS[4]) == 1 then return 0 end +local rawGeneration = redis.call('GET', KEYS[3]) +if rawGeneration then + local decoded, generation = pcall(cjson.decode, rawGeneration) + if not decoded or type(generation) ~= 'table' or generation['value'] ~= ARGV[2] then return 0 end +else + redis.call('PSETEX', KEYS[3], ARGV[4], cjson.encode({value=ARGV[2], expires=tonumber(ARGV[3])})) +end +local decodedEntry, entry = pcall(cjson.decode, ARGV[5]) +if not decodedEntry then return redis.error_reply('Invalid MCP tools migration entry') end +redis.call('PSETEX', KEYS[4], ARGV[7], cjson.encode({value=entry, expires=tonumber(ARGV[6])})) +return 1 +`; + +const RENEW_GENERATION_SCRIPT = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then return 0 end +local decoded, value = pcall(cjson.decode, raw) +if not decoded or type(value) ~= 'table' or value['value'] ~= ARGV[1] then return 0 end +value['expires'] = tonumber(ARGV[2]) +redis.call('PSETEX', KEYS[1], ARGV[3], cjson.encode(value)) +return 1 +`; + +const WRITE_USER_TOOLS_SCRIPT = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then return 0 end +local decoded, value = pcall(cjson.decode, raw) +if not decoded or type(value) ~= 'table' or value['value'] ~= ARGV[1] then return 0 end +value['expires'] = tonumber(ARGV[2]) +redis.call('PSETEX', KEYS[1], ARGV[3], cjson.encode(value)) +local decodedEntry, entry = pcall(cjson.decode, ARGV[4]) +if not decodedEntry then return redis.error_reply('Invalid MCP tools cache entry') end +redis.call('PSETEX', KEYS[2], ARGV[6], cjson.encode({value=entry, expires=tonumber(ARGV[5])})) +return 1 +`; + +const WRITE_GLOBAL_IF_OWNER_SCRIPT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end +local decoded, tools = pcall(cjson.decode, ARGV[2]) +if not decoded then return redis.error_reply('Invalid global tools catalog') end +redis.call('PSETEX', KEYS[2], ARGV[4], cjson.encode({value=tools, expires=tonumber(ARGV[3])})) +return 1 +`; + +const NEXT_APP_TOOLS_REVISION_SCRIPT = ` +local current = redis.call('GET', KEYS[1]) +if not current then + redis.call('PSETEX', KEYS[1], ARGV[2], ARGV[1]) +end +local revision = redis.call('INCR', KEYS[1]) +redis.call('PEXPIRE', KEYS[1], ARGV[2]) +return tostring(revision) +`; + +const WRITE_APP_TOOLS_IF_CURRENT_SCRIPT = ` +local current = redis.call('GET', KEYS[1]) +if current and tonumber(current) > tonumber(ARGV[1]) then return 0 end +local rawCurrentEntry = redis.call('GET', KEYS[2]) +if rawCurrentEntry then + local decodedCurrent, currentEntry = pcall(cjson.decode, rawCurrentEntry) + if decodedCurrent and type(currentEntry) == 'table' and type(currentEntry['value']) == 'table' then + local storedRevision = currentEntry['value']['publicationRevision'] + local storedRevisionNumber = tonumber(storedRevision) + if storedRevisionNumber and storedRevisionNumber > tonumber(ARGV[1]) then return 0 end + end +end +local decodedEntry, entry = pcall(cjson.decode, ARGV[2]) +if not decodedEntry then return redis.error_reply('Invalid MCP app tools cache entry') end +redis.call('PSETEX', KEYS[1], ARGV[5], ARGV[1]) +redis.call('PSETEX', KEYS[2], ARGV[4], cjson.encode({value=entry, expires=tonumber(ARGV[3])})) +return 1 +`; + +const DELETE_GLOBAL_IF_OWNER_SCRIPT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end +redis.call('DEL', KEYS[2]) +return 1 +`; + +export const ToolCacheKeys = { + GLOBAL: 'tools:global', + MCP_APP_SERVER: (serverName: string, configGeneration: string): string => + `tools:mcp:app:${encodeURIComponent(serverName)}:${encodeURIComponent(configGeneration)}`, + MCP_SERVER: (userId: string, serverName: string, configGeneration?: string): string => + configGeneration + ? `tools:mcp:user:{${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}}:${encodeURIComponent(configGeneration)}` + : `tools:mcp:${userId}:${serverName}`, + MCP_SERVER_GENERATION: (userId: string, serverName: string): string => + `tools:metadata:mcp:user-generation:{${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}}`, + MCP_SERVER_LEGACY_FENCE: (userId: string, serverName: string): string => + `tools:metadata:mcp:user-legacy-fence:{${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}}`, +}; + +export interface CatalogCache { + get(key: string): Promise; + set(key: string, value: unknown, ttl?: number): Promise; + delete(key: string): Promise; +} + +interface LockRedisClient { + set(key: string, value: string, mode: 'PX', ttl: number, condition: 'NX'): Promise; + eval(script: string, numberOfKeys: number, ...args: string[]): Promise; +} + +interface KeyvRedisClient { + eval(script: string, options: { keys: string[]; arguments: string[] }): Promise; +} + +export interface CatalogStoreDeps { + getCache: () => CatalogCache; + cacheConfig: { + FORCED_IN_MEMORY_CACHE_NAMESPACES?: string[]; + REDIS_KEY_PREFIX?: string; + GLOBAL_PREFIX_SEPARATOR?: string; + }; + ioredisClient?: LockRedisClient | null; + keyvRedisClient?: KeyvRedisClient | null; + userConnectionIdleTimeout?: number | string; +} + +export interface CachedToolsOptions { + userId?: string; + serverName?: string; + configGeneration?: string; + ttl?: number; +} + +export interface GuardedToolsOptions extends CachedToolsOptions { + userId: string; + serverName: string; + configGeneration: string; + publicationGeneration: string; +} + +interface GuardedEntry { + version: number; + publicationGeneration: string; + tools: LCAvailableTools; +} + +interface AppToolsEntry { + version: number; + publicationRevision: string; + tools: LCAvailableTools; +} + +interface OwnershipFence { + key: string; + token: string; +} + +export interface MCPCatalogStore { + getCachedTools: (options?: CachedToolsOptions) => Promise; + updateCachedGlobalTools: (update: (tools: LCAvailableTools) => LCAvailableTools) => Promise; + setCachedTools: (tools: LCAvailableTools, options?: CachedToolsOptions) => Promise; + setCachedToolsWithinGlobalLock: ( + tools: LCAvailableTools, + options?: CachedToolsOptions, + ) => Promise; + setCachedToolsIfCurrent: ( + tools: LCAvailableTools, + options: GuardedToolsOptions, + ) => Promise; + getMCPToolsCacheGeneration: (scope: { + userId: string; + serverName: string; + }) => Promise; + renewMCPToolsCacheGeneration: (scope: { + userId: string; + serverName: string; + publicationGeneration: string; + }) => Promise; + getCachedAppServerTools: ( + serverName: string, + configGeneration: string, + ) => Promise; + getNextAppToolsPublicationRevision: ( + serverName: string, + configGeneration: string, + ) => Promise; + setCachedAppServerTools: ( + serverName: string, + configGeneration: string, + tools: LCAvailableTools, + publicationRevision?: string, + ttl?: number, + ) => Promise; + runWithGlobalCacheLock: (operation: () => Promise) => Promise; + invalidateCachedTools: (options?: { + userId?: string; + serverName?: string; + invalidateGlobal?: boolean; + }) => Promise; +} + +function isTools(value: unknown): value is LCAvailableTools { + return value != null && typeof value === 'object' && !Array.isArray(value); +} + +function isGuardedEntry(value: unknown): value is GuardedEntry { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const entry = value as Partial; + return ( + entry.version === CACHE_ENTRY_VERSION && + typeof entry.publicationGeneration === 'string' && + isTools(entry.tools) + ); +} + +function isAppToolsEntry(value: unknown): value is AppToolsEntry { + if (value == null || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + const entry = value as Partial; + return ( + entry.version === CACHE_ENTRY_VERSION && + typeof entry.publicationRevision === 'string' && + isTools(entry.tools) + ); +} + +function parseAppToolsRevision(revision: string): number { + const parsed = Number(revision); + if (!Number.isSafeInteger(parsed) || parsed < 0) { + throw new Error(`Invalid MCP app tools publication revision: ${revision}`); + } + return parsed; +} + +export function createMCPCatalogStore(deps: CatalogStoreDeps): MCPCatalogStore { + const generationTtl = Math.max( + Time.ONE_DAY, + Number.isFinite(Number(deps.userConnectionIdleTimeout)) + ? Number(deps.userConnectionIdleTimeout) * 2 + : 0, + ); + const userQueues = new Map>(); + const appRevisionCounters = new Map(); + const appCommittedRevisions = new Map(); + const appRevisionExpiryTimers = new Map>(); + const appWriteQueues = new Map>(); + let globalQueue = Promise.resolve(); + let globalLockToken: string | undefined; + + const sharedRedis = (): boolean => + deps.ioredisClient != null && + deps.keyvRedisClient != null && + !deps.cacheConfig.FORCED_IN_MEMORY_CACHE_NAMESPACES?.includes(CacheKeys.TOOL_CACHE); + + const rawKey = (key: string): string => { + const namespaced = `${CacheKeys.TOOL_CACHE}:${key}`; + return deps.cacheConfig.REDIS_KEY_PREFIX + ? `${deps.cacheConfig.REDIS_KEY_PREFIX}${deps.cacheConfig.GLOBAL_PREFIX_SEPARATOR ?? '::'}${namespaced}` + : namespaced; + }; + const rememberAppRevision = ( + revisions: Map, + scope: string, + revision: number, + ): void => { + revisions.set(scope, revision); + const timerKey = `${revisions === appRevisionCounters ? 'allocated' : 'committed'}:${scope}`; + const previousTimer = appRevisionExpiryTimers.get(timerKey); + if (previousTimer) { + clearTimeout(previousTimer); + } + const timer = setTimeout(() => { + if (revisions.get(scope) === revision) { + revisions.delete(scope); + } + appRevisionExpiryTimers.delete(timerKey); + }, Time.TWELVE_HOURS); + timer.unref?.(); + appRevisionExpiryTimers.set(timerKey, timer); + }; + function withAppWriteQueue(scope: string, operation: () => Promise): Promise { + const previous = appWriteQueues.get(scope) ?? Promise.resolve(); + const result = previous.then(operation, operation); + const tail = result.then( + () => undefined, + () => undefined, + ); + appWriteQueues.set(scope, tail); + void tail.then(() => appWriteQueues.get(scope) === tail && appWriteQueues.delete(scope)); + return result; + } + const redisHashTag = (key: string): string => { + const openingBrace = key.indexOf('{'); + const closingBrace = openingBrace >= 0 ? key.indexOf('}', openingBrace + 1) : -1; + return closingBrace > openingBrace + 1 ? key.slice(openingBrace + 1, closingBrace) : key; + }; + const globalCacheRawKey = rawKey(ToolCacheKeys.GLOBAL); + /** Hashes to the same Redis Cluster slot as the unchanged legacy global catalog key. */ + const globalFenceKey = `tools:global:write-fence:{${redisHashTag(globalCacheRawKey)}}`; + + async function withRedisLock( + lockKey: string, + ttl: number, + operation: (token?: string, leaseExpiresAt?: number) => Promise, + ): Promise { + if (!sharedRedis()) { + return operation(); + } + const redis = deps.ioredisClient!; + const token = randomUUID(); + const deadline = Date.now() + ttl + LOCK_RETRY_MS; + let leaseExpiresAt = 0; + while (true) { + const attemptedLeaseExpiry = Date.now() + ttl; + if ((await redis.set(lockKey, token, 'PX', ttl, 'NX')) === 'OK') { + leaseExpiresAt = attemptedLeaseExpiry; + break; + } + if (Date.now() >= deadline) { + throw new Error(`Timed out waiting for tool cache lock ${lockKey}`); + } + await new Promise((resolve) => setTimeout(resolve, LOCK_RETRY_MS)); + } + try { + return await operation(token, leaseExpiresAt); + } finally { + try { + await redis.eval(RELEASE_LOCK_SCRIPT, 1, lockKey, token); + } catch (error) { + logger.warn(`[MCP Cache] Failed to release tool cache lock ${lockKey}:`, error); + } + } + } + + async function withOwnershipFence( + fenceKey: string, + token: string | undefined, + leaseExpiresAt: number | undefined, + operation: (fence?: OwnershipFence) => Promise, + ): Promise { + if (!sharedRedis() || !token) { + return operation(); + } + const fenceTtl = (leaseExpiresAt ?? 0) - Date.now() - LOCK_FENCE_SAFETY_MS; + if (fenceTtl <= 0) { + throw new Error('Tool cache lock expired before ownership could be fenced'); + } + const claimed = await deps.keyvRedisClient!.eval(CLAIM_OWNERSHIP_FENCE_SCRIPT, { + keys: [fenceKey], + arguments: [token, String(leaseExpiresAt), String(LOCK_FENCE_SAFETY_MS)], + }); + if (Number(claimed) !== 1) { + throw new Error('Tool cache lock expired or was superseded before ownership could be fenced'); + } + try { + return await operation({ key: fenceKey, token }); + } finally { + try { + await deps.keyvRedisClient!.eval(RELEASE_LOCK_SCRIPT, { + keys: [fenceKey], + arguments: [token], + }); + } catch (error) { + logger.warn(`[MCP Cache] Failed to release tool cache fence ${fenceKey}:`, error); + } + } + } + + function withUserQueue( + userId: string, + serverName: string, + operation: (fence?: OwnershipFence) => Promise, + requireOwnershipFence = false, + ): Promise { + const scope = JSON.stringify([userId, serverName]); + const previous = userQueues.get(scope) ?? Promise.resolve(); + const lockKey = `${CacheKeys.TOOL_CACHE}:tools:mcp-write-lock:${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}`; + const generationRawKey = rawKey(ToolCacheKeys.MCP_SERVER_GENERATION(userId, serverName)); + const fenceKey = `tools:mcp:write-fence:{${redisHashTag(generationRawKey)}}`; + const locked = () => + withRedisLock(lockKey, GLOBAL_LOCK_TTL_MS, (token, leaseExpiresAt) => + requireOwnershipFence + ? withOwnershipFence(fenceKey, token, leaseExpiresAt, operation) + : operation(), + ); + const result = previous.then(locked, locked); + const tail = result.then( + () => undefined, + () => undefined, + ); + userQueues.set(scope, tail); + void tail.then(() => userQueues.get(scope) === tail && userQueues.delete(scope)); + return result; + } + + async function renewIfCurrent( + cache: CatalogCache, + key: string, + publicationGeneration: string, + ): Promise { + if (sharedRedis()) { + const renewed = await deps.keyvRedisClient!.eval(RENEW_GENERATION_SCRIPT, { + keys: [rawKey(key)], + arguments: [ + publicationGeneration, + String(Date.now() + generationTtl), + String(generationTtl), + ], + }); + return Number(renewed) === 1; + } + if ((await cache.get(key)) !== publicationGeneration) { + return false; + } + if ((await cache.set(key, publicationGeneration, generationTtl)) === false) { + throw new Error('Tool publication generation cache rejected the lease refresh'); + } + return true; + } + + async function setGlobalWithinLock( + cache: CatalogCache, + tools: LCAvailableTools, + ttl: number, + ): Promise { + if (!sharedRedis()) { + return (await cache.set(ToolCacheKeys.GLOBAL, tools, ttl)) !== false; + } + if (!globalLockToken) { + throw new Error('Global tool cache write requires lock ownership'); + } + const written = await deps.keyvRedisClient!.eval(WRITE_GLOBAL_IF_OWNER_SCRIPT, { + keys: [globalFenceKey, globalCacheRawKey], + arguments: [globalLockToken, JSON.stringify(tools), String(Date.now() + ttl), String(ttl)], + }); + if (Number(written) !== 1) { + throw new Error('Global tool cache lock ownership was lost before write'); + } + return true; + } + + async function deleteGlobalWithinLock(cache: CatalogCache): Promise { + if (!sharedRedis()) { + await cache.delete(ToolCacheKeys.GLOBAL); + return; + } + if (!globalLockToken) { + throw new Error('Global tool cache invalidation requires lock ownership'); + } + const deleted = await deps.keyvRedisClient!.eval(DELETE_GLOBAL_IF_OWNER_SCRIPT, { + keys: [globalFenceKey, globalCacheRawKey], + arguments: [globalLockToken], + }); + if (Number(deleted) !== 1) { + throw new Error('Global tool cache lock ownership was lost before invalidation'); + } + } + + function runWithGlobalCacheLock(operation: () => Promise): Promise { + const locked = () => + withRedisLock(GLOBAL_LOCK_KEY, GLOBAL_LOCK_TTL_MS, async (token, leaseExpiresAt) => { + return withOwnershipFence(globalFenceKey, token, leaseExpiresAt, async (fence) => { + globalLockToken = fence?.token; + try { + return await operation(); + } finally { + globalLockToken = undefined; + } + }); + }); + const result = globalQueue.then(locked, locked); + globalQueue = result.then( + () => undefined, + () => undefined, + ); + return result; + } + + async function getCachedTools( + options: CachedToolsOptions = {}, + ): Promise { + const cache = deps.getCache(); + const { userId, serverName, configGeneration } = options; + if (!userId || !serverName) { + const global = await cache.get(ToolCacheKeys.GLOBAL); + return isTools(global) ? global : null; + } + const toolsKey = ToolCacheKeys.MCP_SERVER(userId, serverName, configGeneration); + let cached = await cache.get(toolsKey); + if (cached == null && configGeneration) { + cached = await withUserQueue( + userId, + serverName, + async (fence) => { + const current = await cache.get(toolsKey); + if (current != null) return current; + const legacyFenceKey = ToolCacheKeys.MCP_SERVER_LEGACY_FENCE(userId, serverName); + if ((await cache.get(legacyFenceKey)) != null) { + return null; + } + const legacy = await cache.get(ToolCacheKeys.MCP_SERVER(userId, serverName)); + if (!isTools(legacy)) return null; + const generationKey = ToolCacheKeys.MCP_SERVER_GENERATION(userId, serverName); + const generation = await cache.get(generationKey); + const publicationGeneration = + typeof generation === 'string' && generation.length > 0 ? generation : randomUUID(); + const migrated: GuardedEntry = { + version: CACHE_ENTRY_VERSION, + publicationGeneration, + tools: legacy, + }; + if (!sharedRedis()) { + if (publicationGeneration !== generation) { + if ( + (await cache.set(generationKey, publicationGeneration, generationTtl)) === false + ) { + throw new Error('Tool publication generation cache rejected the migration fence'); + } + } + if ((await cache.set(toolsKey, migrated, Time.TWELVE_HOURS)) === false) { + throw new Error('Tool cache rejected the legacy user catalog migration'); + } + return migrated; + } + if (!fence) { + throw new Error('Legacy tool migration requires lock ownership'); + } + const toolsTtl = Time.TWELVE_HOURS; + const written = await deps.keyvRedisClient!.eval(MIGRATE_USER_TOOLS_IF_OWNER_SCRIPT, { + keys: [fence.key, rawKey(legacyFenceKey), rawKey(generationKey), rawKey(toolsKey)], + arguments: [ + fence.token, + publicationGeneration, + String(Date.now() + generationTtl), + String(generationTtl), + JSON.stringify(migrated), + String(Date.now() + toolsTtl), + String(toolsTtl), + ], + }); + if (Number(written) < 0) { + throw new Error('Tool cache lock ownership was lost before legacy migration'); + } + return Number(written) === 1 ? migrated : await cache.get(toolsKey); + }, + true, + ); + } + if (!isGuardedEntry(cached)) { + return isTools(cached) ? cached : null; + } + const generation = await cache.get(ToolCacheKeys.MCP_SERVER_GENERATION(userId, serverName)); + return generation === cached.publicationGeneration ? cached.tools : null; + } + + async function updateCachedGlobalTools( + update: (tools: LCAvailableTools) => LCAvailableTools, + ): Promise { + await runWithGlobalCacheLock(async () => { + const cache = deps.getCache(); + const current = await cache.get(ToolCacheKeys.GLOBAL); + const currentTools = isTools(current) ? current : {}; + const next = update(currentTools); + if (isTools(current) && next === current) { + return; + } + if (!(await setGlobalWithinLock(cache, next, Time.TWELVE_HOURS))) { + throw new Error('Global tool cache rejected the migration write'); + } + }); + } + + async function setCachedToolsWithinGlobalLock( + tools: LCAvailableTools, + options: CachedToolsOptions = {}, + ): Promise { + const cache = deps.getCache(); + const ttl = options.ttl ?? Time.TWELVE_HOURS; + if (options.userId && options.serverName) { + return ( + (await cache.set( + ToolCacheKeys.MCP_SERVER(options.userId, options.serverName, options.configGeneration), + tools, + ttl, + )) !== false + ); + } + return setGlobalWithinLock(cache, tools, ttl); + } + + async function setCachedTools( + tools: LCAvailableTools, + options: CachedToolsOptions = {}, + ): Promise { + if (options.userId && options.serverName) { + return withUserQueue(options.userId, options.serverName, () => + setCachedToolsWithinGlobalLock(tools, options), + ); + } + return runWithGlobalCacheLock(() => setCachedToolsWithinGlobalLock(tools, options)); + } + + async function getMCPToolsCacheGeneration(scope: { + userId: string; + serverName: string; + }): Promise { + const cache = deps.getCache(); + const key = ToolCacheKeys.MCP_SERVER_GENERATION(scope.userId, scope.serverName); + const existing = await cache.get(key); + if (typeof existing === 'string' && existing.length > 0) return existing; + return withUserQueue( + scope.userId, + scope.serverName, + async (fence) => { + const current = await cache.get(key); + if (typeof current === 'string' && current.length > 0) return current; + const generation = randomUUID(); + if (!sharedRedis()) { + if ((await cache.set(key, generation, generationTtl)) === false) { + throw new Error('Tool publication generation cache rejected the write'); + } + return generation; + } + if (!fence) { + throw new Error('Tool publication generation creation requires lock ownership'); + } + const created = await deps.keyvRedisClient!.eval(CREATE_GENERATION_IF_OWNER_SCRIPT, { + keys: [fence.key, rawKey(key)], + arguments: [ + fence.token, + generation, + String(Date.now() + generationTtl), + String(generationTtl), + ], + }); + if (Number(created) < 0) { + throw new Error('Tool cache lock ownership was lost before generation creation'); + } + if (Number(created) === 1) { + return generation; + } + const concurrent = await cache.get(key); + if (typeof concurrent === 'string' && concurrent.length > 0) { + return concurrent; + } + throw new Error('Tool publication generation changed during creation'); + }, + true, + ); + } + + async function renewMCPToolsCacheGeneration(scope: { + userId: string; + serverName: string; + publicationGeneration: string; + }): Promise { + const cache = deps.getCache(); + return withUserQueue(scope.userId, scope.serverName, () => + renewIfCurrent( + cache, + ToolCacheKeys.MCP_SERVER_GENERATION(scope.userId, scope.serverName), + scope.publicationGeneration, + ), + ); + } + + async function setCachedToolsIfCurrent( + tools: LCAvailableTools, + options: GuardedToolsOptions, + ): Promise { + const cache = deps.getCache(); + return withUserQueue(options.userId, options.serverName, async () => { + const generationKey = ToolCacheKeys.MCP_SERVER_GENERATION(options.userId, options.serverName); + const guarded: GuardedEntry = { + version: CACHE_ENTRY_VERSION, + publicationGeneration: options.publicationGeneration, + tools, + }; + if (!sharedRedis()) { + if (!(await renewIfCurrent(cache, generationKey, options.publicationGeneration))) { + return false; + } + return ( + (await cache.set( + ToolCacheKeys.MCP_SERVER(options.userId, options.serverName, options.configGeneration), + guarded, + options.ttl ?? Time.TWELVE_HOURS, + )) !== false + ); + } + const ttl = options.ttl ?? Time.TWELVE_HOURS; + const written = await deps.keyvRedisClient!.eval(WRITE_USER_TOOLS_SCRIPT, { + keys: [ + rawKey(generationKey), + rawKey( + ToolCacheKeys.MCP_SERVER(options.userId, options.serverName, options.configGeneration), + ), + ], + arguments: [ + options.publicationGeneration, + String(Date.now() + generationTtl), + String(generationTtl), + JSON.stringify(guarded), + String(Date.now() + ttl), + String(ttl), + ], + }); + return Number(written) === 1; + }); + } + + async function getCachedAppServerTools( + serverName: string, + configGeneration: string, + ): Promise { + const value = await deps + .getCache() + .get(ToolCacheKeys.MCP_APP_SERVER(serverName, configGeneration)); + if (isAppToolsEntry(value)) { + return value.tools; + } + return isTools(value) ? value : null; + } + + async function getNextAppToolsPublicationRevision( + serverName: string, + configGeneration: string, + ): Promise { + const toolsKey = ToolCacheKeys.MCP_APP_SERVER(serverName, configGeneration); + const scope = JSON.stringify([serverName, configGeneration]); + const cached = await deps.getCache().get(toolsKey); + const cachedRevision = isAppToolsEntry(cached) + ? parseAppToolsRevision(cached.publicationRevision) + : 0; + if (!sharedRedis()) { + const revision = Math.max(appRevisionCounters.get(scope) ?? 0, cachedRevision) + 1; + rememberAppRevision(appRevisionCounters, scope, revision); + return String(revision); + } + const toolsRawKey = rawKey(toolsKey); + const revisionKey = rawKey(`tools:mcp:app-revision:{${redisHashTag(toolsRawKey)}}`); + const revision = await deps.keyvRedisClient!.eval(NEXT_APP_TOOLS_REVISION_SCRIPT, { + keys: [revisionKey], + arguments: [String(cachedRevision), String(Time.TWELVE_HOURS)], + }); + const parsedRevision = parseAppToolsRevision(String(revision)); + return String(parsedRevision); + } + + async function setCachedAppServerTools( + serverName: string, + configGeneration: string, + tools: LCAvailableTools, + publicationRevision = '0', + ttl = Time.TWELVE_HOURS, + ): Promise { + const nextRevision = parseAppToolsRevision(publicationRevision); + const toolsKey = ToolCacheKeys.MCP_APP_SERVER(serverName, configGeneration); + const entry: AppToolsEntry = { + version: CACHE_ENTRY_VERSION, + publicationRevision, + tools, + }; + const scope = JSON.stringify([serverName, configGeneration]); + if (!sharedRedis()) { + return withAppWriteQueue(scope, async () => { + const cached = await deps.getCache().get(toolsKey); + const cachedRevision = isAppToolsEntry(cached) + ? parseAppToolsRevision(cached.publicationRevision) + : 0; + const currentRevision = Math.max(appCommittedRevisions.get(scope) ?? 0, cachedRevision); + if (currentRevision > nextRevision) { + return false; + } + if ((await deps.getCache().set(toolsKey, entry, ttl)) === false) { + throw new Error('App tool cache rejected the write'); + } + rememberAppRevision(appCommittedRevisions, scope, nextRevision); + return true; + }); + } + const toolsRawKey = rawKey(toolsKey); + const committedRevisionKey = rawKey( + `tools:mcp:app-committed-revision:{${redisHashTag(toolsRawKey)}}`, + ); + const written = await deps.keyvRedisClient!.eval(WRITE_APP_TOOLS_IF_CURRENT_SCRIPT, { + keys: [committedRevisionKey, toolsRawKey], + arguments: [ + publicationRevision, + JSON.stringify(entry), + String(Date.now() + ttl), + String(ttl), + String(ttl), + ], + }); + return Number(written) === 1; + } + + async function invalidateCachedTools( + options: { + userId?: string; + serverName?: string; + invalidateGlobal?: boolean; + } = {}, + ): Promise { + const cache = deps.getCache(); + if (options.invalidateGlobal) { + await runWithGlobalCacheLock(() => deleteGlobalWithinLock(cache)); + } + const { userId, serverName } = options; + if (userId && serverName) { + await withUserQueue(userId, serverName, async () => { + if ( + (await cache.set( + ToolCacheKeys.MCP_SERVER_LEGACY_FENCE(userId, serverName), + true, + generationTtl, + )) === false + ) { + throw new Error('Tool cache rejected the legacy migration fence'); + } + if ( + (await cache.set( + ToolCacheKeys.MCP_SERVER_GENERATION(userId, serverName), + randomUUID(), + generationTtl, + )) === false + ) { + throw new Error('Tool publication generation cache rejected invalidation'); + } + await cache.delete(ToolCacheKeys.MCP_SERVER(userId, serverName)); + }); + } + } + + return { + getCachedTools, + updateCachedGlobalTools, + setCachedTools, + setCachedToolsIfCurrent, + getMCPToolsCacheGeneration, + renewMCPToolsCacheGeneration, + setCachedToolsWithinGlobalLock, + getCachedAppServerTools, + getNextAppToolsPublicationRevision, + setCachedAppServerTools, + runWithGlobalCacheLock, + invalidateCachedTools, + }; +} diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index 121ec7480d..56e5dff936 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -5,12 +5,15 @@ import { fetch as undiciFetch, Agent, ProxyAgent } from 'undici'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; import { WebSocketClientTransport } from '@modelcontextprotocol/sdk/client/websocket.js'; -import { ResourceListChangedNotificationSchema } from '@modelcontextprotocol/sdk/types.js'; import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; import { StdioClientTransport, getDefaultEnvironment, } from '@modelcontextprotocol/sdk/client/stdio.js'; +import { + ResourceListChangedNotificationSchema, + ToolListChangedNotificationSchema, +} from '@modelcontextprotocol/sdk/types.js'; import type { RequestInit as UndiciRequestInit, RequestInfo as UndiciRequestInfo, @@ -21,6 +24,7 @@ import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; import type { MCPOAuthTokens } from './oauth/types'; import type * as t from './types'; import { createSSRFSafeUndiciConnect, isSSRFTarget, resolveHostnameSSRF } from '~/auth'; +import { reserveMCPToolsChangedRevision } from './toolsChanged'; import { isOAuthServer, sanitizeUrlForLogging } from './utils'; import { runOutsideTracing } from '~/utils/tracing'; import { isAddressAllowed } from '~/auth/domain'; @@ -121,6 +125,8 @@ const SSE_CONNECT_TIMEOUT = 120000; const DEFAULT_INIT_TIMEOUT = 30000; /** Upper bound on the spec-mandated Streamable HTTP session DELETE so teardown never blocks on a hung server */ const SESSION_TERMINATION_TIMEOUT = 5000; +const TOOL_LIST_REFRESH_RETRY_BASE_MS = 250; +const TOOL_LIST_REFRESH_RETRY_MAX_MS = 30_000; /** Max 307/308 redirects to follow per request (prevents redirect loops) */ const MAX_REDIRECTS = 5; const DEFAULT_MCP_STREAMABLE_HTTP_MAX_RESPONSE_BYTES = 16 * 1024 * 1024; @@ -1124,6 +1130,11 @@ interface MCPConnectionParams { /** Result of an MCP `tools/list` request: one page of tools plus an optional pagination cursor. */ type MCPListToolsResult = Awaited>; +export interface MCPToolsSnapshot { + tools: MCPListToolsResult['tools']; + complete: boolean; +} + export class MCPConnection extends EventEmitter { public client: Client; private options: t.MCPOptions; @@ -1148,6 +1159,21 @@ export class MCPConnection extends EventEmitter { private readonly allowedAddresses?: string[] | null; private readonly ephemeralConnection: boolean; private readonly proxyConfig?: MCPProxyConfig; + private toolListChangeGeneration = 0; + private handledToolListChangeGeneration = 0; + private toolListRefreshFailures = 0; + private toolListRefreshPromise: Promise | null = null; + private toolListRefreshRetryTimer: ReturnType | null = null; + private toolListRefreshEpoch = 0; + private toolListRefreshSuspended = false; + private publishedToolListSnapshot: { + epoch: number; + generation: number; + tools: MCPListToolsResult['tools']; + } | null = null; + + private hasConnected = false; + private isDisposed = false; iconPath?: string; timeout?: number; sseReadTimeout?: number; @@ -1750,10 +1776,25 @@ export class MCPConnection extends EventEmitter { this.on('connectionChange', (state: t.ConnectionState) => { this.connectionState = state; if (state === 'connected') { + const isReconnect = this.hasConnected; + this.hasConnected = true; + this.toolListRefreshSuspended = false; this.isReconnecting = false; this.isInitializing = false; this.shouldStopReconnecting = false; this.reconnectAttempts = 0; + if (isReconnect) { + if (this.client.getServerCapabilities()?.tools != null) { + this.toolListChangeGeneration++; + } else { + this.handledToolListChangeGeneration = this.toolListChangeGeneration; + this.toolListRefreshFailures = 0; + this.dispatchToolsChanged([]); + } + } + if (this.handledToolListChangeGeneration < this.toolListChangeGeneration) { + this.startToolListRefresh(); + } /** * // FOR DEBUGGING * // this.client.setRequestHandler(PingRequestSchema, async (request, extra) => { @@ -1765,7 +1806,15 @@ export class MCPConnection extends EventEmitter { * // return {}; * // }); */ - } else if (state === 'error' && !this.isReconnecting && !this.isInitializing) { + } else { + /** Invalidate work started on the previous transport. A reconnect can complete before an + * old `tools/list` request settles, so checking connectionState alone is not sufficient. */ + this.toolListRefreshEpoch++; + this.toolListRefreshSuspended = true; + this.clearToolListRefreshRetry(); + } + + if (state === 'error' && !this.isReconnecting && !this.isInitializing) { this.handleReconnection().catch((error) => { logger.error(`${this.getLogPrefix()} Reconnection handler failed:`, error); }); @@ -1773,6 +1822,7 @@ export class MCPConnection extends EventEmitter { }); this.subscribeToResources(); + this.subscribeToToolListChanges(); } private async handleReconnection(): Promise { @@ -1854,6 +1904,136 @@ export class MCPConnection extends EventEmitter { }); } + /** + * A server that builds tools at runtime tells us so instead of us polling for it. + * + * The spec's list-changed flow is: the server notifies, the client asks for the list again. Until + * this was handled the tool list stayed at whatever it was when the connection came up, so a + * server that adds a tool mid-session was invisible until a restart (#7117). + */ + private subscribeToToolListChanges(): void { + this.client.setNotificationHandler(ToolListChangedNotificationSchema, async () => { + logger.debug(`${this.getLogPrefix()} Server reported a changed tool list`); + await this.refreshToolList(); + }); + } + + /** Queues a live tool-list refresh through the same coalescing path used by notifications. */ + public async refreshToolList(): Promise { + this.toolListChangeGeneration++; + this.clearToolListRefreshRetry(); + this.startToolListRefresh(); + await this.toolListRefreshPromise; + } + + private clearToolListRefreshRetry(): void { + if (this.toolListRefreshRetryTimer) { + clearTimeout(this.toolListRefreshRetryTimer); + this.toolListRefreshRetryTimer = null; + } + } + + private scheduleToolListRefreshRetry(): void { + if ( + this.isDisposed || + this.toolListRefreshSuspended || + this.connectionState !== 'connected' || + this.toolListRefreshRetryTimer + ) { + return; + } + + const delay = Math.min( + TOOL_LIST_REFRESH_RETRY_BASE_MS * Math.pow(2, Math.max(0, this.toolListRefreshFailures - 1)), + TOOL_LIST_REFRESH_RETRY_MAX_MS, + ); + this.toolListRefreshRetryTimer = setTimeout(() => { + this.toolListRefreshRetryTimer = null; + this.startToolListRefresh(); + }, delay); + this.toolListRefreshRetryTimer.unref?.(); + } + + private startToolListRefresh(): void { + if ( + this.isDisposed || + this.toolListRefreshSuspended || + this.connectionState !== 'connected' || + this.toolListRefreshPromise + ) { + return; + } + + this.toolListRefreshPromise = this.refreshChangedTools().finally(() => { + this.toolListRefreshPromise = null; + if ( + !this.toolListRefreshRetryTimer && + this.handledToolListChangeGeneration < this.toolListChangeGeneration + ) { + this.startToolListRefresh(); + } + }); + } + + private async refreshChangedTools(): Promise { + const refreshEpoch = this.toolListRefreshEpoch; + while (this.handledToolListChangeGeneration < this.toolListChangeGeneration) { + const targetGeneration = this.toolListChangeGeneration; + let publicationRevision: string | undefined; + try { + publicationRevision = await reserveMCPToolsChangedRevision({ + serverName: this.serverName, + serverConfig: this.options, + userId: this.userId, + }); + } catch (error) { + this.toolListRefreshFailures++; + logger.error( + `${this.getLogPrefix()} Failed to reserve tool-list publication order:`, + error, + ); + this.scheduleToolListRefreshRetry(); + return; + } + const snapshot = + this.client.getServerCapabilities()?.tools == null + ? { tools: [], complete: true } + : await this.fetchToolsSnapshot(); + if ( + this.toolListRefreshEpoch !== refreshEpoch || + this.toolListRefreshSuspended || + this.connectionState !== 'connected' + ) { + return; + } + if (!snapshot.complete) { + this.toolListRefreshFailures++; + this.scheduleToolListRefreshRetry(); + return; + } + + this.toolListRefreshFailures = 0; + this.handledToolListChangeGeneration = targetGeneration; + this.publishedToolListSnapshot = { + epoch: refreshEpoch, + generation: targetGeneration, + tools: snapshot.tools, + }; + this.dispatchToolsChanged(snapshot.tools, publicationRevision); + } + } + + private dispatchToolsChanged( + tools: MCPListToolsResult['tools'], + publicationRevision?: string, + ): void { + try { + this.emit('toolsChanged', tools, publicationRevision); + } catch (error) { + logger.error(`${this.getLogPrefix()} Failed to dispatch refreshed tools:`, error); + } + } + async connectClient(): Promise { if (this.connectionState === 'connected') { return; @@ -2231,6 +2411,9 @@ export class MCPConnection extends EventEmitter { } public async disconnect(resetCycleTracking = true, forceAgentClose = false): Promise { + this.toolListRefreshEpoch++; + this.toolListRefreshSuspended = true; + this.clearToolListRefreshRetry(); try { if (this.transport) { await this.terminateStreamableSession(); @@ -2253,6 +2436,8 @@ export class MCPConnection extends EventEmitter { /** Permanently tears down a connection that will never be reused. */ public async dispose(): Promise { + this.isDisposed = true; + this.clearToolListRefreshRetry(); this.shouldStopReconnecting = true; this.removeAllListeners(); await this.disconnect(true, true); @@ -2279,6 +2464,15 @@ export class MCPConnection extends EventEmitter { * and the method never throws. */ async fetchTools(): Promise { + return (await this.fetchToolsSnapshot()).tools; + } + + /** + * Fetches a bounded tool snapshot while preserving whether every requested page succeeded. + * Notification refreshes use `complete` to avoid replacing a known-good cache with an empty or + * partial list after a transient `tools/list` failure. + */ + public async fetchToolsSnapshot(): Promise { const maxPages = mcpConfig.TOOLS_LIST_MAX_PAGES; const maxTools = mcpConfig.TOOLS_LIST_MAX_TOOLS; const maxBytes = mcpConfig.TOOLS_LIST_MAX_BYTES; @@ -2297,31 +2491,31 @@ export class MCPConnection extends EventEmitter { ); if (exhaustedBudget != null) { this.warnToolsListBudgetExceeded(exhaustedBudget, allTools.length); - return allTools; + return { tools: allTools, complete: true }; } const remainingMs = deadline - Date.now(); if (remainingMs <= 0) { this.warnToolsListBudgetExceeded('time', allTools.length); - return allTools; + return { tools: allTools, complete: false }; } const result = await this.listToolsPage(cursor, remainingMs); if (result == null) { /** Request failed mid-pagination: return the pages already fetched instead of discarding them. */ - return allTools; + return { tools: allTools, complete: false }; } for (const tool of result.tools) { if (allTools.length >= maxTools) { this.warnToolsListBudgetExceeded('tool count', allTools.length); - return allTools; + return { tools: allTools, complete: true }; } const toolBytes = getApproximateToolBytes(tool); if (totalBytes + toolBytes > maxBytes) { this.warnToolsListBudgetExceeded('size', allTools.length); - return allTools; + return { tools: allTools, complete: true }; } allTools.push(tool); @@ -2330,7 +2524,7 @@ export class MCPConnection extends EventEmitter { const { nextCursor } = result; if (nextCursor == null) { - return allTools; + return { tools: allTools, complete: true }; } const nextPageBudget = getToolsListBudgetExceededReason( @@ -2341,14 +2535,14 @@ export class MCPConnection extends EventEmitter { ); if (nextPageBudget != null) { this.warnToolsListBudgetExceeded(nextPageBudget, allTools.length); - return allTools; + return { tools: allTools, complete: true }; } if (seenCursors.has(nextCursor)) { logger.warn( `${this.getLogPrefix()} MCP server returned a repeated tools/list cursor; stopping pagination after ${page} page(s).`, ); - return allTools; + return { tools: allTools, complete: false }; } seenCursors.add(nextCursor); @@ -2358,7 +2552,51 @@ export class MCPConnection extends EventEmitter { logger.warn( `${this.getLogPrefix()} Reached the tools/list pagination limit of ${maxPages} page(s); some tools may be omitted. Set MCP_TOOLS_LIST_MAX_PAGES higher if this server legitimately exposes more.`, ); - return allTools; + return { tools: allTools, complete: true }; + } + + /** + * Returns a complete snapshot that cannot precede a concurrent `list_changed` refresh. + * If the notification refresh cannot complete, the caller receives an incomplete snapshot + * instead of publishing a request result that may already be stale. + */ + public async fetchOrderedToolsSnapshot(): Promise { + const startEpoch = this.toolListRefreshEpoch; + const startGeneration = this.toolListChangeGeneration; + const snapshot = await this.fetchToolsSnapshot(); + + if ( + startEpoch === this.toolListRefreshEpoch && + startGeneration === this.toolListChangeGeneration + ) { + return snapshot; + } + + while ( + startEpoch === this.toolListRefreshEpoch && + this.handledToolListChangeGeneration < this.toolListChangeGeneration + ) { + this.startToolListRefresh(); + const refresh = this.toolListRefreshPromise; + if (!refresh) { + break; + } + await refresh; + if (this.toolListRefreshRetryTimer) { + break; + } + } + + const published = this.publishedToolListSnapshot; + if ( + published?.epoch === this.toolListRefreshEpoch && + published.generation === this.toolListChangeGeneration && + this.handledToolListChangeGeneration === this.toolListChangeGeneration + ) { + return { tools: published.tools, complete: true }; + } + + return { tools: [], complete: false }; } private warnToolsListBudgetExceeded(reason: string, toolCount: number): void { diff --git a/packages/api/src/mcp/registry/MCPServerInspector.ts b/packages/api/src/mcp/registry/MCPServerInspector.ts index 06281eb0e8..f191bc230e 100644 --- a/packages/api/src/mcp/registry/MCPServerInspector.ts +++ b/packages/api/src/mcp/registry/MCPServerInspector.ts @@ -178,7 +178,11 @@ export class MCPServerInspector { serverName: string, connection: MCPConnection, ): Promise { - const tools = await connection.fetchTools(); + const snapshot = await connection.fetchOrderedToolsSnapshot(); + if (!snapshot.complete) { + throw new Error(`Incomplete tools/list snapshot for MCP server ${serverName}`); + } + const { tools } = snapshot; const toolFunctions: t.LCAvailableTools = {}; /** Model-facing key: must match the runtime instance name, which embeds diff --git a/packages/api/src/mcp/registry/MCPServersRegistry.ts b/packages/api/src/mcp/registry/MCPServersRegistry.ts index fa0bdb4e13..99d3d0c41f 100644 --- a/packages/api/src/mcp/registry/MCPServersRegistry.ts +++ b/packages/api/src/mcp/registry/MCPServersRegistry.ts @@ -298,6 +298,15 @@ export class MCPServersRegistry { return base ? { ...candidate, source: base.source } : candidate; } + /** Returns whether an effective config exactly matches the operator-owned base config. */ + public async isAppServerConfig( + serverName: string, + effectiveConfig: t.ParsedServerConfig, + ): Promise { + const baseConfig = await this.getServerConfig(serverName); + return baseConfig != null && deepEqual(baseConfig, effectiveConfig); + } + /** * Returns the full server config map after merging YAML cache, Config-tier overrides, * and User-DB entries. @@ -491,7 +500,12 @@ export class MCPServersRegistry { return { serverName, config: updatedConfig }; } - public async updateServer( + /** + * Inspects an update without mutating its backing repository. Callers that must + * coordinate an external fence with persistence can prepare first, fence, and + * then commit the returned config. + */ + public async inspectServerUpdate( serverName: string, config: t.MCPOptions, storageLocation: 'CACHE' | 'DB', @@ -532,11 +546,37 @@ export class MCPServersRegistry { } throw new MCPInspectionFailedError(serverName, error as Error); } + return parsedConfig; + } + + /** Persists a previously inspected update without opening a second MCP connection. */ + public async commitServerUpdate( + serverName: string, + parsedConfig: t.ParsedServerConfig, + storageLocation: 'CACHE' | 'DB', + userId?: string, + ): Promise { + const configRepo = this.getConfigRepository(storageLocation); await configRepo.update(serverName, parsedConfig, userId); await this.invalidateServerReadCaches(serverName, userId); return parsedConfig; } + public async updateServer( + serverName: string, + config: t.MCPOptions, + storageLocation: 'CACHE' | 'DB', + userId?: string, + ): Promise { + const parsedConfig = await this.inspectServerUpdate( + serverName, + config, + storageLocation, + userId, + ); + return await this.commitServerUpdate(serverName, parsedConfig, storageLocation, userId); + } + /** * Ensures that config-source MCP servers (from admin Config overrides) are initialized. * Identifies servers in `resolvedMcpConfig` that are not from YAML, lazily initializes diff --git a/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts b/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts index a20261bfb5..7b5d62d03f 100644 --- a/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts +++ b/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts @@ -410,6 +410,9 @@ describe('MCPServerInspector', () => { // Mock server with no tools mockConnection.fetchTools = jest.fn().mockResolvedValue([]); + mockConnection.fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: [], complete: true }); const result = await MCPServerInspector.inspect('test_server', rawConfig, mockConnection); @@ -495,27 +498,30 @@ describe('MCPServerInspector', () => { describe('getToolFunctions()', () => { it('should convert MCP tools to LibreChat tool functions format', async () => { - mockConnection.fetchTools = jest.fn().mockResolvedValue([ - { - name: 'file_read', - description: 'Read a file', - inputSchema: { - type: 'object', - properties: { path: { type: 'string' } }, - }, - }, - { - name: 'file_write', - description: 'Write a file', - inputSchema: { - type: 'object', - properties: { - path: { type: 'string' }, - content: { type: 'string' }, + mockConnection.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + complete: true, + tools: [ + { + name: 'file_read', + description: 'Read a file', + inputSchema: { + type: 'object', + properties: { path: { type: 'string' } }, }, }, - }, - ]); + { + name: 'file_write', + description: 'Write a file', + inputSchema: { + type: 'object', + properties: { + path: { type: 'string' }, + content: { type: 'string' }, + }, + }, + }, + ], + }); const result = await MCPServerInspector.getToolFunctions('my_server', mockConnection); @@ -549,7 +555,9 @@ describe('MCPServerInspector', () => { }); it('should handle empty tools list', async () => { - mockConnection.fetchTools = jest.fn().mockResolvedValue([]); + mockConnection.fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: [], complete: true }); const result = await MCPServerInspector.getToolFunctions('my_server', mockConnection); @@ -557,13 +565,16 @@ describe('MCPServerInspector', () => { }); it('builds keys with the normalized server name (model-facing contract)', async () => { - mockConnection.fetchTools = jest.fn().mockResolvedValue([ - { - name: 'file_read', - description: 'Read a file', - inputSchema: { type: 'object', properties: {} }, - }, - ]); + mockConnection.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + complete: true, + tools: [ + { + name: 'file_read', + description: 'Read a file', + inputSchema: { type: 'object', properties: {} }, + }, + ], + }); const result = await MCPServerInspector.getToolFunctions('My Server', mockConnection); @@ -571,5 +582,16 @@ describe('MCPServerInspector', () => { expect(Object.keys(result)).toEqual([key]); expect(result[key]['function'].name).toBe(key); }); + + it('rejects an incomplete snapshot before it can replace cached tools', async () => { + mockConnection.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + tools: [{ name: 'partial', inputSchema: { type: 'object' } }], + complete: false, + }); + + await expect( + MCPServerInspector.getToolFunctions('my_server', mockConnection), + ).rejects.toThrow('Incomplete tools/list snapshot for MCP server my_server'); + }); }); }); diff --git a/packages/api/src/mcp/registry/__tests__/MCPServersRegistry.test.ts b/packages/api/src/mcp/registry/__tests__/MCPServersRegistry.test.ts index cf0b18a739..ef1ec81e75 100644 --- a/packages/api/src/mcp/registry/__tests__/MCPServersRegistry.test.ts +++ b/packages/api/src/mcp/registry/__tests__/MCPServersRegistry.test.ts @@ -163,6 +163,26 @@ describe('MCPServersRegistry', () => { }); }); + describe('isAppServerConfig', () => { + it('rejects a same-name tenant override that inherited the YAML source tag', async () => { + const baseConfig = { + ...testParsedConfig, + source: 'yaml' as const, + url: 'https://base.example.com/mcp', + type: 'streamable-http' as const, + }; + await registry['cacheConfigsRepo'].add('shared', baseConfig); + + await expect(registry.isAppServerConfig('shared', baseConfig)).resolves.toBe(true); + await expect( + registry.isAppServerConfig('shared', { + ...baseConfig, + url: 'https://tenant.example.com/mcp', + }), + ).resolves.toBe(false); + }); + }); + describe('addServer', () => { it('should pass user source to inspector before storing DB servers', async () => { const inspectSpy = jest.spyOn(MCPServerInspector, 'inspect'); @@ -383,6 +403,29 @@ describe('MCPServersRegistry', () => { } }); + it('separates update inspection from persistence', async () => { + await registry.addServer('cache_server', testParsedConfig, 'CACHE'); + const updatedConfig = { ...testParsedConfig, command: 'python' } as t.ParsedServerConfig; + + const inspected = await registry.inspectServerUpdate( + 'cache_server', + updatedConfig, + 'CACHE', + ); + + const beforeCommit = await registry['cacheConfigsRepo'].get('cache_server'); + expect(beforeCommit && 'command' in beforeCommit ? beforeCommit.command : undefined).toBe( + 'node', + ); + + await registry.commitServerUpdate('cache_server', inspected, 'CACHE'); + + const afterCommit = await registry['cacheConfigsRepo'].get('cache_server'); + expect(afterCommit && 'command' in afterCommit ? afterCommit.command : undefined).toBe( + 'python', + ); + }); + it('should route removeServer to cache repository', async () => { await registry.addServer('cache_server', testParsedConfig, 'CACHE'); // Verify server exists in underlying cache repository (not via getServerConfig to avoid populating read-through cache) diff --git a/packages/api/src/mcp/registry/__tests__/mcpConnectionsMock.helper.ts b/packages/api/src/mcp/registry/__tests__/mcpConnectionsMock.helper.ts index 5e028c4c81..225af4916a 100644 --- a/packages/api/src/mcp/registry/__tests__/mcpConnectionsMock.helper.ts +++ b/packages/api/src/mcp/registry/__tests__/mcpConnectionsMock.helper.ts @@ -32,7 +32,10 @@ export function createMockConnection(serverName: string): jest.Mocked; } diff --git a/packages/api/src/mcp/tools.spec.ts b/packages/api/src/mcp/tools.spec.ts index 5e036f483a..9f314cff91 100644 --- a/packages/api/src/mcp/tools.spec.ts +++ b/packages/api/src/mcp/tools.spec.ts @@ -1,6 +1,7 @@ -import { Constants } from 'librechat-data-provider'; +import { Constants, normalizeServerName } from 'librechat-data-provider'; import type { LCAvailableTools, ParsedServerConfig } from './types'; -import type { MCPToolInput, MCPToolCacheDeps } from './tools'; +import type { MCPToolCacheDeps, MCPToolInput } from './tools'; +import { getMCPAppToolsPublicationGeneration } from './toolsChanged'; import { createMCPToolCacheService } from './tools'; const requestScopedConfig: ParsedServerConfig = { @@ -15,413 +16,479 @@ const cacheableConfig: ParsedServerConfig = { source: 'yaml', }; +const tenantConfig: ParsedServerConfig = { + ...cacheableConfig, + source: 'config', +}; + +const toolName = (name: string, server: string) => + `${name}${Constants.mcp_delimiter}${normalizeServerName(server)}`; + +const makeTool = (name: string) => ({ + type: 'function' as const, + ['function']: { name, description: '', parameters: { type: 'object' as const, properties: {} } }, +}); + function createMockDeps(overrides: Partial = {}): MCPToolCacheDeps { return { getCachedTools: jest.fn().mockResolvedValue(null), setCachedTools: jest.fn().mockResolvedValue(true), + getCachedAppServerTools: jest.fn().mockResolvedValue(null), + setCachedAppServerTools: jest.fn().mockResolvedValue(true), getServerConfig: jest.fn().mockResolvedValue(undefined), ...overrides, }; } +function createSharedCacheDeps(params: { + config: ParsedServerConfig; + app?: boolean; + appCache?: Map; + userCache?: Map; +}): MCPToolCacheDeps { + const { config, app = true } = params; + const appCache = params.appCache ?? new Map(); + const userCache = params.userCache ?? new Map(); + const appKey = (serverName: string, generation: string) => + JSON.stringify([serverName, generation]); + const userKey = (userId: string, serverName: string, generation?: string) => + JSON.stringify([userId, serverName, generation]); + + return { + getCachedTools: jest.fn(async ({ userId, serverName, configGeneration } = {}) => { + if (!userId || !serverName) { + return null; + } + return userCache.get(userKey(userId, serverName, configGeneration)) ?? null; + }), + setCachedTools: jest.fn(async (tools, { userId, serverName, configGeneration } = {}) => { + if (userId && serverName) { + userCache.set(userKey(userId, serverName, configGeneration), tools); + } + return true; + }), + setCachedToolsIfCurrent: jest.fn(async (tools, { userId, serverName, configGeneration }) => { + userCache.set(userKey(userId, serverName, configGeneration), tools); + return true; + }), + getCachedAppServerTools: jest.fn( + async (serverName, generation) => appCache.get(appKey(serverName, generation)) ?? null, + ), + setCachedAppServerTools: jest.fn(async (serverName, generation, tools) => { + appCache.set(appKey(serverName, generation), tools); + return true; + }), + getServerConfig: jest.fn().mockResolvedValue(config), + getAllServerConfigs: jest.fn().mockResolvedValue(app ? { dynamic: config } : {}), + isAppServerConfig: jest.fn().mockResolvedValue(app), + }; +} + describe('createMCPToolCacheService', () => { - describe('updateMCPServerTools', () => { - it('returns empty object for null tools', async () => { - const deps = createMockDeps(); - const { updateMCPServerTools } = createMCPToolCacheService(deps); + describe('configuration-addressed app catalogs', () => { + it('restores the static catalog without discovering app server configs', async () => { + const staticTools = { builtin: makeTool('builtin') }; + const updateCachedGlobalTools = jest.fn(async (update) => update({})); + const getAllServerConfigs = jest.fn().mockResolvedValue({ alpha: cacheableConfig }); + const service = createMCPToolCacheService( + createMockDeps({ updateCachedGlobalTools, getAllServerConfigs }), + ); - const result = await updateMCPServerTools({ userId: 'u1', serverName: 'srv', tools: null }); + await service.syncStaticTools(staticTools); - expect(result).toEqual({}); - expect(deps.setCachedTools).not.toHaveBeenCalled(); + expect(updateCachedGlobalTools).toHaveBeenCalledTimes(1); + expect(updateCachedGlobalTools.mock.calls[0][0]({})).toEqual(staticTools); + expect(getAllServerConfigs).not.toHaveBeenCalled(); }); - it('replaces a stale cache entry when the server returns an empty tools array', async () => { - const deps = createMockDeps(); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - - const result = await updateMCPServerTools({ userId: 'u1', serverName: 'srv', tools: [] }); - - expect(result).toEqual({}); - expect(deps.setCachedTools).toHaveBeenCalledWith({}, { userId: 'u1', serverName: 'srv' }); - }); - - it('builds MODEL-FACING keys with the normalized server name, store keyed raw', async () => { - /** Tool keys become builder tool ids, agent.tools entries, tool_options - * keys, and definition names, and must equal the runtime instance name, - * which embeds `normalizeServerName(serverName)`. */ - const deps = createMockDeps(); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - const tools: MCPToolInput[] = [ - { name: 'search', description: 'Search', inputSchema: { type: 'object', properties: {} } }, - ]; - - const result = await updateMCPServerTools({ - userId: 'u1', - serverName: 'Connector: Company', - tools, - }); - - const expectedKey = `search${Constants.mcp_delimiter}Connector__Company`; - expect(result[expectedKey]).toBeDefined(); - expect(result[expectedKey]['function'].name).toBe(expectedKey); - expect(deps.setCachedTools).toHaveBeenCalledWith(result, { - userId: 'u1', - serverName: 'Connector: Company', - }); - }); - - it('constructs tool names with mcp_delimiter and caches them', async () => { - const deps = createMockDeps(); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - const tools: MCPToolInput[] = [ - { - name: 'search', - description: 'Search docs', - inputSchema: { type: 'object', properties: {} }, - }, - ]; - - const result = await updateMCPServerTools({ userId: 'u1', serverName: 'brave', tools }); - - const expectedKey = `search${Constants.mcp_delimiter}brave`; - expect(result[expectedKey]).toBeDefined(); - expect(result[expectedKey].type).toBe('function'); - expect(result[expectedKey]['function'].name).toBe(expectedKey); - expect(result[expectedKey]['function'].description).toBe('Search docs'); - expect(deps.setCachedTools).toHaveBeenCalledWith(result, { - userId: 'u1', - serverName: 'brave', - }); - }); - - it('builds tool names without caching when the resolved config is request-scoped', async () => { + it('removes legacy MCP entries from the shared global catalog during rollout', async () => { + const alphaConfig = { ...cacheableConfig, toolFunctions: {} }; + const builtin = 'code_interpreter'; + const staticDelimiterTool = `get${Constants.mcp_delimiter}status`; + let globalTools: LCAvailableTools = {}; + const staticTools = { + [builtin]: makeTool(builtin), + [staticDelimiterTool]: makeTool(staticDelimiterTool), + }; const deps = createMockDeps({ - getServerConfig: jest.fn().mockResolvedValue(requestScopedConfig), - }); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - const tools: MCPToolInput[] = [ - { - name: 'search', - description: 'Search request-scoped docs', - inputSchema: { type: 'object', properties: {} }, - }, - ]; - - const result = await updateMCPServerTools({ - userId: 'u1', - serverName: 'body-scoped', - tools, + getAllServerConfigs: jest.fn().mockResolvedValue({ alpha: alphaConfig }), + updateCachedGlobalTools: jest.fn(async (update) => { + globalTools = update(globalTools); + }), }); - const expectedKey = `search${Constants.mcp_delimiter}body-scoped`; - expect(result[expectedKey]).toBeDefined(); - expect(deps.getServerConfig).toHaveBeenCalledWith('body-scoped', 'u1'); - expect(deps.setCachedTools).not.toHaveBeenCalled(); + await createMCPToolCacheService(deps).mergeAppTools({}, staticTools); + + expect(globalTools).toEqual(staticTools); + expect(deps.updateCachedGlobalTools).toHaveBeenCalledTimes(1); }); - it('uses a provided serverConfig without calling the resolver', async () => { - const deps = createMockDeps(); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - const tools: MCPToolInput[] = [{ name: 'search' }]; - - await updateMCPServerTools({ - userId: 'u1', - serverName: 'body-scoped', - tools, - serverConfig: requestScopedConfig, - }); - - expect(deps.getServerConfig).not.toHaveBeenCalled(); - expect(deps.setCachedTools).not.toHaveBeenCalled(); - }); - - it('fails open and caches when config resolution throws', async () => { - const deps = createMockDeps({ - getServerConfig: jest.fn().mockRejectedValue(new Error('registry not initialized')), - }); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - const tools: MCPToolInput[] = [{ name: 'search' }]; - - await updateMCPServerTools({ userId: 'u1', serverName: 'srv', tools }); - - expect(deps.setCachedTools).toHaveBeenCalled(); - }); - - it('propagates setCachedTools errors', async () => { - const deps = createMockDeps({ - setCachedTools: jest.fn().mockRejectedValue(new Error('Redis down')), - }); - const { updateMCPServerTools } = createMCPToolCacheService(deps); - const tools: MCPToolInput[] = [{ name: 'tool1' }]; + it('writes an authoritative empty snapshot under the publishing config generation', async () => { + const setCachedAppServerTools = jest.fn().mockResolvedValue(true); + const deps = createMockDeps({ setCachedAppServerTools }); + const service = createMCPToolCacheService(deps); + const generation = getMCPAppToolsPublicationGeneration(cacheableConfig); await expect( - updateMCPServerTools({ userId: 'u1', serverName: 'srv', tools }), + service.replaceAppServerTools({ + serverName: 'dynamic', + serverTools: {}, + publicationGeneration: generation, + publicationRevision: '1', + }), + ).resolves.toBe(true); + + expect(setCachedAppServerTools).toHaveBeenCalledWith('dynamic', generation, {}, '1'); + expect(deps.setCachedTools).not.toHaveBeenCalled(); + }); + + it('propagates a rejected app-slice write', async () => { + const deps = createMockDeps({ + setCachedAppServerTools: jest.fn().mockRejectedValue(new Error('Redis down')), + }); + + await expect( + createMCPToolCacheService(deps).replaceAppServerTools({ + serverName: 'dynamic', + serverTools: {}, + publicationGeneration: 'config-generation', + publicationRevision: '1', + }), ).rejects.toThrow('Redis down'); }); - }); - describe('mergeAppTools', () => { - it('no-ops when appTools is empty', async () => { + it('does not publish a live app snapshot without pre-fetch ordering', async () => { const deps = createMockDeps(); - const { mergeAppTools } = createMCPToolCacheService(deps); - - await mergeAppTools({}); - - expect(deps.getCachedTools).not.toHaveBeenCalled(); - expect(deps.setCachedTools).not.toHaveBeenCalled(); - }); - - it('merges app tools with existing cached tools', async () => { - const existing: LCAvailableTools = { - old: { - type: 'function', - ['function']: { - name: 'old', - description: '', - parameters: { type: 'object', properties: {} }, - }, - }, - }; - const deps = createMockDeps({ getCachedTools: jest.fn().mockResolvedValue(existing) }); - const { mergeAppTools } = createMCPToolCacheService(deps); - const appTools: LCAvailableTools = { - new: { - type: 'function', - ['function']: { - name: 'new', - description: '', - parameters: { type: 'object', properties: {} }, - }, - }, - }; - - await mergeAppTools(appTools); - - expect(deps.setCachedTools).toHaveBeenCalledWith( - expect.objectContaining({ old: existing.old, new: appTools.new }), - ); - }); - - it('handles null cache (cold start) by defaulting to empty', async () => { - const deps = createMockDeps({ getCachedTools: jest.fn().mockResolvedValue(null) }); - const { mergeAppTools } = createMCPToolCacheService(deps); - const appTools: LCAvailableTools = { - tool: { - type: 'function', - ['function']: { - name: 'tool', - description: '', - parameters: { type: 'object', properties: {} }, - }, - }, - }; - - await mergeAppTools(appTools); - - expect(deps.setCachedTools).toHaveBeenCalledWith( - expect.objectContaining({ tool: appTools.tool }), - ); - }); - - it('propagates getCachedTools errors', async () => { - const deps = createMockDeps({ - getCachedTools: jest.fn().mockRejectedValue(new Error('cache read failed')), - }); - const { mergeAppTools } = createMCPToolCacheService(deps); await expect( - mergeAppTools({ - t: { - type: 'function', - ['function']: { - name: 't', - description: '', - parameters: { type: 'object', properties: {} }, - }, - }, + createMCPToolCacheService(deps).replaceAppServerTools({ + serverName: 'dynamic', + serverTools: {}, + publicationGeneration: 'config-generation', }), - ).rejects.toThrow('cache read failed'); - }); - }); + ).resolves.toBe(false); - describe('cacheMCPServerTools', () => { - const serverTools: LCAvailableTools = { - tool: { - type: 'function', - ['function']: { - name: 'tool', - description: '', - parameters: { type: 'object', properties: {} }, - }, - }, - }; - - it('no-ops when serverTools is empty', async () => { - const deps = createMockDeps(); - const { cacheMCPServerTools } = createMCPToolCacheService(deps); - - await cacheMCPServerTools({ userId: 'u1', serverName: 'srv', serverTools: {} }); - - expect(deps.setCachedTools).not.toHaveBeenCalled(); + expect(deps.setCachedAppServerTools).not.toHaveBeenCalled(); }); - it('caches server tools with userId and serverName', async () => { - const deps = createMockDeps(); - const { cacheMCPServerTools } = createMCPToolCacheService(deps); - - await cacheMCPServerTools({ userId: 'u1', serverName: 'brave', serverTools }); - - expect(deps.setCachedTools).toHaveBeenCalledWith(serverTools, { - userId: 'u1', - serverName: 'brave', - }); - }); - - it('skips caching for request-scoped servers', async () => { + it('rejects a tool boundary owned by another app server', async () => { + const shadowed = toolName('search', 'foo_bar'); const deps = createMockDeps({ - getServerConfig: jest.fn().mockResolvedValue(requestScopedConfig), + getAllServerConfigs: jest.fn().mockResolvedValue({ + 'foo bar': cacheableConfig, + foo_bar: cacheableConfig, + }), }); - const { cacheMCPServerTools } = createMCPToolCacheService(deps); - - await cacheMCPServerTools({ userId: 'u1', serverName: 'body-scoped', serverTools }); - - expect(deps.setCachedTools).not.toHaveBeenCalled(); - }); - - it('propagates setCachedTools errors', async () => { - const deps = createMockDeps({ - setCachedTools: jest.fn().mockRejectedValue(new Error('write failed')), - }); - const { cacheMCPServerTools } = createMCPToolCacheService(deps); await expect( - cacheMCPServerTools({ userId: 'u1', serverName: 'srv', serverTools }), - ).rejects.toThrow('write failed'); + createMCPToolCacheService(deps).replaceAppServerTools({ + serverName: 'foo bar', + serverTools: { [shadowed]: makeTool(shadowed) }, + publicationGeneration: 'config-generation', + publicationRevision: '1', + }), + ).rejects.toThrow('belongs to app server foo_bar'); + }); + + it('isolates old and new replicas instead of electing the first publisher', async () => { + const appCache = new Map(); + const oldConfig = cacheableConfig; + const newConfig = { ...cacheableConfig, url: 'https://mcp.example.com/v2/mcp' }; + const oldService = createMCPToolCacheService( + createSharedCacheDeps({ config: oldConfig, appCache }), + ); + const newService = createMCPToolCacheService( + createSharedCacheDeps({ config: newConfig, appCache }), + ); + const oldTools = { [toolName('old', 'dynamic')]: makeTool(toolName('old', 'dynamic')) }; + const newTools = { [toolName('new', 'dynamic')]: makeTool(toolName('new', 'dynamic')) }; + + await newService.replaceAppServerTools({ + serverName: 'dynamic', + serverTools: newTools, + publicationGeneration: getMCPAppToolsPublicationGeneration(newConfig), + publicationRevision: '1', + }); + await oldService.replaceAppServerTools({ + serverName: 'dynamic', + serverTools: oldTools, + publicationGeneration: getMCPAppToolsPublicationGeneration(oldConfig), + publicationRevision: '1', + }); + + await expect(newService.getMCPServerTools('user', 'dynamic')).resolves.toEqual(newTools); + await expect(oldService.getMCPServerTools('user', 'dynamic')).resolves.toEqual(oldTools); + expect(appCache).toHaveProperty('size', 2); + }); + + it('recovers safely after shared cache loss even if a stale replica publishes first', async () => { + const appCache = new Map(); + const oldConfig = cacheableConfig; + const newConfig = { ...cacheableConfig, url: 'https://mcp.example.com/v2/mcp' }; + const oldService = createMCPToolCacheService( + createSharedCacheDeps({ config: oldConfig, appCache }), + ); + const newService = createMCPToolCacheService( + createSharedCacheDeps({ config: newConfig, appCache }), + ); + const stale = { [toolName('stale', 'dynamic')]: makeTool(toolName('stale', 'dynamic')) }; + const current = { + [toolName('current', 'dynamic')]: makeTool(toolName('current', 'dynamic')), + }; + + appCache.clear(); + await oldService.replaceAppServerTools({ + serverName: 'dynamic', + serverTools: stale, + publicationGeneration: getMCPAppToolsPublicationGeneration(oldConfig), + publicationRevision: '1', + }); + await expect(newService.getMCPServerTools('user', 'dynamic')).resolves.toBeNull(); + + await newService.replaceAppServerTools({ + serverName: 'dynamic', + serverTools: current, + publicationGeneration: getMCPAppToolsPublicationGeneration(newConfig), + publicationRevision: '1', + }); + await expect(newService.getMCPServerTools('user', 'dynamic')).resolves.toEqual(current); + }); + + it('splits startup tools into independently addressed server slices', async () => { + const alphaConfig = { ...cacheableConfig, toolFunctions: {} }; + const betaConfig = { ...cacheableConfig, url: 'https://beta.example.com', toolFunctions: {} }; + const setCachedAppServerTools = jest.fn().mockResolvedValue(true); + const deps = createMockDeps({ + getAllServerConfigs: jest.fn().mockResolvedValue({ alpha: alphaConfig, beta: betaConfig }), + setCachedAppServerTools, + }); + const alpha = toolName('one', 'alpha'); + const beta = toolName('two', 'beta'); + + await createMCPToolCacheService(deps).mergeAppTools( + { + [alpha]: makeTool(alpha), + [beta]: makeTool(beta), + }, + {}, + ); + + expect(setCachedAppServerTools).toHaveBeenCalledWith( + 'alpha', + getMCPAppToolsPublicationGeneration(alphaConfig), + { [alpha]: makeTool(alpha) }, + ); + expect(setCachedAppServerTools).toHaveBeenCalledWith( + 'beta', + getMCPAppToolsPublicationGeneration(betaConfig), + { [beta]: makeTool(beta) }, + ); + expect(deps.setCachedTools).not.toHaveBeenCalled(); + }); + + it('does not replace a known-good slice when startup inspection is incomplete', async () => { + const setCachedAppServerTools = jest.fn().mockResolvedValue(true); + const deps = createMockDeps({ + getAllServerConfigs: jest.fn().mockResolvedValue({ + complete: { ...cacheableConfig, toolFunctions: {} }, + incomplete: { ...cacheableConfig, toolFunctions: undefined }, + }), + setCachedAppServerTools, + }); + + await createMCPToolCacheService(deps).mergeAppTools({}, {}); + + expect(setCachedAppServerTools).toHaveBeenCalledTimes(1); + expect(setCachedAppServerTools).toHaveBeenCalledWith('complete', expect.any(String), {}); }); }); - describe('getMCPServerTools', () => { - const cachedTools: LCAvailableTools = { - tool: { - type: 'function', - ['function']: { - name: 'tool', - description: '', - parameters: { type: 'object', properties: {} }, - }, - }, - }; - - it('returns cached tools for cacheable servers', async () => { + describe('user catalog fencing', () => { + it('passes both connection and config generations to the guarded write', async () => { + const setCachedToolsIfCurrent = jest.fn().mockResolvedValue(true); const deps = createMockDeps({ - getCachedTools: jest.fn().mockResolvedValue(cachedTools), - getServerConfig: jest.fn().mockResolvedValue(cacheableConfig), + getServerConfig: jest.fn().mockResolvedValue(tenantConfig), + getAllServerConfigs: jest.fn().mockResolvedValue({}), + setCachedToolsIfCurrent, }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - const result = await getMCPServerTools('u1', 'brave'); + await createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'tenant', + tools: [{ name: 'search' }], + publicationGeneration: 'connection-generation', + }); - expect(result).toEqual(cachedTools); - expect(deps.getCachedTools).toHaveBeenCalledWith({ userId: 'u1', serverName: 'brave' }); + expect(setCachedToolsIfCurrent).toHaveBeenCalledWith(expect.any(Object), { + userId: 'u1', + serverName: 'tenant', + configGeneration: getMCPAppToolsPublicationGeneration(tenantConfig), + publicationGeneration: 'connection-generation', + }); }); - it('treats a cached empty catalog as a miss so discovery remains enabled', async () => { + it('does not return definitions rejected by the publication-generation fence', async () => { const deps = createMockDeps({ - getCachedTools: jest.fn().mockResolvedValue({}), - getServerConfig: jest.fn().mockResolvedValue(cacheableConfig), + getServerConfig: jest.fn().mockResolvedValue(tenantConfig), + getAllServerConfigs: jest.fn().mockResolvedValue({}), + setCachedToolsIfCurrent: jest.fn().mockResolvedValue(false), }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - const result = await getMCPServerTools('u1', 'brave'); - - expect(result).toBeNull(); + await expect( + createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'tenant', + tools: [{ name: 'stale' }], + publicationGeneration: 'stale-generation', + }), + ).resolves.toBeNull(); }); - it('heals stale raw-keyed cache entries to the normalized key format at read time', async () => { - /** Entries written before keys embedded the normalized server name would - * otherwise make the server's tools vanish for up to the cache TTL — - * the definitions-only loader treats this map as authoritative and - * never reconnects on a per-key miss. */ - const staleTools: LCAvailableTools = { - [`search${Constants.mcp_delimiter}Connector: Company`]: { - type: 'function', - ['function']: { - name: `search${Constants.mcp_delimiter}Connector: Company`, - description: 'Search', - parameters: { type: 'object', properties: {} }, - }, - }, - }; - const deps = createMockDeps({ - getCachedTools: jest.fn().mockResolvedValue(staleTools), - getServerConfig: jest.fn().mockResolvedValue(cacheableConfig), + it('keeps a late old-config publication invisible to current readers', async () => { + const userCache = new Map(); + const oldConfig = tenantConfig; + const newConfig = { ...tenantConfig, url: 'https://mcp.example.com/v2/mcp' }; + const oldService = createMCPToolCacheService( + createSharedCacheDeps({ config: oldConfig, app: false, userCache }), + ); + const newService = createMCPToolCacheService( + createSharedCacheDeps({ config: newConfig, app: false, userCache }), + ); + const oldTools = { [toolName('old', 'dynamic')]: makeTool(toolName('old', 'dynamic')) }; + const newTools = { [toolName('new', 'dynamic')]: makeTool(toolName('new', 'dynamic')) }; + + await newService.cacheMCPServerTools({ + userId: 'u1', + serverName: 'dynamic', + serverTools: newTools, + publicationGeneration: 'new-connection', + }); + await oldService.cacheMCPServerTools({ + userId: 'u1', + serverName: 'dynamic', + serverTools: oldTools, + publicationGeneration: 'old-connection', }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - const result = await getMCPServerTools('u1', 'Connector: Company'); - - const healedKey = `search${Constants.mcp_delimiter}Connector__Company`; - expect(Object.keys(result ?? {})).toEqual([healedKey]); - expect(result?.[healedKey]['function'].name).toBe(healedKey); + await expect(newService.getMCPServerTools('u1', 'dynamic')).resolves.toEqual(newTools); }); - it('returns the same reference for safe server names (no heal pass)', async () => { + it('does not fall back to an unfenced write when a guard is configured', async () => { + const setCachedToolsIfCurrent = jest.fn().mockResolvedValue(true); const deps = createMockDeps({ - getCachedTools: jest.fn().mockResolvedValue(cachedTools), - getServerConfig: jest.fn().mockResolvedValue(cacheableConfig), + getServerConfig: jest.fn().mockResolvedValue(tenantConfig), + getAllServerConfigs: jest.fn().mockResolvedValue({}), + setCachedToolsIfCurrent, }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - const result = await getMCPServerTools('u1', 'brave'); + await createMCPToolCacheService(deps).cacheMCPServerTools({ + userId: 'u1', + serverName: 'tenant', + serverTools: {}, + }); - expect(result).toBe(cachedTools); + expect(setCachedToolsIfCurrent).not.toHaveBeenCalled(); + expect(deps.setCachedTools).not.toHaveBeenCalled(); + }); + }); + + describe('tool construction and reads', () => { + it('returns empty for a null tool list without caching', async () => { + const deps = createMockDeps(); + const result = await createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'srv', + tools: null, + }); + + expect(result).toEqual({}); + expect(deps.setCachedTools).not.toHaveBeenCalled(); }); - it('returns null for request-scoped servers without reading the cache', async () => { + it('builds model-facing names with the normalized server name', async () => { + const deps = createMockDeps(); + const tools: MCPToolInput[] = [{ name: 'search', description: 'Search' }]; + const result = await createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'Connector: Company', + tools, + }); + const expected = toolName('search', 'Connector: Company'); + + expect(result?.[expected]?.['function'].name).toBe(expected); + expect(deps.setCachedTools).toHaveBeenCalledWith(result, { + userId: 'u1', + serverName: 'Connector: Company', + configGeneration: undefined, + }); + }); + + it('builds request-scoped tools without caching them', async () => { const deps = createMockDeps({ - getCachedTools: jest.fn().mockResolvedValue(cachedTools), getServerConfig: jest.fn().mockResolvedValue(requestScopedConfig), }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - - const result = await getMCPServerTools('u1', 'body-scoped'); - - expect(result).toBeNull(); - expect(deps.getCachedTools).not.toHaveBeenCalled(); - }); - - it('uses a provided serverConfig without calling the resolver', async () => { - const deps = createMockDeps({ - getCachedTools: jest.fn().mockResolvedValue(cachedTools), + const result = await createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'body-scoped', + tools: [{ name: 'search' }], }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - const result = await getMCPServerTools('u1', 'body-scoped', requestScopedConfig); - - expect(result).toBeNull(); - expect(deps.getServerConfig).not.toHaveBeenCalled(); - expect(deps.getCachedTools).not.toHaveBeenCalled(); + expect(result?.[toolName('search', 'body-scoped')]).toBeDefined(); + expect(deps.setCachedTools).not.toHaveBeenCalled(); + expect(deps.setCachedAppServerTools).not.toHaveBeenCalled(); }); - it('returns null when the cache is empty', async () => { + it('treats a missing app slice differently from an authoritative empty slice', async () => { + const getCachedAppServerTools = jest + .fn() + .mockResolvedValueOnce(null) + .mockResolvedValueOnce({}); + const deps = createMockDeps({ + getServerConfig: jest.fn().mockResolvedValue(cacheableConfig), + getAllServerConfigs: jest.fn().mockResolvedValue({ dynamic: cacheableConfig }), + getCachedAppServerTools, + }); + const service = createMCPToolCacheService(deps); + + await expect(service.getMCPServerTools('u1', 'dynamic')).resolves.toBeNull(); + await expect(service.getMCPServerTools('u1', 'dynamic')).resolves.toEqual({}); + }); + + it('heals raw server names in a configuration-addressed user slice', async () => { + const staleName = `search${Constants.mcp_delimiter}Connector: Company`; + const staleTools = { [staleName]: makeTool(staleName) }; + const deps = createMockDeps({ + getServerConfig: jest.fn().mockResolvedValue(tenantConfig), + getAllServerConfigs: jest.fn().mockResolvedValue({}), + getCachedTools: jest.fn().mockResolvedValue(staleTools), + }); + + const result = await createMCPToolCacheService(deps).getMCPServerTools( + 'u1', + 'Connector: Company', + ); + const healed = toolName('search', 'Connector: Company'); + + expect(Object.keys(result ?? {})).toEqual([healed]); + expect(result?.[healed]['function'].name).toBe(healed); + }); + + it('returns null without reading cache for request-scoped servers', async () => { const deps = createMockDeps(); - const { getMCPServerTools } = createMCPToolCacheService(deps); - - const result = await getMCPServerTools('u1', 'brave'); - - expect(result).toBeNull(); + await expect( + createMCPToolCacheService(deps).getMCPServerTools('u1', 'body-scoped', requestScopedConfig), + ).resolves.toBeNull(); + expect(deps.getCachedTools).not.toHaveBeenCalled(); + expect(deps.getCachedAppServerTools).not.toHaveBeenCalled(); }); - it('returns null instead of throwing when the cache read fails', async () => { + it('returns null when a cache read fails', async () => { const deps = createMockDeps({ getCachedTools: jest.fn().mockRejectedValue(new Error('cache unavailable')), }); - const { getMCPServerTools } = createMCPToolCacheService(deps); - - const result = await getMCPServerTools('u1', 'brave'); - - expect(result).toBeNull(); + await expect( + createMCPToolCacheService(deps).getMCPServerTools('u1', 'server'), + ).resolves.toBeNull(); }); }); }); diff --git a/packages/api/src/mcp/tools.ts b/packages/api/src/mcp/tools.ts index 0cb3255789..35dfd39fcd 100644 --- a/packages/api/src/mcp/tools.ts +++ b/packages/api/src/mcp/tools.ts @@ -1,40 +1,75 @@ import { logger } from '@librechat/data-schemas'; -import { Constants, normalizeServerName } from 'librechat-data-provider'; +import { Constants, buildServerNameAliases, normalizeServerName } from 'librechat-data-provider'; +import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import type { JsonSchemaType } from '@librechat/agents'; import type { LCAvailableTools, LCFunctionTool, ParsedServerConfig } from './types'; -import { requiresEphemeralUserConnection } from './utils'; +import { canUseAppConnection, requiresEphemeralUserConnection } from './utils'; +import { getMCPAppToolsPublicationGeneration } from './toolsChanged'; +import { normalizeJsonSchema, resolveJsonSchemaRefs } from './zod'; -export interface MCPToolInput { - name: string; - description?: string; - inputSchema?: JsonSchemaType; -} +export type MCPToolInput = Pick & Partial>; export interface MCPToolCacheDeps { getCachedTools: (options?: { userId?: string; serverName?: string; + configGeneration?: string; }) => Promise; + updateCachedGlobalTools?: ( + update: (tools: LCAvailableTools) => LCAvailableTools, + ) => Promise; setCachedTools: ( tools: LCAvailableTools, - options?: { userId?: string; serverName?: string }, + options?: { userId?: string; serverName?: string; configGeneration?: string }, + ) => Promise; + setCachedToolsIfCurrent?: ( + tools: LCAvailableTools, + options: { + userId: string; + serverName: string; + configGeneration: string; + publicationGeneration: string; + }, + ) => Promise; + getCachedAppServerTools: ( + serverName: string, + configGeneration: string, + ) => Promise; + setCachedAppServerTools: ( + serverName: string, + configGeneration: string, + tools: LCAvailableTools, + publicationRevision?: string, ) => Promise; getServerConfig: (serverName: string, userId?: string) => Promise; + getAllServerConfigs?: () => Promise>; + isAppServerConfig?: (serverName: string, effectiveConfig: ParsedServerConfig) => Promise; } export interface MCPToolCacheService { updateMCPServerTools: (params: { - userId: string; + userId?: string; serverName: string; tools: MCPToolInput[] | null; serverConfig?: ParsedServerConfig; - }) => Promise; - mergeAppTools: (appTools: LCAvailableTools) => Promise; + publicationGeneration?: string; + publicationRevision?: string; + }) => Promise; + syncStaticTools: (staticTools: LCAvailableTools) => Promise; + mergeAppTools: (appTools: LCAvailableTools, staticTools: LCAvailableTools) => Promise; + replaceAppServerTools: (params: { + serverName: string; + serverTools: LCAvailableTools; + publicationGeneration?: string; + publicationRevision?: string; + }) => Promise; cacheMCPServerTools: (params: { userId: string; serverName: string; serverTools: LCAvailableTools; serverConfig?: ParsedServerConfig; + publicationGeneration?: string; + publicationRevision?: string; }) => Promise; getMCPServerTools: ( userId: string, @@ -43,42 +78,147 @@ export interface MCPToolCacheService { ) => Promise; } -export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheService { - const { getCachedTools, setCachedTools, getServerConfig } = deps; +interface AppServerBoundary { + serverName: string; + suffix: string; +} - /** - * Request-scoped servers resolve runtime user/request placeholders per - * connection, so their definitions must never enter the persistent tool - * cache. Fails open: an unresolvable config is treated as cacheable, - * preserving pre-gating behavior for servers the registry cannot see. - * The resolver sees only base registry configs — callers holding merged - * Config-overlay configs must pass them. All writers do, so an entry that - * predates gating or an overlay change survives at most one cache TTL. - */ - async function isRequestScoped( - userId: string, +export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheService { + const { + getCachedTools, + updateCachedGlobalTools, + setCachedTools, + setCachedToolsIfCurrent, + getCachedAppServerTools, + setCachedAppServerTools, + getServerConfig, + getAllServerConfigs, + isAppServerConfig, + } = deps; + + async function writeCachedTools( + tools: LCAvailableTools, + options?: { userId?: string; serverName?: string; configGeneration?: string }, + ): Promise { + const success = options ? await setCachedTools(tools, options) : await setCachedTools(tools); + if (success === false) { + throw new Error('Tool cache rejected the write'); + } + } + + async function isAppSharedConfig( serverName: string, - serverConfig?: ParsedServerConfig, + config: ParsedServerConfig | undefined, ): Promise { + if (!config || !canUseAppConnection(config)) { + return false; + } + if (isAppServerConfig) { + return isAppServerConfig(serverName, config); + } + if (!getAllServerConfigs) { + return true; + } try { - const config = serverConfig ?? (await getServerConfig(serverName, userId)); - return config ? requiresEphemeralUserConnection(config) : false; + const appConfigs = await getAllServerConfigs(); + return appConfigs[serverName] != null; } catch (error) { logger.debug( - `[MCP Cache] Could not resolve config for ${serverName} (user: ${userId}), treating as cacheable:`, + `[MCP Cache] Could not verify app ownership for ${serverName}; using user scope:`, error, ); return false; } } + async function resolveCacheConfig( + userId: string | undefined, + serverName: string, + serverConfig?: ParsedServerConfig, + ): Promise { + if (serverConfig) { + return serverConfig; + } + try { + return await getServerConfig(serverName, userId); + } catch (error) { + logger.debug( + `[MCP Cache] Could not resolve config for ${serverName} (user: ${userId}), preserving legacy cache scope:`, + error, + ); + return undefined; + } + } + + function buildAppServerBoundaries(serverNames: readonly string[]): AppServerBoundary[] { + const names = Array.from(new Set(serverNames)); + const boundaryOwners = new Map(); + for (const rawName of names) { + if (normalizeServerName(rawName) !== rawName) { + boundaryOwners.set(`${Constants.mcp_delimiter}${rawName}`, rawName); + } + } + for (const [normalizedName, rawName] of buildServerNameAliases(names)) { + boundaryOwners.set(`${Constants.mcp_delimiter}${normalizedName}`, rawName); + } + + return Array.from(boundaryOwners, ([suffix, rawName]) => ({ + serverName: rawName, + suffix, + })).sort((left, right) => right.suffix.length - left.suffix.length); + } + + async function getAppServerNames(): Promise { + if (!getAllServerConfigs) { + return []; + } + return Object.entries(await getAllServerConfigs()) + .filter(([, config]) => canUseAppConnection(config)) + .map(([name]) => name); + } + + async function getAppServerBoundaries(serverName: string): Promise { + const names = await getAppServerNames(); + if (!names.includes(serverName)) { + names.push(serverName); + } + return buildAppServerBoundaries(names); + } + + function resolveToolServerName( + toolName: string, + boundaries: readonly AppServerBoundary[], + ): string | null { + for (const boundary of boundaries) { + if (toolName.endsWith(boundary.suffix)) { + return boundary.serverName; + } + } + return null; + } + + function getAppServerSlice( + tools: LCAvailableTools, + serverName: string, + boundaries: readonly AppServerBoundary[], + ): LCAvailableTools { + return Object.fromEntries( + Object.entries(tools).filter( + ([name]) => resolveToolServerName(name, boundaries) === serverName, + ), + ); + } + async function updateMCPServerTools(params: { - userId: string; + userId?: string; serverName: string; tools: MCPToolInput[] | null; serverConfig?: ParsedServerConfig; - }): Promise { - const { userId, serverName, tools, serverConfig } = params; + publicationGeneration?: string; + publicationRevision?: string; + }): Promise { + const { userId, serverName, tools, serverConfig, publicationGeneration, publicationRevision } = + params; try { const serverTools: LCAvailableTools = {}; const mcpDelimiter = Constants.mcp_delimiter; @@ -88,16 +228,6 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS return serverTools; } - if (tools.length === 0) { - if (!(await isRequestScoped(userId, serverName, serverConfig))) { - await setCachedTools(serverTools, { userId, serverName }); - logger.debug( - `[MCP Cache] Cleared stale tools for server ${serverName} (user: ${userId})`, - ); - } - return serverTools; - } - /** Cache keys are MODEL-FACING: they become builder tool ids, agent.tools * entries, tool_options keys, and definition names, and must equal the * runtime instance name (`createToolInstance` in MCP.js), which embeds @@ -111,22 +241,65 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS ['function']: { name, description: tool.description ?? '', - parameters: tool.inputSchema ?? ({ type: 'object', properties: {} } as JsonSchemaType), + parameters: tool.inputSchema + ? (normalizeJsonSchema(resolveJsonSchemaRefs(tool.inputSchema)) as JsonSchemaType) + : ({ type: 'object', properties: {} } as JsonSchemaType), }, }; serverTools[name] = entry; } - if (await isRequestScoped(userId, serverName, serverConfig)) { + const resolvedConfig = await resolveCacheConfig(userId, serverName, serverConfig); + const configGeneration = resolvedConfig + ? getMCPAppToolsPublicationGeneration(resolvedConfig) + : undefined; + if (resolvedConfig && requiresEphemeralUserConnection(resolvedConfig)) { logger.debug( `[MCP Cache] Built ${tools.length} tools for request-scoped server ${serverName} (user: ${userId}) without caching`, ); return serverTools; } - await setCachedTools(serverTools, { userId, serverName }); + if (userId && !(await isAppSharedConfig(serverName, resolvedConfig))) { + if (setCachedToolsIfCurrent) { + if (!publicationGeneration || !configGeneration) { + logger.debug( + `[MCP Cache] Skipped unfenced or unaddressed tool publication for ${serverName} (user: ${userId})`, + ); + return null; + } + const current = await setCachedToolsIfCurrent(serverTools, { + userId, + serverName, + configGeneration, + publicationGeneration, + }); + if (!current) { + logger.debug( + `[MCP Cache] Ignored stale tool publication for ${serverName} (user: ${userId})`, + ); + return null; + } + } else { + await writeCachedTools(serverTools, { userId, serverName, configGeneration }); + } + } else { + const appConfigGeneration = + userId == null + ? (publicationGeneration ?? configGeneration) + : (configGeneration ?? publicationGeneration); + const replaced = await replaceAppServerTools({ + serverName, + serverTools, + publicationGeneration: appConfigGeneration, + publicationRevision, + }); + if (!replaced) { + return null; + } + } logger.debug( - `[MCP Cache] Updated ${tools.length} tools for server ${serverName} (user: ${userId})`, + `[MCP Cache] Updated ${tools.length} tools for server ${serverName}${userId ? ` (user: ${userId})` : ' (app-level)'}`, ); return serverTools; } catch (error) { @@ -138,41 +311,160 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS } } - async function mergeAppTools(appTools: LCAvailableTools): Promise { + async function mergeAppTools( + appTools: LCAvailableTools, + staticTools: LCAvailableTools, + ): Promise { try { const count = Object.keys(appTools).length; - if (!count) { - return; - } - const cachedTools = (await getCachedTools()) ?? {}; - const mergedTools: LCAvailableTools = { ...cachedTools, ...appTools }; - await setCachedTools(mergedTools); - logger.debug(`Merged ${count} app-level tools`); + const appConfigs = getAllServerConfigs + ? Object.entries(await getAllServerConfigs()).filter(([, config]) => + canUseAppConnection(config), + ) + : []; + const boundaries = buildAppServerBoundaries(appConfigs.map(([serverName]) => serverName)); + await syncStaticTools(staticTools); + await Promise.all( + appConfigs + .filter(([, config]) => config.toolFunctions != null) + .map(async ([serverName, config]) => { + const serverTools = getAppServerSlice(appTools, serverName, boundaries); + const configGeneration = getMCPAppToolsPublicationGeneration(config); + await setCachedAppServerTools(serverName, configGeneration, serverTools); + }), + ); + logger.debug(`Synchronized ${count} app-level MCP tools`); } catch (error) { logger.error('Failed to merge app-level tools:', error); throw error; } } + async function syncStaticTools(staticTools: LCAvailableTools): Promise { + await updateCachedGlobalTools?.(() => staticTools); + } + + /** + * Replaces one server's configuration-addressed app-level snapshot. Old and new replicas may + * publish concurrently without overwriting each other; readers select the current config key. + */ + async function replaceAppServerTools(params: { + serverName: string; + serverTools: LCAvailableTools; + publicationGeneration?: string; + publicationRevision?: string; + }): Promise { + const { serverName, serverTools, publicationGeneration, publicationRevision } = params; + try { + const boundaries = await getAppServerBoundaries(serverName); + for (const name of Object.keys(serverTools)) { + const owner = resolveToolServerName(name, boundaries); + if (owner && owner !== serverName) { + throw new Error(`Tool ${name} belongs to app server ${owner}, not ${serverName}`); + } + } + let configGeneration = publicationGeneration; + if (!configGeneration) { + const config = await resolveCacheConfig(undefined, serverName); + configGeneration = config ? getMCPAppToolsPublicationGeneration(config) : undefined; + } + if (!configGeneration) { + logger.debug(`[MCP Cache] Skipped unaddressed app-level publication for ${serverName}`); + return false; + } + if (!publicationRevision) { + logger.debug(`[MCP Cache] Skipped unordered app-level publication for ${serverName}`); + return false; + } + const replaced = await setCachedAppServerTools( + serverName, + configGeneration, + serverTools, + publicationRevision, + ); + if (replaced === false) { + logger.debug( + `[MCP Cache] Ignored superseded app-level tools for ${serverName} at revision ${publicationRevision ?? '0'}`, + ); + return false; + } + logger.debug( + `[MCP Cache] Replaced app-level tools for ${serverName} with ${Object.keys(serverTools).length} tool(s)`, + ); + return true; + } catch (error) { + logger.error(`[MCP Cache] Failed to replace app-level tools for ${serverName}:`, error); + throw error; + } + } + async function cacheMCPServerTools(params: { userId: string; serverName: string; serverTools: LCAvailableTools; serverConfig?: ParsedServerConfig; + publicationGeneration?: string; + publicationRevision?: string; }): Promise { - const { userId, serverName, serverTools, serverConfig } = params; + const { + userId, + serverName, + serverTools, + serverConfig, + publicationGeneration, + publicationRevision, + } = params; try { const count = Object.keys(serverTools).length; - if (!count) { - return; - } - if (await isRequestScoped(userId, serverName, serverConfig)) { + const resolvedConfig = await resolveCacheConfig(userId, serverName, serverConfig); + const configGeneration = resolvedConfig + ? getMCPAppToolsPublicationGeneration(resolvedConfig) + : undefined; + if (resolvedConfig && requiresEphemeralUserConnection(resolvedConfig)) { logger.debug( `[MCP Cache] Skipped caching ${count} tools for request-scoped server ${serverName} (user: ${userId})`, ); return; } - await setCachedTools(serverTools, { userId, serverName }); + if (await isAppSharedConfig(serverName, resolvedConfig)) { + const appConfigGeneration = + userId == null + ? (publicationGeneration ?? configGeneration) + : (configGeneration ?? publicationGeneration); + const replaced = await replaceAppServerTools({ + serverName, + serverTools, + publicationGeneration: appConfigGeneration, + publicationRevision, + }); + if (!replaced) { + return; + } + logger.debug(`Refreshed app-level MCP tools for ${serverName}`); + return; + } + if (setCachedToolsIfCurrent) { + if (!publicationGeneration || !configGeneration) { + logger.debug( + `[MCP Cache] Skipped unfenced or unaddressed discovered tools for ${serverName} (user: ${userId})`, + ); + return; + } + const current = await setCachedToolsIfCurrent(serverTools, { + userId, + serverName, + configGeneration, + publicationGeneration, + }); + if (!current) { + logger.debug( + `[MCP Cache] Ignored stale discovered tools for ${serverName} (user: ${userId})`, + ); + return; + } + } else { + await writeCachedTools(serverTools, { userId, serverName, configGeneration }); + } logger.debug(`Cached ${count} MCP server tools for ${serverName} (user: ${userId})`); } catch (error) { logger.error(`Failed to cache MCP server tools for ${serverName} (user: ${userId}):`, error); @@ -223,12 +515,27 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS serverName: string, serverConfig?: ParsedServerConfig, ): Promise { - if (await isRequestScoped(userId, serverName, serverConfig)) { + const resolvedConfig = await resolveCacheConfig(userId, serverName, serverConfig); + if (resolvedConfig && requiresEphemeralUserConnection(resolvedConfig)) { return null; } try { - const cached = (await getCachedTools({ userId, serverName })) ?? null; - if (!cached || Object.keys(cached).length === 0) { + if (await isAppSharedConfig(serverName, resolvedConfig)) { + if (!resolvedConfig) { + return null; + } + const configGeneration = getMCPAppToolsPublicationGeneration(resolvedConfig); + const serverTools = await getCachedAppServerTools(serverName, configGeneration); + if (serverTools == null) { + return null; + } + return normalizeCachedToolKeys(serverTools, serverName); + } + const configGeneration = resolvedConfig + ? getMCPAppToolsPublicationGeneration(resolvedConfig) + : undefined; + const cached = (await getCachedTools({ userId, serverName, configGeneration })) ?? null; + if (!cached) { return null; } return normalizeCachedToolKeys(cached, serverName); @@ -238,5 +545,12 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS } } - return { updateMCPServerTools, mergeAppTools, cacheMCPServerTools, getMCPServerTools }; + return { + updateMCPServerTools, + syncStaticTools, + mergeAppTools, + replaceAppServerTools, + cacheMCPServerTools, + getMCPServerTools, + }; } diff --git a/packages/api/src/mcp/toolsChanged.spec.ts b/packages/api/src/mcp/toolsChanged.spec.ts new file mode 100644 index 0000000000..a3167dceb1 --- /dev/null +++ b/packages/api/src/mcp/toolsChanged.spec.ts @@ -0,0 +1,258 @@ +import type { MCPToolsChangedEvent } from './toolsChanged'; +import type { ParsedServerConfig } from './types'; +import { + setMCPToolsChangedHandler, + setMCPToolsChangedGenerationHandler, + setMCPToolsChangedGenerationRenewalHandler, + setMCPToolsChangedRevisionHandler, + getMCPToolsChangedGeneration, + renewMCPToolsChangedGeneration, + reserveMCPToolsChangedRevision, + hasMCPToolsChangedHandler, + cancelMCPToolsChanged, + notifyMCPToolsChanged, + getMCPAppToolsPublicationGeneration, +} from './toolsChanged'; + +const createEvent = (name = 'one'): MCPToolsChangedEvent => ({ + serverName: 'dynamic', + serverConfig: { type: 'streamable-http', url: 'https://mcp.example.com' }, + tools: [{ name, inputSchema: { type: 'object' } }], +}); + +describe('MCP tools-changed dispatch', () => { + afterEach(() => { + setMCPToolsChangedHandler(null); + setMCPToolsChangedGenerationHandler(null); + setMCPToolsChangedGenerationRenewalHandler(null); + setMCPToolsChangedRevisionHandler(null); + jest.useRealTimers(); + }); + + it('reports whether a handler is registered', () => { + expect(hasMCPToolsChangedHandler()).toBe(false); + setMCPToolsChangedHandler(jest.fn()); + expect(hasMCPToolsChangedHandler()).toBe(true); + setMCPToolsChangedHandler(null); + expect(hasMCPToolsChangedHandler()).toBe(false); + }); + + it('derives stable app publication generations from connection-relevant config', () => { + const first: ParsedServerConfig = { + type: 'sse', + url: 'https://mcp.example.com/sse', + headers: { Authorization: 'Bearer token', Accept: 'text/event-stream' }, + updatedAt: 1, + toolFunctions: {}, + }; + const equivalent: ParsedServerConfig = { + headers: { Accept: 'text/event-stream', Authorization: 'Bearer token' }, + url: 'https://mcp.example.com/sse', + type: 'sse', + updatedAt: 2, + inspectionFailed: true, + }; + const changed: ParsedServerConfig = { + ...equivalent, + url: 'https://mcp.example.com/v2/sse', + }; + + expect(getMCPAppToolsPublicationGeneration(first)).toBe( + getMCPAppToolsPublicationGeneration(equivalent), + ); + expect(getMCPAppToolsPublicationGeneration(first)).not.toBe( + getMCPAppToolsPublicationGeneration(changed), + ); + }); + + it('includes the resolved runtime environment in app publication generations', () => { + const variable = 'LIBRECHAT_MCP_CATALOG_ORIGIN_TEST'; + const original = process.env[variable]; + const config: ParsedServerConfig = { + type: 'streamable-http', + url: `\${${variable}}/mcp`, + }; + + try { + process.env[variable] = 'https://old.example.com'; + const oldGeneration = getMCPAppToolsPublicationGeneration(config); + process.env[variable] = 'https://new.example.com'; + const newGeneration = getMCPAppToolsPublicationGeneration(config); + + expect(newGeneration).not.toBe(oldGeneration); + } finally { + if (original === undefined) { + delete process.env[variable]; + } else { + process.env[variable] = original; + } + } + }); + + it('captures a connection-bound publication generation from the app layer', async () => { + const generationHandler = jest.fn().mockResolvedValue('generation-a'); + setMCPToolsChangedGenerationHandler(generationHandler); + + await expect( + getMCPToolsChangedGeneration({ userId: 'user-1', serverName: 'dynamic' }), + ).resolves.toBe('generation-a'); + expect(generationHandler).toHaveBeenCalledWith({ + userId: 'user-1', + serverName: 'dynamic', + }); + }); + + it('renews a current connection-bound publication generation through the app layer', async () => { + const renewalHandler = jest.fn().mockResolvedValue(true); + setMCPToolsChangedGenerationRenewalHandler(renewalHandler); + const scope = { + userId: 'user-1', + serverName: 'dynamic', + publicationGeneration: 'generation-a', + }; + + await expect(renewMCPToolsChangedGeneration(scope)).resolves.toBe(true); + expect(renewalHandler).toHaveBeenCalledWith(scope); + }); + + it('reserves app revisions by runtime config and skips user scopes', async () => { + const revisionHandler = jest.fn().mockResolvedValue('7'); + const serverConfig: ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/mcp', + }; + setMCPToolsChangedRevisionHandler(revisionHandler); + + await expect( + reserveMCPToolsChangedRevision({ serverName: 'dynamic', serverConfig }), + ).resolves.toBe('7'); + await expect( + reserveMCPToolsChangedRevision({ serverName: 'dynamic', serverConfig, userId: 'user-1' }), + ).resolves.toBeUndefined(); + expect(revisionHandler).toHaveBeenCalledTimes(1); + expect(revisionHandler).toHaveBeenCalledWith({ + serverName: 'dynamic', + configGeneration: getMCPAppToolsPublicationGeneration(serverConfig), + }); + }); + + it('passes a complete server snapshot and user scope to the handler', async () => { + const handler = jest.fn(); + const event = { ...createEvent(), userId: 'user-1' }; + setMCPToolsChangedHandler(handler); + + await notifyMCPToolsChanged(event); + + expect(handler).toHaveBeenCalledWith(event); + }); + + it('awaits an async handler before returning', async () => { + let finished = false; + setMCPToolsChangedHandler(async () => { + await Promise.resolve(); + finished = true; + }); + + await notifyMCPToolsChanged(createEvent()); + + expect(finished).toBe(true); + }); + + it('coalesces an in-flight burst and publishes the newest snapshot last', async () => { + let releaseFirst: (() => void) | undefined; + const firstBlocked = new Promise((resolve) => { + releaseFirst = resolve; + }); + const published: string[] = []; + const handler = jest.fn(async (event: MCPToolsChangedEvent) => { + published.push(event.tools[0].name); + if (event.tools[0].name === 'one') { + await firstBlocked; + } + }); + setMCPToolsChangedHandler(handler); + + const first = notifyMCPToolsChanged(createEvent('one')); + await Promise.resolve(); + const second = notifyMCPToolsChanged(createEvent('two')); + const third = notifyMCPToolsChanged(createEvent('three')); + releaseFirst?.(); + await Promise.all([first, second, third]); + + expect(published).toEqual(['one', 'three']); + }); + + it('retries a failed cache publication without rejecting the notification handler', async () => { + jest.useFakeTimers(); + const handler = jest + .fn, [MCPToolsChangedEvent]>() + .mockRejectedValueOnce(new Error('Redis down')) + .mockResolvedValue(undefined); + setMCPToolsChangedHandler(handler); + + await expect(notifyMCPToolsChanged(createEvent())).resolves.toBeUndefined(); + expect(handler).toHaveBeenCalledTimes(1); + + await jest.advanceTimersByTimeAsync(250); + + expect(handler).toHaveBeenCalledTimes(2); + }); + + it('stops dispatching when the handler is unregistered during an in-flight failure', async () => { + let rejectPublication: ((error: Error) => void) | undefined; + const publication = new Promise((_, reject) => { + rejectPublication = reject; + }); + const handler = jest.fn(() => publication); + setMCPToolsChangedHandler(handler); + + const notification = notifyMCPToolsChanged(createEvent()); + await Promise.resolve(); + setMCPToolsChangedHandler(null); + rejectPublication?.(new Error('publisher shutting down')); + + await expect(notification).resolves.toBeUndefined(); + await Promise.resolve(); + expect(handler).toHaveBeenCalledTimes(1); + }); + + it('cancels a queued retry before cache invalidation', async () => { + jest.useFakeTimers(); + const handler = jest.fn().mockRejectedValue(new Error('Redis down')); + const event = { ...createEvent(), userId: 'user-1' }; + setMCPToolsChangedHandler(handler); + + await notifyMCPToolsChanged(event); + await cancelMCPToolsChanged(event); + await jest.advanceTimersByTimeAsync(30_000); + + expect(handler).toHaveBeenCalledTimes(1); + }); + + it('drains an in-flight publication before disconnect returns', async () => { + let releasePublication: (() => void) | undefined; + const publication = new Promise((resolve) => { + releasePublication = resolve; + }); + const event = { ...createEvent(), userId: 'user-1' }; + setMCPToolsChangedHandler(() => publication); + + const notification = notifyMCPToolsChanged(event); + await Promise.resolve(); + let drained = false; + const cancellation = cancelMCPToolsChanged(event).then(() => { + drained = true; + }); + await Promise.resolve(); + expect(drained).toBe(false); + + releasePublication?.(); + await Promise.all([notification, cancellation]); + + expect(drained).toBe(true); + }); + + it('does nothing when no handler is registered', async () => { + await expect(notifyMCPToolsChanged(createEvent())).resolves.toBeUndefined(); + }); +}); diff --git a/packages/api/src/mcp/toolsChanged.ts b/packages/api/src/mcp/toolsChanged.ts new file mode 100644 index 0000000000..be250889eb --- /dev/null +++ b/packages/api/src/mcp/toolsChanged.ts @@ -0,0 +1,265 @@ +import { createHash } from 'crypto'; +import { logger } from '@librechat/data-schemas'; +import { MCPOptionsSchema } from 'librechat-data-provider'; +import type { Tool } from '@modelcontextprotocol/sdk/types.js'; +import type { MCPOptions, ParsedServerConfig } from './types'; +import { processMCPEnv } from '../utils/env'; + +const RETRY_BASE_DELAY_MS = 250; +const RETRY_MAX_DELAY_MS = 30_000; + +type StableConfigValue = + | string + | number + | boolean + | null + | StableConfigValue[] + | { [key: string]: StableConfigValue | undefined }; + +function sortConfigValue(value: StableConfigValue): StableConfigValue { + if (Array.isArray(value)) { + return value.map(sortConfigValue); + } + if (value == null || typeof value !== 'object') { + return value; + } + const sorted: { [key: string]: StableConfigValue | undefined } = {}; + for (const key of Object.keys(value).sort()) { + sorted[key] = sortConfigValue(value[key] as StableConfigValue); + } + return sorted; +} + +/** Returns a stable token for the connection-relevant portion of an MCP config. */ +export function getMCPAppToolsPublicationGeneration(config: ParsedServerConfig): string { + /** App replicas can resolve the same stored config through different process environments during + * a rolling deployment. Address the catalog by the effective runtime config so an old replica's + * live connection cannot publish into the new replica's slice. DB-sourced configs deliberately + * remain literal because processMCPEnv derives that rule from dbId. */ + const runtimeConfig = processMCPEnv({ options: config }); + const parsedConfig = MCPOptionsSchema.parse(runtimeConfig) as StableConfigValue; + return createHash('sha256') + .update(JSON.stringify(sortConfigValue(parsedConfig))) + .digest('hex'); +} + +/** A complete tool-list snapshot and the cache scope it belongs to. */ +export interface MCPToolsChangedEvent { + serverName: string; + tools: Tool[]; + serverConfig: MCPOptions; + userId?: string; + /** Connection-bound token used to fence stale cross-replica cache publications. */ + publicationGeneration?: string; + /** Monotonic ticket assigned before an app-level tools/list request begins. */ + publicationRevision?: string; +} + +export type MCPToolsChangedHandler = (event: MCPToolsChangedEvent) => Promise | void; + +interface PendingToolsChange { + latest: MCPToolsChangedEvent; + generation: number; + handledGeneration: number; + failures: number; + refreshPromise: Promise | null; + retryTimer: ReturnType | null; +} + +let handler: MCPToolsChangedHandler | null = null; +const pendingChanges = new Map(); + +type MCPToolsChangedScope = Pick; + +export type MCPToolsChangedGenerationHandler = ( + scope: MCPToolsChangedScope, +) => Promise | string | undefined; + +let generationHandler: MCPToolsChangedGenerationHandler | null = null; + +export type MCPToolsChangedGenerationRenewalHandler = ( + scope: MCPToolsChangedScope & { publicationGeneration: string }, +) => Promise | boolean; + +let generationRenewalHandler: MCPToolsChangedGenerationRenewalHandler | null = null; + +export type MCPToolsChangedRevisionHandler = (scope: { + serverName: string; + configGeneration: string; +}) => Promise | string; + +let revisionHandler: MCPToolsChangedRevisionHandler | null = null; + +function getChangeKey(event: MCPToolsChangedScope): string { + return JSON.stringify([event.userId ?? null, event.serverName]); +} + +function clearRetryTimer(change: PendingToolsChange): void { + if (change.retryTimer) { + clearTimeout(change.retryTimer); + change.retryTimer = null; + } +} + +function scheduleRetry(key: string, change: PendingToolsChange): void { + if (change.retryTimer || !handler || pendingChanges.get(key) !== change) { + return; + } + + const delay = Math.min( + RETRY_BASE_DELAY_MS * Math.pow(2, Math.max(0, change.failures - 1)), + RETRY_MAX_DELAY_MS, + ); + change.retryTimer = setTimeout(() => { + change.retryTimer = null; + startDispatch(key, change); + }, delay); + change.retryTimer.unref?.(); +} + +async function dispatchPendingChange(key: string, change: PendingToolsChange): Promise { + while (handler && change.handledGeneration < change.generation) { + const targetGeneration = change.generation; + const event = change.latest; + try { + await handler(event); + change.handledGeneration = targetGeneration; + change.failures = 0; + } catch (error) { + change.failures++; + logger.error( + `[MCP][${event.serverName}] Failed to publish tools after list_changed; retrying:`, + error, + ); + scheduleRetry(key, change); + return; + } + } +} + +function startDispatch(key: string, change: PendingToolsChange): Promise { + if (change.refreshPromise) { + return change.refreshPromise; + } + + change.refreshPromise = dispatchPendingChange(key, change).finally(() => { + change.refreshPromise = null; + if (!handler || pendingChanges.get(key) !== change) { + return; + } + if (change.handledGeneration >= change.generation) { + pendingChanges.delete(key); + } else if (!change.retryTimer) { + return startDispatch(key, change); + } + }); + return change.refreshPromise; +} + +/** Registers the app-layer publisher for refreshed MCP tool snapshots. */ +export function setMCPToolsChangedHandler(fn: MCPToolsChangedHandler | null): void { + handler = fn; + if (!fn) { + for (const change of pendingChanges.values()) { + clearRetryTimer(change); + } + pendingChanges.clear(); + } +} + +export function hasMCPToolsChangedHandler(): boolean { + return handler != null; +} + +/** Registers the app-layer provider for connection-bound publication generations. */ +export function setMCPToolsChangedGenerationHandler( + fn: MCPToolsChangedGenerationHandler | null, +): void { + generationHandler = fn; +} + +/** Registers the app-layer lease renewer for active durable user connections. */ +export function setMCPToolsChangedGenerationRenewalHandler( + fn: MCPToolsChangedGenerationRenewalHandler | null, +): void { + generationRenewalHandler = fn; +} + +/** Registers the shared app-catalog revision allocator. */ +export function setMCPToolsChangedRevisionHandler(fn: MCPToolsChangedRevisionHandler | null): void { + revisionHandler = fn; +} + +/** Captures the current cache generation before a durable user connection is created. */ +export async function getMCPToolsChangedGeneration( + scope: MCPToolsChangedScope, +): Promise { + return generationHandler?.(scope); +} + +/** Renews a connection's publication lease without allowing a stale generation to revive. */ +export async function renewMCPToolsChangedGeneration( + scope: MCPToolsChangedScope & { publicationGeneration: string }, +): Promise { + return generationRenewalHandler?.(scope); +} + +/** Reserves ordering before an app-level tools/list request starts. */ +export async function reserveMCPToolsChangedRevision(scope: { + serverName: string; + serverConfig: ParsedServerConfig; + userId?: string; +}): Promise { + if (scope.userId || !revisionHandler) { + return undefined; + } + return revisionHandler({ + serverName: scope.serverName, + configGeneration: getMCPAppToolsPublicationGeneration(scope.serverConfig), + }); +} + +/** + * Publishes the latest snapshot for a server. Concurrent notifications are single-flighted and + * cache-write failures retain the latest snapshot for bounded-backoff retries. + */ +export async function notifyMCPToolsChanged(event: MCPToolsChangedEvent): Promise { + if (!handler) { + logger.debug( + `[MCP][${event.serverName}] Tool list changed but no handler is registered; tools stay as they were`, + ); + return; + } + + const key = getChangeKey(event); + let change = pendingChanges.get(key); + if (!change) { + change = { + latest: event, + generation: 0, + handledGeneration: 0, + failures: 0, + refreshPromise: null, + retryTimer: null, + }; + pendingChanges.set(key, change); + } + + change.latest = event; + change.generation++; + clearRetryTimer(change); + await startDispatch(key, change); +} + +/** Cancels queued retries and drains an in-flight publication before cache invalidation. */ +export async function cancelMCPToolsChanged(scope: MCPToolsChangedScope): Promise { + const key = getChangeKey(scope); + const change = pendingChanges.get(key); + if (!change) { + return; + } + pendingChanges.delete(key); + clearRetryTimer(change); + change.generation = change.handledGeneration; + await change.refreshPromise; +} diff --git a/packages/api/src/mcp/utils.ts b/packages/api/src/mcp/utils.ts index d56194a63a..07c296ffbd 100644 --- a/packages/api/src/mcp/utils.ts +++ b/packages/api/src/mcp/utils.ts @@ -1,5 +1,6 @@ import { Constants, + MCPOptionsSchema, normalizeServerName, normalizeMCPToolKey, buildServerNameAliases, @@ -10,6 +11,18 @@ import type { RequestBody } from '~/types'; export const mcpToolPattern: RegExp = new RegExp(`^.+${Constants.mcp_delimiter}.+$`); +function isMCPServerConfig(config: unknown): config is ParsedServerConfig { + return MCPOptionsSchema.safeParse(config).success; +} + +/** Validates an effective MCP config without stripping its server-managed metadata. */ +export function validateMCPServerConfig(config: unknown): ParsedServerConfig { + if (!isMCPServerConfig(config)) { + throw new Error('Invalid effective MCP server configuration'); + } + return config; +} + /** * Prefix of the lazily-expanded MCP placeholder `mcp_all`, * pushed into an agent's `tools` for overlay/user-connection servers whose @@ -184,7 +197,10 @@ type PlaceholderValue = | readonly PlaceholderValue[] | { readonly [key: string]: PlaceholderValue }; -type UserScopedConnectionConfig = Pick & { +type UserScopedConnectionConfig = Pick< + ParsedServerConfig, + 'requiresOAuth' | 'source' | 'dbId' | 'startup' +> & { args?: string[]; /** Loosened from the parsed shapes so raw (pre-inspection) configs qualify; * scoping predicates only check key presence */ @@ -372,6 +388,13 @@ export function requiresUserScopedConnection(config: UserScopedConnectionConfig) ); } +/** Whether a server can share one operator-owned connection across all users. */ +export function canUseAppConnection(config: UserScopedConnectionConfig): boolean { + return ( + config.startup !== false && !isUserSourced(config) && !requiresUserScopedConnection(config) + ); +} + /** * Returns the names of `customUserVars` declared on the server config for which * the user has not supplied a non-blank value (unset, empty, or whitespace-only