From 3f6e07adf126aa7454d7204e5b9b4a47bdd0075a Mon Sep 17 00:00:00 2001
From: Hannes Rudolph <hrudolph@gmail.com>
Date: Fri, 4 Apr 2025 21:14:05 -0600
Subject: [PATCH] refactor(mcp): Implement reference counting for McpHub
 lifecycle (#2310)

Problem:
Closing auxiliary windows (e.g., from "open in editor") could prematurely dispose the shared singleton McpHub instance, shutting down MCP servers unexpectedly while other panels (like the sidebar) might still be active.

Solution:
Implemented reference counting directly within the McpHub class to manage its own lifecycle based on active clients (ClineProvider instances).
- Added `refCount`, `registerClient()`, and `unregisterClient()` methods to McpHub.
- McpHub now disposes itself only when the last registered client unregisters (`refCount` reaches 0).
- ClineProvider instances now call `mcpHub.registerClient()` upon initialization and `mcpHub.unregisterClient()` upon disposal.
- Removed the direct `mcpHub.dispose()` call from ClineProvider.dispose.

Benefit:
This ensures the shared McpHub instance remains active as long as at least one ClineProvider instance is using it. Cleanup now correctly occurs only when the last provider is closed or during full extension deactivation. This centralizes the resource's lifecycle logic within the resource class itself.

Files Changed:
- src/services/mcp/McpHub.ts
- src/core/webview/ClineProvider.ts
---
 src/core/webview/ClineProvider.ts |  3 ++-
 src/services/mcp/McpHub.ts        | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts
index 060e79a68d..c38126ec0f 100644
--- a/src/core/webview/ClineProvider.ts
+++ b/src/core/webview/ClineProvider.ts
@@ -115,6 +115,7 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 		McpServerManager.getInstance(this.context, this)
 			.then((hub) => {
 				this.mcpHub = hub
+				this.mcpHub.registerClient()
 			})
 			.catch((error) => {
 				this.outputChannel.appendLine(`Failed to initialize MCP Hub: ${error}`)
@@ -221,7 +222,7 @@ export class ClineProvider extends EventEmitter<ClineProviderEvents> implements
 
 		this.workspaceTracker?.dispose()
 		this.workspaceTracker = undefined
-		this.mcpHub?.dispose()
+		await this.mcpHub?.unregisterClient()
 		this.mcpHub = undefined
 		this.customModesManager?.dispose()
 		this.outputChannel.appendLine("Disposed all disposables")
diff --git a/src/services/mcp/McpHub.ts b/src/services/mcp/McpHub.ts
index 037ff6a9f9..31d0dd8020 100644
--- a/src/services/mcp/McpHub.ts
+++ b/src/services/mcp/McpHub.ts
@@ -109,6 +109,7 @@ export class McpHub {
 	private isDisposed: boolean = false
 	connections: McpConnection[] = []
 	isConnecting: boolean = false
+	private refCount: number = 0 // Reference counter for active clients
 
 	constructor(provider: ClineProvider) {
 		this.providerRef = new WeakRef(provider)
@@ -118,6 +119,27 @@ export class McpHub {
 		this.initializeGlobalMcpServers()
 		this.initializeProjectMcpServers()
 	}
+	/**
+	 * Registers a client (e.g., ClineProvider) using this hub.
+	 * Increments the reference count.
+	 */
+	public registerClient(): void {
+		this.refCount++
+		console.log(`McpHub: Client registered. Ref count: ${this.refCount}`)
+	}
+
+	/**
+	 * Unregisters a client. Decrements the reference count.
+	 * If the count reaches zero, disposes the hub.
+	 */
+	public async unregisterClient(): Promise<void> {
+		this.refCount--
+		console.log(`McpHub: Client unregistered. Ref count: ${this.refCount}`)
+		if (this.refCount <= 0) {
+			console.log("McpHub: Last client unregistered. Disposing hub.")
+			await this.dispose()
+		}
+	}
 
 	/**
 	 * Validates and normalizes server configuration
@@ -1247,6 +1269,12 @@ export class McpHub {
 	}
 
 	async dispose(): Promise<void> {
+		// Prevent multiple disposals
+		if (this.isDisposed) {
+			console.log("McpHub: Already disposed.")
+			return
+		}
+		console.log("McpHub: Disposing...")
 		this.isDisposed = true
 		this.removeAllFileWatchers()
 		for (const connection of this.connections) {