From 11dc33eab793a6259b422168d180d2ea37d5a8f5 Mon Sep 17 00:00:00 2001
From: Jack Wotherspoon <jackwoth@google.com>
Date: Tue, 24 Mar 2026 13:53:21 -0700
Subject: [PATCH] fix(core): prevent premature MCP discovery completion
 (#23637)

---
 .../core/src/tools/mcp-client-manager.test.ts | 45 +++++++++++++++++++
 packages/core/src/tools/mcp-client-manager.ts | 11 +++--
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/packages/core/src/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts
index 84d3e138ce..a96f3f7d29 100644
--- a/packages/core/src/tools/mcp-client-manager.test.ts
+++ b/packages/core/src/tools/mcp-client-manager.test.ts
@@ -147,6 +147,51 @@ describe('McpClientManager', () => {
     expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
   });
 
+  it('should NOT set COMPLETED prematurely when startConfiguredMcpServers finishes before parallel extensions', async () => {
+    mockConfig.getMcpServers.mockReturnValue({});
+    const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
+
+    let resolveExtension: (value: void) => void;
+    const extensionPromise = new Promise<void>((resolve) => {
+      resolveExtension = resolve;
+    });
+
+    mockedMcpClient.connect.mockImplementation(async () => {
+      await extensionPromise;
+    });
+
+    const extensionStartPromise = manager.startExtension({
+      name: 'test-extension',
+      mcpServers: {
+        'extension-server': { command: 'node' },
+      },
+      isActive: true,
+      version: '1.0.0',
+      path: '/some-path',
+      contextFiles: [],
+      id: '123',
+    });
+
+    // Wait for the state to become IN_PROGRESS (since maybeDiscoverMcpServer is async)
+    await vi.waitFor(() => {
+      if (manager.getDiscoveryState() !== MCPDiscoveryState.IN_PROGRESS) {
+        throw new Error('Discovery state is not IN_PROGRESS');
+      }
+    });
+
+    expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.IN_PROGRESS);
+
+    await manager.startConfiguredMcpServers();
+
+    // discoveryState should still be IN_PROGRESS because the extension is still starting
+    expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.IN_PROGRESS);
+
+    resolveExtension!(undefined);
+    await extensionStartPromise;
+
+    expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.COMPLETED);
+  });
+
   it('should mark discovery completed when all configured servers are blocked', async () => {
     mockConfig.getMcpServers.mockReturnValue({
       'test-server': { command: 'node' },
diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts
index 666b6d5321..3e7ef75d4c 100644
--- a/packages/core/src/tools/mcp-client-manager.ts
+++ b/packages/core/src/tools/mcp-client-manager.ts
@@ -554,8 +554,10 @@ export class McpClientManager {
     );
 
     if (Object.keys(servers).length === 0) {
-      this.discoveryState = MCPDiscoveryState.COMPLETED;
-      this.eventEmitter?.emit('mcp-client-update', this.clients);
+      if (!this.discoveryPromise) {
+        this.discoveryState = MCPDiscoveryState.COMPLETED;
+        this.eventEmitter?.emit('mcp-client-update', this.clients);
+      }
       return;
     }
 
@@ -574,7 +576,10 @@ export class McpClientManager {
     // If every configured server was skipped (for example because all are
     // disabled by user settings), no discovery promise is created. In that
     // case we must still mark discovery complete or the UI will wait forever.
-    if (this.discoveryState === MCPDiscoveryState.IN_PROGRESS) {
+    if (
+      this.discoveryState === MCPDiscoveryState.IN_PROGRESS &&
+      !this.discoveryPromise
+    ) {
       this.discoveryState = MCPDiscoveryState.COMPLETED;
       this.eventEmitter?.emit('mcp-client-update', this.clients);
     }