Browse Source

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
Hannes Rudolph 9 months ago
parent
commit
3f6e07adf1
2 changed files with 30 additions and 1 deletions
  1. 2 1
      src/core/webview/ClineProvider.ts
  2. 28 0
      src/services/mcp/McpHub.ts

+ 2 - 1
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")

+ 28 - 0
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) {