Browse Source

Merge pull request #852 from MuriloFP/mcp-server-fixes

MCP Server Bug Fixes Issues #833 and #834
Matt Rubens 1 year ago
parent
commit
324c5c8728
4 changed files with 125 additions and 8 deletions
  1. 6 4
      src/core/Cline.ts
  2. 29 2
      src/core/webview/ClineProvider.ts
  3. 7 2
      src/extension.ts
  4. 83 0
      src/services/mcp/McpServerManager.ts

+ 6 - 4
src/core/Cline.ts

@@ -832,7 +832,7 @@ export class Cline {
 		this.lastApiRequestTime = Date.now()
 
 		if (mcpEnabled ?? true) {
-			mcpHub = this.providerRef.deref()?.mcpHub
+			mcpHub = this.providerRef.deref()?.getMcpHub()
 			if (!mcpHub) {
 				throw new Error("MCP hub not available")
 			}
@@ -1013,7 +1013,7 @@ export class Cline {
 					// (have to do this for partial and complete since sending content in thinking tags to markdown renderer will automatically be removed)
 					// Remove end substrings of <thinking or </thinking (below xml parsing is only for opening tags)
 					// (this is done with the xml parsing below now, but keeping here for reference)
-					// content = content.replace(/<\/?t(?:h(?:i(?:n(?:k(?:i(?:n(?:g)?)?)?)?$/, "")
+					// content = content.replace(/<\/?t(?:h(?:i(?:n(?:k(?:i(?:n(?:g)?)?)?$/, "")
 					// Remove all instances of <thinking> (with optional line break after) and </thinking> (with optional line break before)
 					// - Needs to be separate since we dont want to remove the line break before the first tag
 					// - Needs to happen before the xml parsing below
@@ -2267,7 +2267,8 @@ export class Cline {
 								await this.say("mcp_server_request_started") // same as browser_action_result
 								const toolResult = await this.providerRef
 									.deref()
-									?.mcpHub?.callTool(server_name, tool_name, parsedArguments)
+									?.getMcpHub()
+									?.callTool(server_name, tool_name, parsedArguments)
 
 								// TODO: add progress indicator and ability to parse images and non-text responses
 								const toolResultPretty =
@@ -2335,7 +2336,8 @@ export class Cline {
 								await this.say("mcp_server_request_started")
 								const resourceResult = await this.providerRef
 									.deref()
-									?.mcpHub?.readResource(server_name, uri)
+									?.getMcpHub()
+									?.readResource(server_name, uri)
 								const resourceResultPretty =
 									resourceResult?.contents
 										.map((item) => {

+ 29 - 2
src/core/webview/ClineProvider.ts

@@ -36,6 +36,7 @@ import { EXPERIMENT_IDS, experiments as Experiments, experimentDefault, Experime
 import { CustomSupportPrompts, supportPrompt } from "../../shared/support-prompt"
 
 import { ACTION_NAMES } from "../CodeActionProvider"
+import { McpServerManager } from "../../services/mcp/McpServerManager"
 
 /*
 https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default/weather-webview/src/providers/WeatherViewProvider.ts
@@ -136,7 +137,7 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 	private isViewLaunched = false
 	private cline?: Cline
 	private workspaceTracker?: WorkspaceTracker
-	mcpHub?: McpHub
+	protected mcpHub?: McpHub // Change from private to protected
 	private latestAnnouncementId = "jan-21-2025-custom-modes" // update to some unique identifier when we add a new announcement
 	configManager: ConfigManager
 	customModesManager: CustomModesManager
@@ -148,11 +149,19 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		this.outputChannel.appendLine("ClineProvider instantiated")
 		ClineProvider.activeInstances.add(this)
 		this.workspaceTracker = new WorkspaceTracker(this)
-		this.mcpHub = new McpHub(this)
 		this.configManager = new ConfigManager(this.context)
 		this.customModesManager = new CustomModesManager(this.context, async () => {
 			await this.postStateToWebview()
 		})
+
+		// Initialize MCP Hub through the singleton manager
+		McpServerManager.getInstance(this.context, this)
+			.then((hub) => {
+				this.mcpHub = hub
+			})
+			.catch((error) => {
+				this.outputChannel.appendLine(`Failed to initialize MCP Hub: ${error}`)
+			})
 	}
 
 	/*
@@ -181,6 +190,9 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 		this.customModesManager?.dispose()
 		this.outputChannel.appendLine("Disposed all disposables")
 		ClineProvider.activeInstances.delete(this)
+
+		// Unregister from McpServerManager
+		McpServerManager.unregisterProvider(this)
 	}
 
 	public static getVisibleInstance(): ClineProvider | undefined {
@@ -601,6 +613,15 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 								this.postMessageToWebview({ type: "openRouterModels", openRouterModels })
 							}
 						})
+
+						// If MCP Hub is already initialized, update the webview with current server list
+						if (this.mcpHub) {
+							this.postMessageToWebview({
+								type: "mcpServers",
+								mcpServers: this.mcpHub.getServers(),
+							})
+						}
+
 						// gui relies on model info to be up-to-date to provide the most accurate pricing, so we need to fetch the latest details on launch.
 						// we do this for all users since many users switch between api providers and if they were to switch back to openrouter it would be showing outdated model info if we hadn't retrieved the latest at this point
 						// (see normalizeApiConfiguration > openrouter)
@@ -2103,6 +2124,7 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 			autoApprovalEnabled: autoApprovalEnabled ?? false,
 			customModes: await this.customModesManager.getCustomModes(),
 			experiments: experiments ?? experimentDefault,
+			mcpServers: this.mcpHub?.getServers() ?? [],
 		}
 	}
 
@@ -2538,4 +2560,9 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 	get messages() {
 		return this.cline?.clineMessages || []
 	}
+
+	// Add public getter
+	public getMcpHub(): McpHub | undefined {
+		return this.mcpHub
+	}
 }

+ 7 - 2
src/extension.ts

@@ -6,6 +6,7 @@ import "./utils/path" // Necessary to have access to String.prototype.toPosix.
 import { CodeActionProvider } from "./core/CodeActionProvider"
 import { DIFF_VIEW_URI_SCHEME } from "./integrations/editor/DiffViewProvider"
 import { handleUri, registerCommands, registerCodeActions, registerTerminalActions } from "./activate"
+import { McpServerManager } from "./services/mcp/McpServerManager"
 
 /**
  * Built using https://github.com/microsoft/vscode-webview-ui-toolkit
@@ -16,10 +17,12 @@ import { handleUri, registerCommands, registerCodeActions, registerTerminalActio
  */
 
 let outputChannel: vscode.OutputChannel
+let extensionContext: vscode.ExtensionContext
 
 // This method is called when your extension is activated.
 // Your extension is activated the very first time the command is executed.
 export function activate(context: vscode.ExtensionContext) {
+	extensionContext = context
 	outputChannel = vscode.window.createOutputChannel("Roo-Code")
 	context.subscriptions.push(outputChannel)
 	outputChannel.appendLine("Roo-Code extension activated")
@@ -83,7 +86,9 @@ export function activate(context: vscode.ExtensionContext) {
 	return createClineAPI(outputChannel, sidebarProvider)
 }
 
-// This method is called when your extension is deactivated.
-export function deactivate() {
+// This method is called when your extension is deactivated
+export async function deactivate() {
 	outputChannel.appendLine("Roo-Code extension deactivated")
+	// Clean up MCP server manager
+	await McpServerManager.cleanup(extensionContext)
 }

+ 83 - 0
src/services/mcp/McpServerManager.ts

@@ -0,0 +1,83 @@
+import * as vscode from "vscode"
+import { McpHub } from "./McpHub"
+import { ClineProvider } from "../../core/webview/ClineProvider"
+
+/**
+ * Singleton manager for MCP server instances.
+ * Ensures only one set of MCP servers runs across all webviews.
+ */
+export class McpServerManager {
+	private static instance: McpHub | null = null
+	private static readonly GLOBAL_STATE_KEY = "mcpHubInstanceId"
+	private static providers: Set<ClineProvider> = new Set()
+	private static initializationPromise: Promise<McpHub> | null = null
+
+	/**
+	 * Get the singleton McpHub instance.
+	 * Creates a new instance if one doesn't exist.
+	 * Thread-safe implementation using a promise-based lock.
+	 */
+	static async getInstance(context: vscode.ExtensionContext, provider: ClineProvider): Promise<McpHub> {
+		// Register the provider
+		this.providers.add(provider)
+
+		// If we already have an instance, return it
+		if (this.instance) {
+			return this.instance
+		}
+
+		// If initialization is in progress, wait for it
+		if (this.initializationPromise) {
+			return this.initializationPromise
+		}
+
+		// Create a new initialization promise
+		this.initializationPromise = (async () => {
+			try {
+				// Double-check instance in case it was created while we were waiting
+				if (!this.instance) {
+					this.instance = new McpHub(provider)
+					// Store a unique identifier in global state to track the primary instance
+					await context.globalState.update(this.GLOBAL_STATE_KEY, Date.now().toString())
+				}
+				return this.instance
+			} finally {
+				// Clear the initialization promise after completion or error
+				this.initializationPromise = null
+			}
+		})()
+
+		return this.initializationPromise
+	}
+
+	/**
+	 * Remove a provider from the tracked set.
+	 * This is called when a webview is disposed.
+	 */
+	static unregisterProvider(provider: ClineProvider): void {
+		this.providers.delete(provider)
+	}
+
+	/**
+	 * Notify all registered providers of server state changes.
+	 */
+	static notifyProviders(message: any): void {
+		this.providers.forEach((provider) => {
+			provider.postMessageToWebview(message).catch((error) => {
+				console.error("Failed to notify provider:", error)
+			})
+		})
+	}
+
+	/**
+	 * Clean up the singleton instance and all its resources.
+	 */
+	static async cleanup(context: vscode.ExtensionContext): Promise<void> {
+		if (this.instance) {
+			await this.instance.dispose()
+			this.instance = null
+			await context.globalState.update(this.GLOBAL_STATE_KEY, undefined)
+		}
+		this.providers.clear()
+	}
+}