diff --git a/packages/api/src/mcp/UserConnectionManager.ts b/packages/api/src/mcp/UserConnectionManager.ts index b52d66795c..67505ee9a7 100644 --- a/packages/api/src/mcp/UserConnectionManager.ts +++ b/packages/api/src/mcp/UserConnectionManager.ts @@ -229,9 +229,10 @@ export abstract class UserConnectionManager { * against the config and connection state the teardown left behind, instead of failing. */ public async getUserConnection(opts: t.UserMCPConnectionOptions): Promise { + let attemptOptions = opts; for (let attempt = 0; ; attempt++) { try { - return await this.establishUserConnection(opts); + return await this.establishUserConnection(attemptOptions); } catch (error) { if ( !(error instanceof ConnectionCreationCancelledError) || @@ -242,10 +243,38 @@ export abstract class UserConnectionManager { logger.info( `[MCP][User: ${opts.user?.id}][${opts.serverName}] Connection creation raced a teardown; re-establishing`, ); + attemptOptions = await this.refreshFencedServerConfig(attemptOptions); } } } + /** + * Callers hand in a config they resolved before the teardown, so a re-established attempt + * re-reads it: a registry-backed server picks up the committed update, and a deleted user + * server resolves to nothing and fails the attempt instead of reconnecting to the removed + * endpoint. A caller-supplied config the registry cannot resolve on its own — a config-source + * server not cached yet — is kept, since no user mutation removes it. + */ + private async refreshFencedServerConfig( + opts: t.UserMCPConnectionOptions, + ): Promise { + const userId = opts.user?.id; + if (!userId) { + return opts; + } + const serverConfig = await MCPServersRegistry.getInstance().getServerConfig( + opts.serverName, + userId, + ); + if (serverConfig) { + return { ...opts, serverConfig }; + } + if (opts.serverConfig && !isUserSourced(opts.serverConfig)) { + return opts; + } + return { ...opts, serverConfig: undefined }; + } + private async establishUserConnection(opts: t.UserMCPConnectionOptions): Promise { const { serverName, forceNew, user } = opts; const userId = user?.id; diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index 3b8c758c4c..792ed46668 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -3363,7 +3363,7 @@ describe('MCPManager', () => { expect(manager.getUserConnections(userId)?.get(serverName)).toBe(reestablishedConnection); }); - it('re-reads the server config when a teardown fences an in-flight creation', async () => { + it('re-reads the committed config when a teardown fences a caller-resolved creation', async () => { const fencedConnection = newUserConnection(); const reestablishedConnection = newUserConnection(); const staleConfig: t.ParsedServerConfig = { @@ -3383,7 +3383,12 @@ describe('MCPManager', () => { .mockResolvedValue(reestablishedConnection); const manager = await MCPManager.createInstance(newMCPServersConfig()); - const creation = manager.getUserConnection({ serverName, user: mockUser }); + /** `MCPManager.getConnection` resolves the config once and hands it to every attempt. */ + const creation = manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig: staleConfig, + }); while ((MCPConnectionFactory.create as jest.Mock).mock.calls.length === 0) { await new Promise((resolve) => setImmediate(resolve)); } @@ -3405,6 +3410,81 @@ describe('MCPManager', () => { ); }); + it('fails a fenced creation whose user server was deleted during the teardown', async () => { + const fencedConnection = newUserConnection(); + const deletedConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/deleted', + source: 'user', + dbId: 'server-1', + }; + 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(deletedConfig); + (MCPConnectionFactory.create as jest.Mock).mockReturnValue(factoryResult); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const creation = manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig: deletedConfig, + }); + while ((MCPConnectionFactory.create as jest.Mock).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(undefined); + await manager.disconnectUserConnection(userId, serverName); + resolveConnection?.(fencedConnection); + + await expect(creation).rejects.toThrow(`Configuration for server "${serverName}" not found`); + expect(fencedConnection.dispose).toHaveBeenCalledTimes(1); + expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(1); + expect(manager.getUserConnections(userId)?.has(serverName) ?? false).toBe(false); + }); + + it('keeps a config-source config the registry cannot resolve when re-establishing', async () => { + const fencedConnection = newUserConnection(); + const reestablishedConnection = newUserConnection(); + const configSourced: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/config-sourced', + source: 'config', + }; + 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(undefined); + (MCPConnectionFactory.create as jest.Mock) + .mockReturnValueOnce(factoryResult) + .mockResolvedValue(reestablishedConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const creation = manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig: configSourced, + }); + while ((MCPConnectionFactory.create as jest.Mock).mock.calls.length === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } + + await manager.disconnectUserConnection(userId, serverName); + resolveConnection?.(fencedConnection); + + await expect(creation).resolves.toBe(reestablishedConnection); + expect(fencedConnection.dispose).toHaveBeenCalledTimes(1); + expect(MCPConnectionFactory.create).toHaveBeenLastCalledWith( + expect.objectContaining({ serverConfig: expect.objectContaining(configSourced) }), + expect.anything(), + ); + }); + it('fails a creation that a teardown keeps cancelling on every attempt', async () => { mockAppConnections({ has: jest.fn().mockResolvedValue(false) }); (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue({