🚉 feat: MCP Registry Individual Server Init (2) (#9940)

* initialize servers sequentially

* adjust for exported properties that are not nullable anymore

* use underscore separator

* mock with set

* customize init timeout via env var

* refactor for readability, use loaded conns for tool functions

* address PR comments

* clean up fire-and-forget

* fix tests

packages/api/src/mcp/__tests__/MCPServersRegistry.test.ts

@@ -113,6 +113,7 @@ describe('MCPServersRegistry - Initialize Function', () => {
       get: jest.fn(),
       getLoaded: jest.fn(),
       disconnectAll: jest.fn(),
+      disconnect: jest.fn().mockResolvedValue(undefined),
     } as unknown as jest.Mocked<ConnectionsRepository>;
 
     mockConnectionsRepo.get.mockImplementation((serverName: string) => {
@@ -160,6 +161,7 @@ describe('MCPServersRegistry - Initialize Function', () => {
   });
 
   afterEach(() => {
+    delete process.env.MCP_INIT_TIMEOUT_MS;
     jest.clearAllMocks();
   });
 
@@ -179,15 +181,14 @@ describe('MCPServersRegistry - Initialize Function', () => {
       const registry = new MCPServersRegistry(rawConfigs);
 
       // Verify initial state
-      expect(registry.oauthServers).toBeNull();
-      expect(registry.serverInstructions).toBeNull();
-      expect(registry.toolFunctions).toBeNull();
-      expect(registry.appServerConfigs).toBeNull();
+      expect(registry.oauthServers.size).toBe(0);
+      expect(registry.serverInstructions).toEqual({});
+      expect(registry.toolFunctions).toEqual({});
+      expect(registry.appServerConfigs).toEqual({});
 
       await registry.initialize();
 
       // Test oauthServers Set
-      expect(registry.oauthServers).toBeInstanceOf(Set);
       expect(registry.oauthServers).toEqual(
         new Set(['oauth_server', 'oauth_predefined', 'oauth_startup_enabled']),
       );
@@ -228,18 +229,49 @@ describe('MCPServersRegistry - Initialize Function', () => {
       expect(registry.toolFunctions).toEqual(expectedToolFunctions);
     });
 
-    it('should handle errors gracefully and continue initialization', async () => {
+    it('should handle errors gracefully and continue initialization of other servers', async () => {
       const registry = new MCPServersRegistry(rawConfigs);
 
-      // Make one server throw an error
-      mockDetectOAuthRequirement.mockRejectedValueOnce(new Error('OAuth detection failed'));
+      // Make one specific server throw an error during OAuth detection
+      mockDetectOAuthRequirement.mockImplementation((url: string) => {
+        if (url === 'https://api.github.com/mcp') {
+          return Promise.reject(new Error('OAuth detection failed'));
+        }
+        // Return normal responses for other servers
+        const oauthResults: Record<string, OAuthDetectionResult> = {
+          'https://api.disabled.com/mcp': {
+            requiresOAuth: false,
+            method: 'no-metadata-found',
+            metadata: null,
+          },
+          'https://api.public.com/mcp': {
+            requiresOAuth: false,
+            method: 'no-metadata-found',
+            metadata: null,
+          },
+        };
+        return Promise.resolve(
+          oauthResults[url] ?? {
+            requiresOAuth: false,
+            method: 'no-metadata-found',
+            metadata: null,
+          },
+        );
+      });
 
       await registry.initialize();
 
-      // Should still initialize successfully
+      // Should still initialize successfully for other servers
       expect(registry.oauthServers).toBeInstanceOf(Set);
       expect(registry.toolFunctions).toBeDefined();
 
+      // The failed server should not be in oauthServers (since it failed OAuth detection)
+      expect(registry.oauthServers.has('oauth_server')).toBe(false);
+
+      // But other servers should still be processed successfully
+      expect(registry.appServerConfigs).toHaveProperty('stdio_server');
+      expect(registry.appServerConfigs).toHaveProperty('non_oauth_server');
+
       // Error should be logged as a warning at the higher level
       expect(mockLogger.warn).toHaveBeenCalledWith(
         expect.stringContaining('[MCP][oauth_server] Failed to initialize server:'),
@@ -247,12 +279,15 @@ describe('MCPServersRegistry - Initialize Function', () => {
       );
     });
 
-    it('should disconnect all connections after initialization', async () => {
+    it('should disconnect individual connections after each server initialization', async () => {
       const registry = new MCPServersRegistry(rawConfigs);
 
       await registry.initialize();
 
-      expect(mockConnectionsRepo.disconnectAll).toHaveBeenCalledTimes(1);
+      // Verify disconnect was called for each server during initialization
+      // All servers attempt to connect during initialization for metadata gathering
+      const serverNames = Object.keys(rawConfigs);
+      expect(mockConnectionsRepo.disconnect).toHaveBeenCalledTimes(serverNames.length);
     });
 
     it('should log configuration updates for each startup-enabled server', async () => {
@@ -357,5 +392,204 @@ describe('MCPServersRegistry - Initialize Function', () => {
       // Verify getInstructions was called for both "true" cases
       expect(mockClient.getInstructions).toHaveBeenCalledTimes(2);
     });
+
+    it('should use Promise.allSettled for individual server initialization', async () => {
+      const registry = new MCPServersRegistry(rawConfigs);
+
+      // Spy on Promise.allSettled to verify it's being used
+      const allSettledSpy = jest.spyOn(Promise, 'allSettled');
+
+      await registry.initialize();
+
+      // Verify Promise.allSettled was called with an array of server initialization promises
+      expect(allSettledSpy).toHaveBeenCalledWith(expect.arrayContaining([expect.any(Promise)]));
+
+      // Verify it was called with the correct number of server promises
+      const serverNames = Object.keys(rawConfigs);
+      expect(allSettledSpy).toHaveBeenCalledWith(
+        expect.arrayContaining(new Array(serverNames.length).fill(expect.any(Promise))),
+      );
+
+      allSettledSpy.mockRestore();
+    });
+
+    it('should isolate server failures and not affect other servers', async () => {
+      const registry = new MCPServersRegistry(rawConfigs);
+
+      // Make multiple servers fail in different ways
+      mockConnectionsRepo.get.mockImplementation((serverName: string) => {
+        if (serverName === 'stdio_server') {
+          // First server fails
+          throw new Error('Connection failed for stdio_server');
+        }
+        if (serverName === 'websocket_server') {
+          // Second server fails
+          throw new Error('Connection failed for websocket_server');
+        }
+        // Other servers succeed
+        const connection = mockConnections.get(serverName);
+        if (!connection) {
+          throw new Error(`Connection not found for server: ${serverName}`);
+        }
+        return Promise.resolve(connection);
+      });
+
+      await registry.initialize();
+
+      // Despite failures, initialization should complete
+      expect(registry.oauthServers).toBeInstanceOf(Set);
+      expect(registry.toolFunctions).toBeDefined();
+
+      // Successful servers should still be processed
+      expect(registry.appServerConfigs).toHaveProperty('non_oauth_server');
+
+      // Failed servers should not crash the whole initialization
+      expect(mockLogger.warn).toHaveBeenCalledWith(
+        expect.stringContaining('[MCP][stdio_server] Failed to fetch server capabilities:'),
+        expect.any(Error),
+      );
+      expect(mockLogger.warn).toHaveBeenCalledWith(
+        expect.stringContaining('[MCP][websocket_server] Failed to fetch server capabilities:'),
+        expect.any(Error),
+      );
+    });
+
+    it('should properly clean up connections even when some servers fail', async () => {
+      const registry = new MCPServersRegistry(rawConfigs);
+
+      // Track disconnect failures but suppress unhandled rejections
+      const disconnectErrors: Error[] = [];
+      mockConnectionsRepo.disconnect.mockImplementation((serverName: string) => {
+        if (serverName === 'stdio_server') {
+          const error = new Error('Disconnect failed');
+          disconnectErrors.push(error);
+          return Promise.reject(error).catch(() => {}); // Suppress unhandled rejection
+        }
+        return Promise.resolve();
+      });
+
+      await registry.initialize();
+
+      // Should still attempt to disconnect all servers during initialization
+      const serverNames = Object.keys(rawConfigs);
+      expect(mockConnectionsRepo.disconnect).toHaveBeenCalledTimes(serverNames.length);
+      expect(disconnectErrors).toHaveLength(1);
+    });
+
+    it('should timeout individual server initialization after configured timeout', async () => {
+      const timeout = 2000;
+      // Create registry with a short timeout for testing
+      process.env.MCP_INIT_TIMEOUT_MS = `${timeout}`;
+
+      const registry = new MCPServersRegistry(rawConfigs);
+
+      // Make one server hang indefinitely during OAuth detection
+      mockDetectOAuthRequirement.mockImplementation((url: string) => {
+        if (url === 'https://api.github.com/mcp') {
+          // Slow init
+          return new Promise((res) => setTimeout(res, timeout * 2));
+        }
+        // Return normal responses for other servers
+        return Promise.resolve({
+          requiresOAuth: false,
+          method: 'no-metadata-found',
+          metadata: null,
+        });
+      });
+
+      const start = Date.now();
+      await registry.initialize();
+      const duration = Date.now() - start;
+
+      // Should complete within reasonable time despite one server hanging
+      // Allow some buffer for test execution overhead
+      expect(duration).toBeLessThan(timeout * 1.5);
+
+      // The timeout should prevent the hanging server from blocking initialization
+      // Other servers should still be processed successfully
+      expect(registry.appServerConfigs).toHaveProperty('stdio_server');
+      expect(registry.appServerConfigs).toHaveProperty('non_oauth_server');
+    }, 10_000); // 10 second Jest timeout
+
+    it('should skip tool function fetching if connection was not established', async () => {
+      const testConfig: t.MCPServers = {
+        server_with_connection: {
+          type: 'stdio',
+          args: [],
+          command: 'test-command',
+        },
+        server_without_connection: {
+          type: 'stdio',
+          args: [],
+          command: 'failing-command',
+        },
+      };
+
+      const registry = new MCPServersRegistry(testConfig);
+
+      const mockClient = {
+        listTools: jest.fn().mockResolvedValue({
+          tools: [
+            {
+              name: 'test_tool',
+              description: 'Test tool',
+              inputSchema: { type: 'object', properties: {} },
+            },
+          ],
+        }),
+        getInstructions: jest.fn().mockReturnValue(undefined),
+        getServerCapabilities: jest.fn().mockReturnValue({ tools: {} }),
+      };
+      const mockConnection = {
+        client: mockClient,
+      } as unknown as jest.Mocked<MCPConnection>;
+
+      mockConnectionsRepo.get.mockImplementation((serverName: string) => {
+        if (serverName === 'server_with_connection') {
+          return Promise.resolve(mockConnection);
+        }
+        throw new Error('Connection failed');
+      });
+
+      // Mock getLoaded to return connections map - the real implementation returns all loaded connections at once
+      mockConnectionsRepo.getLoaded.mockResolvedValue(
+        new Map([['server_with_connection', mockConnection]]),
+      );
+
+      mockDetectOAuthRequirement.mockResolvedValue({
+        requiresOAuth: false,
+        method: 'no-metadata-found',
+        metadata: null,
+      });
+
+      await registry.initialize();
+
+      expect(registry.toolFunctions).toHaveProperty('test_tool_mcp_server_with_connection');
+      expect(Object.keys(registry.toolFunctions)).toHaveLength(1);
+    });
+
+    it('should handle getLoaded returning empty map gracefully', async () => {
+      const testConfig: t.MCPServers = {
+        test_server: {
+          type: 'stdio',
+          args: [],
+          command: 'test-command',
+        },
+      };
+
+      const registry = new MCPServersRegistry(testConfig);
+
+      mockConnectionsRepo.get.mockRejectedValue(new Error('All connections failed'));
+      mockConnectionsRepo.getLoaded.mockResolvedValue(new Map());
+      mockDetectOAuthRequirement.mockResolvedValue({
+        requiresOAuth: false,
+        method: 'no-metadata-found',
+        metadata: null,
+      });
+
+      await registry.initialize();
+
+      expect(registry.toolFunctions).toEqual({});
+    });
   });
 });