Register tool functions of mcp servers again after restart of the frontend (#14723)

JonasHelming · web-flow · commit 8c638f2618cf · 2025-01-16T10:22:58.000+01:00
fixed #14701
diff --git a/packages/ai-mcp/src/browser/mcp-command-contribution.ts b/packages/ai-mcp/src/browser/mcp-command-contribution.ts
@@ -17,10 +17,7 @@ import { AICommandHandlerFactory } from '@theia/ai-core/lib/browser/ai-command-h
 import { CommandContribution, CommandRegistry, MessageService } from '@theia/core';
 import { QuickInputService } from '@theia/core/lib/browser';
 import { inject, injectable } from '@theia/core/shared/inversify';
-import { MCPServerManager } from '../common/mcp-server-manager';
-import { ToolInvocationRegistry, ToolRequest } from '@theia/ai-core';
-
-type MCPTool = Awaited<ReturnType<MCPServerManager['getTools']>>['tools'][number];
+import { MCPFrontendService } from './mcp-frontend-service';
 
 export const StartMCPServer = {
     id: 'mcp.startserver',
@@ -42,11 +39,8 @@ export class MCPCommandContribution implements CommandContribution {
     @inject(MessageService)
     protected messageService: MessageService;
 
-    @inject(MCPServerManager)
-    protected readonly mcpServerManager: MCPServerManager;
-
-    @inject(ToolInvocationRegistry)
-    protected readonly toolInvocationRegistry: ToolInvocationRegistry;
+    @inject(MCPFrontendService)
+    protected readonly mcpFrontendService: MCPFrontendService;
 
     private async getMCPServerSelection(serverNames: string[]): Promise<string | undefined> {
         if (!serverNames || serverNames.length === 0) {
@@ -61,7 +55,7 @@ export class MCPCommandContribution implements CommandContribution {
         commandRegistry.registerCommand(StopMCPServer, this.commandHandlerFactory({
             execute: async () => {
                 try {
-                    const startedServers = await this.mcpServerManager.getStartedServers();
+                    const startedServers = await this.mcpFrontendService.getStartedServers();
                     if (!startedServers || startedServers.length === 0) {
                         this.messageService.error('No MCP servers running.');
                         return;
@@ -70,8 +64,7 @@ export class MCPCommandContribution implements CommandContribution {
                     if (!selection) {
                         return;
                     }
-                    this.toolInvocationRegistry.unregisterAllTools(`mcp_${selection}`);
-                    this.mcpServerManager.stopServer(selection);
+                    await this.mcpFrontendService.stopServer(selection);
                 } catch (error) {
                     console.error('Error while stopping MCP server:', error);
                 }
@@ -81,8 +74,8 @@ export class MCPCommandContribution implements CommandContribution {
         commandRegistry.registerCommand(StartMCPServer, this.commandHandlerFactory({
             execute: async () => {
                 try {
-                    const servers = await this.mcpServerManager.getServerNames();
-                    const startedServers = await this.mcpServerManager.getStartedServers();
+                    const servers = await this.mcpFrontendService.getServerNames();
+                    const startedServers = await this.mcpFrontendService.getStartedServers();
                     const startableServers = servers.filter(server => !startedServers.includes(server));
                     if (!startableServers || startableServers.length === 0) {
                         if (startedServers && startedServers.length > 0) {
@@ -97,13 +90,8 @@ export class MCPCommandContribution implements CommandContribution {
                     if (!selection) {
                         return;
                     }
-                    this.mcpServerManager.startServer(selection);
-                    const { tools } = await this.mcpServerManager.getTools(selection);
-                    const toolRequests: ToolRequest[] = tools.map(tool => this.convertToToolRequest(tool, selection));
-
-                    for (const toolRequest of toolRequests) {
-                        this.toolInvocationRegistry.registerTool(toolRequest);
-                    }
+                    await this.mcpFrontendService.startServer(selection);
+                    const { tools } = await this.mcpFrontendService.getTools(selection);
                     const toolNames = tools.map(tool => tool.name || 'Unnamed Tool').join(', ');
                     this.messageService.info(
                         `MCP server "${selection}" successfully started. Registered tools: ${toolNames || 'No tools available.'}`
@@ -115,29 +103,4 @@ export class MCPCommandContribution implements CommandContribution {
             }
         }));
     }
-
-    convertToToolRequest(tool: MCPTool, serverName: string): ToolRequest {
-        const id = `mcp_${serverName}_${tool.name}`;
-
-        return {
-            id: id,
-            name: id,
-            providerName: `mcp_${serverName}`,
-            parameters: ToolRequest.isToolRequestParameters(tool.inputSchema) ? {
-                type: tool.inputSchema.type,
-                properties: tool.inputSchema.properties,
-                required: tool.inputSchema.required
-            } : undefined,
-            description: tool.description,
-            handler: async (arg_string: string) => {
-                try {
-                    return await this.mcpServerManager.callTool(serverName, tool.name, arg_string);
-                } catch (error) {
-                    console.error(`Error in tool handler for ${tool.name} on MCP server ${serverName}:`, error);
-                    throw error;
-                }
-            },
-        };
-    }
-
 }
diff --git a/packages/ai-mcp/src/browser/mcp-frontend-application-contribution.ts b/packages/ai-mcp/src/browser/mcp-frontend-application-contribution.ts
@@ -18,6 +18,7 @@ import { inject, injectable } from '@theia/core/shared/inversify';
 import { MCPServerDescription, MCPServerManager } from '../common';
 import { MCP_SERVERS_PREF } from './mcp-preferences';
 import { JSONObject } from '@theia/core/shared/@phosphor/coreutils';
+import { MCPFrontendService } from './mcp-frontend-service';
 
 interface MCPServersPreferenceValue {
     command: string;
@@ -60,6 +61,9 @@ export class McpFrontendApplicationContribution implements FrontendApplicationCo
     @inject(MCPServerManager)
     protected manager: MCPServerManager;
 
+    @inject(MCPFrontendService)
+    protected frontendMCPService: MCPFrontendService;
+
     protected prevServers: Map<string, MCPServerDescription> = new Map();
 
     onStart(): void {
@@ -77,6 +81,7 @@ export class McpFrontendApplicationContribution implements FrontendApplicationCo
                 }
             });
         });
+        this.frontendMCPService.registerToolsForAllStartedServers();
     }
 
     protected handleServerChanges(newServers: MCPServersPreference): void {
diff --git a/packages/ai-mcp/src/browser/mcp-frontend-module.ts b/packages/ai-mcp/src/browser/mcp-frontend-module.ts
@@ -21,6 +21,7 @@ import { FrontendApplicationContribution, PreferenceContribution, RemoteConnecti
 import { MCPServerManager, MCPServerManagerPath } from '../common/mcp-server-manager';
 import { McpServersPreferenceSchema } from './mcp-preferences';
 import { McpFrontendApplicationContribution } from './mcp-frontend-application-contribution';
+import { MCPFrontendService } from './mcp-frontend-service';
 
 export default new ContainerModule(bind => {
     bind(PreferenceContribution).toConstantValue({ schema: McpServersPreferenceSchema });
@@ -30,4 +31,5 @@ export default new ContainerModule(bind => {
         const connection = ctx.container.get<ServiceConnectionProvider>(RemoteConnectionProvider);
         return connection.createProxy<MCPServerManager>(MCPServerManagerPath);
     }).inSingletonScope();
+    bind(MCPFrontendService).toSelf().inSingletonScope();
 });
diff --git a/packages/ai-mcp/src/browser/mcp-frontend-service.ts b/packages/ai-mcp/src/browser/mcp-frontend-service.ts
@@ -0,0 +1,87 @@
+// *****************************************************************************
+// Copyright (C) 2024 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+import { injectable, inject } from '@theia/core/shared/inversify';
+import { MCPServer, MCPServerManager } from '../common/mcp-server-manager';
+import { ToolInvocationRegistry, ToolRequest } from '@theia/ai-core';
+
+@injectable()
+export class MCPFrontendService {
+    @inject(MCPServerManager)
+    protected readonly mcpServerManager: MCPServerManager;
+
+    @inject(ToolInvocationRegistry)
+    protected readonly toolInvocationRegistry: ToolInvocationRegistry;
+
+    async startServer(serverName: string): Promise<void> {
+        await this.mcpServerManager.startServer(serverName);
+        this.registerTools(serverName);
+    }
+
+    async registerToolsForAllStartedServers(): Promise<void> {
+        const startedServers = await this.getStartedServers();
+        for (const serverName of startedServers) {
+            await this.registerTools(serverName);
+        }
+    }
+
+    async registerTools(serverName: string): Promise<void> {
+        const { tools } = await this.getTools(serverName);
+        const toolRequests: ToolRequest[] = tools.map(tool => this.convertToToolRequest(tool, serverName));
+        toolRequests.forEach(toolRequest =>
+            this.toolInvocationRegistry.registerTool(toolRequest)
+        );
+    }
+
+    async stopServer(serverName: string): Promise<void> {
+        this.toolInvocationRegistry.unregisterAllTools(`mcp_${serverName}`);
+        await this.mcpServerManager.stopServer(serverName);
+    }
+
+    getStartedServers(): Promise<string[]> {
+        return this.mcpServerManager.getStartedServers();
+    }
+
+    getServerNames(): Promise<string[]> {
+        return this.mcpServerManager.getServerNames();
+    }
+
+    getTools(serverName: string): ReturnType<MCPServer['getTools']> {
+        return this.mcpServerManager.getTools(serverName);
+    }
+
+    private convertToToolRequest(tool: Awaited<ReturnType<MCPServerManager['getTools']>>['tools'][number], serverName: string): ToolRequest {
+        const id = `mcp_${serverName}_${tool.name}`;
+        return {
+            id: id,
+            name: id,
+            providerName: `mcp_${serverName}`,
+            parameters: ToolRequest.isToolRequestParameters(tool.inputSchema) ? {
+                type: tool.inputSchema.type,
+                properties: tool.inputSchema.properties,
+                required: tool.inputSchema.required
+            } : undefined,
+            description: tool.description,
+            handler: async (arg_string: string) => {
+                try {
+                    return await this.mcpServerManager.callTool(serverName, tool.name, arg_string);
+                } catch (error) {
+                    console.error(`Error in tool handler for ${tool.name} on MCP server ${serverName}:`, error);
+                    throw error;
+                }
+            },
+        };
+    }
+}
