migrate mcpServers (#4010)

Co-authored-by: Elephant Lumps <celestial_vault@Elephants-MacBook-Pro.local>

proto/mcp.proto

@@ -20,6 +20,9 @@ service McpService {
   // Subscribe to MCP marketplace catalog updates
   rpc subscribeToMcpMarketplaceCatalog(EmptyRequest) returns (stream McpMarketplaceCatalog);
   rpc getLatestMcpServers(Empty) returns (McpServers);
+  
+  // Subscribe to MCP server updates
+  rpc subscribeToMcpServers(EmptyRequest) returns (stream McpServers);
 }
 
 message ToggleMcpServerRequest {

src/core/controller/mcp/subscribeToMcpServers.ts

@@ -0,0 +1,76 @@
+import { Controller } from "../index"
+import { EmptyRequest } from "@shared/proto/common"
+import { McpServers } from "@shared/proto/mcp"
+import { StreamingResponseHandler, getRequestRegistry } from "../grpc-handler"
+import { convertMcpServersToProtoMcpServers } from "@shared/proto-conversions/mcp/mcp-server-conversion"
+
+// Keep track of active subscriptions
+const activeMcpServersSubscriptions = new Set<StreamingResponseHandler>()
+
+/**
+ * Subscribe to MCP servers events
+ * @param controller The controller instance
+ * @param request The empty request
+ * @param responseStream The streaming response handler
+ * @param requestId The ID of the request (passed by the gRPC handler)
+ */
+export async function subscribeToMcpServers(
+	controller: Controller,
+	request: EmptyRequest,
+	responseStream: StreamingResponseHandler,
+	requestId?: string,
+): Promise<void> {
+	// Add this subscription to the active subscriptions
+	activeMcpServersSubscriptions.add(responseStream)
+
+	// Register cleanup when the connection is closed
+	const cleanup = () => {
+		activeMcpServersSubscriptions.delete(responseStream)
+	}
+
+	// Register the cleanup function with the request registry if we have a requestId
+	if (requestId) {
+		getRequestRegistry().registerRequest(requestId, cleanup, { type: "mcpServers_subscription" }, responseStream)
+	}
+
+	// Send initial state if available
+	if (controller.mcpHub) {
+		const mcpServers = controller.mcpHub.getServers()
+		if (mcpServers.length > 0) {
+			try {
+				const protoServers = McpServers.create({
+					mcpServers: convertMcpServersToProtoMcpServers(mcpServers),
+				})
+				await responseStream(
+					protoServers,
+					false, // Not the last message
+				)
+			} catch (error) {
+				console.error("Error sending initial MCP servers:", error)
+				activeMcpServersSubscriptions.delete(responseStream)
+			}
+		}
+	}
+}
+
+/**
+ * Send an MCP servers update to all active subscribers
+ * @param mcpServers The MCP servers to send
+ */
+export async function sendMcpServersUpdate(mcpServers: McpServers): Promise<void> {
+	// Send the event to all active subscribers
+	const promises = Array.from(activeMcpServersSubscriptions).map(async (responseStream) => {
+		try {
+			await responseStream(
+				mcpServers,
+				false, // Not the last message
+			)
+		} catch (error) {
+			console.error("Error sending MCP servers update:", error)
+			// Remove the subscription if there was an error
+			activeMcpServersSubscriptions.delete(responseStream)
+		}
+	})
+
+	await Promise.all(promises)
+}

src/core/controller/mcp/updateMcpTimeout.ts

@@ -12,7 +12,6 @@ export async function updateMcpTimeout(controller: Controller, request: UpdateMc
 	try {
 		if (request.serverName && typeof request.serverName === "string" && typeof request.timeout === "number") {
 			const mcpServers = await controller.mcpHub?.updateServerTimeoutRPC(request.serverName, request.timeout)
-			console.log("mcpServers", mcpServers)
 			const convertedMcpServers = convertMcpServersToProtoMcpServers(mcpServers)
 			console.log("convertedMcpServers", convertedMcpServers)
 			return McpServers.create({ mcpServers: convertedMcpServers })

src/services/mcp/McpHub.ts

@@ -10,6 +10,8 @@ import {
 	ListToolsResultSchema,
 	ReadResourceResultSchema,
 } from "@modelcontextprotocol/sdk/types.js"
+import { sendMcpServersUpdate } from "@core/controller/mcp/subscribeToMcpServers"
+import { convertMcpServersToProtoMcpServers } from "@shared/proto-conversions/mcp/mcp-server-conversion"
 import chokidar, { FSWatcher } from "chokidar"
 import { setTimeout as setTimeoutPromise } from "node:timers/promises"
 import deepEqual from "fast-deep-equal"
@@ -606,9 +608,13 @@ export class McpHub {
 		const content = await fs.readFile(settingsPath, "utf-8")
 		const config = JSON.parse(content)
 		const serverOrder = Object.keys(config.mcpServers || {})
-		await this.postMessageToWebview({
-			type: "mcpServers",
-			mcpServers: this.getSortedMcpServers(serverOrder),
+
+		// Get sorted servers
+		const sortedServers = this.getSortedMcpServers(serverOrder)
+
+		// Send update using gRPC stream
+		await sendMcpServersUpdate({
+			mcpServers: convertMcpServersToProtoMcpServers(sortedServers),
 		})
 	}
 

src/shared/ExtensionMessage.ts

@@ -19,7 +19,6 @@ export interface ExtensionMessage {
 		| "selectedImages"
 		| "openAiModels"
 		| "requestyModels"
-		| "mcpServers"
 		| "mcpDownloadDetails"
 		| "userCreditsBalance"
 		| "userCreditsUsage"

webview-ui/src/context/ExtensionStateContext.tsx

@@ -26,6 +26,7 @@ import {
 	requestyDefaultModelInfo,
 } from "../../../src/shared/api"
 import { McpMarketplaceCatalog, McpServer, McpViewTab } from "../../../src/shared/mcp"
+import { convertProtoMcpServersToMcpServers } from "@shared/proto-conversions/mcp/mcp-server-conversion"
 import { convertTextMateToHljs } from "../utils/textMateToHljs"
 import { vscode } from "../utils/vscode"
 import { OpenRouterCompatibleModelInfo } from "@shared/proto/models"
@@ -223,10 +224,6 @@ export const ExtensionStateContextProvider: React.FC<{
 				})
 				break
 			}
-			case "mcpServers": {
-				setMcpServers(message.mcpServers ?? [])
-				break
-			}
 		}
 	}, [])
 
@@ -256,6 +253,7 @@ export const ExtensionStateContextProvider: React.FC<{
 			relinquishControlCallbacks.current.delete(callback)
 		}
 	}, [])
+	const mcpServersSubscriptionRef = useRef<(() => void) | null>(null)
 
 	// Subscribe to state updates and UI events using the gRPC streaming API
 	useEffect(() => {
@@ -385,6 +383,22 @@ export const ExtensionStateContextProvider: React.FC<{
 			onComplete: () => {},
 		})
 
+		// Subscribe to MCP servers updates
+		mcpServersSubscriptionRef.current = McpServiceClient.subscribeToMcpServers(EmptyRequest.create(), {
+			onResponse: (response) => {
+				console.log("[DEBUG] Received MCP servers update from gRPC stream")
+				if (response.mcpServers) {
+					setMcpServers(convertProtoMcpServersToMcpServers(response.mcpServers))
+				}
+			},
+			onError: (error) => {
+				console.error("Error in MCP servers subscription:", error)
+			},
+			onComplete: () => {
+				console.log("MCP servers subscription completed")
+			},
+		})
+
 		// Subscribe to workspace file updates
 		workspaceUpdatesUnsubscribeRef.current = FileServiceClient.subscribeToWorkspaceUpdates(EmptyRequest.create({}), {
 			onResponse: (response) => {
@@ -592,6 +606,11 @@ export const ExtensionStateContextProvider: React.FC<{
 				relinquishControlUnsubscribeRef.current()
 				relinquishControlUnsubscribeRef.current = null
 			}
+
+			if (mcpServersSubscriptionRef.current) {
+				mcpServersSubscriptionRef.current()
+				mcpServersSubscriptionRef.current = null
+			}
 		}
 	}, [])