From 361b0cfa2bd476db2d63ef166d8606f5a2f07234 Mon Sep 17 00:00:00 2001 From: celestial-vault <58194240+celestial-vault@users.noreply.github.com> Date: Sat, 7 Jun 2025 15:15:59 -0700 Subject: [PATCH] migrate mcpServers (#4010) Co-authored-by: Elephant Lumps --- proto/mcp.proto | 3 + .../controller/mcp/subscribeToMcpServers.ts | 76 +++++++++++++++++++ src/core/controller/mcp/updateMcpTimeout.ts | 1 - src/services/mcp/McpHub.ts | 12 ++- src/shared/ExtensionMessage.ts | 1 - .../src/context/ExtensionStateContext.tsx | 27 ++++++- 6 files changed, 111 insertions(+), 9 deletions(-) create mode 100644 src/core/controller/mcp/subscribeToMcpServers.ts diff --git a/proto/mcp.proto b/proto/mcp.proto index 63a030924f..f656b275d8 100644 --- a/proto/mcp.proto +++ b/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 { diff --git a/src/core/controller/mcp/subscribeToMcpServers.ts b/src/core/controller/mcp/subscribeToMcpServers.ts new file mode 100644 index 0000000000..10c51598c3 --- /dev/null +++ b/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() + +/** + * 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 { + // 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 { + // 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) +} diff --git a/src/core/controller/mcp/updateMcpTimeout.ts b/src/core/controller/mcp/updateMcpTimeout.ts index 3db67a16e8..327dcfc480 100644 --- a/src/core/controller/mcp/updateMcpTimeout.ts +++ b/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 }) diff --git a/src/services/mcp/McpHub.ts b/src/services/mcp/McpHub.ts index d7d28a06db..dbfdb8f638 100644 --- a/src/services/mcp/McpHub.ts +++ b/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), }) } diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index bb061155c5..eabe5a19fa 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -19,7 +19,6 @@ export interface ExtensionMessage { | "selectedImages" | "openAiModels" | "requestyModels" - | "mcpServers" | "mcpDownloadDetails" | "userCreditsBalance" | "userCreditsUsage" diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index 03857f4367..ec9cff1ef9 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/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 + } } }, [])