Skip to content

Commit 4859c83

Browse files
bingggbinggg
committed
refactor(tool-registration): move conditionalRegisterTool logic into wrapServerWithTelemetry 🚀
1 parent 6088827 commit 4859c83

9 files changed

Lines changed: 35 additions & 45 deletions

File tree

‎mcp/src/index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ export type {
2727

2828
export { getLoginState, logout } from "./auth.js";
2929

30-
export { isCloudMode, enableCloudMode, getCloudModeStatus, shouldRegisterTool, conditionalRegisterTool } from "./utils/cloud-mode.js";
30+
export { isCloudMode, enableCloudMode, getCloudModeStatus, shouldRegisterTool } from "./utils/cloud-mode.js";
3131

3232
export {
3333
getCloudBaseManager,

‎mcp/src/tools/download.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import * as https from "https";
88
import * as http from "http";
99
import { URL } from "url";
1010
import * as net from "net";
11-
import { conditionalRegisterTool } from '../utils/cloud-mode.js';
11+
1212
import * as dns from "dns";
1313
import { getCloudBaseManager } from '../cloudbase-manager.js'
1414
import { ExtendedMcpServer } from '../server.js';
@@ -245,8 +245,7 @@ function downloadFile(url: string): Promise<{
245245

246246
export function registerDownloadTools(server: ExtendedMcpServer) {
247247
// downloadRemoteFile - 下载远程文件 (cloud-incompatible)
248-
conditionalRegisterTool(
249-
server,
248+
server.registerTool(
250249
"downloadRemoteFile",
251250
{
252251
title: "下载远程文件",

‎mcp/src/tools/functions.ts‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { z } from "zod";
22
import { getCloudBaseManager } from '../cloudbase-manager.js'
33
import { ExtendedMcpServer } from '../server.js';
4-
import { conditionalRegisterTool } from '../utils/cloud-mode.js';
4+
55
import path from 'path';
66

77
// 支持的 Node.js 运行时列表
@@ -96,8 +96,7 @@ export function registerFunctionTools(server: ExtendedMcpServer) {
9696
);
9797

9898
// createFunction - 创建云函数 (cloud-incompatible)
99-
conditionalRegisterTool(
100-
server,
99+
server.registerTool(
101100
"createFunction",
102101
{
103102
title: "创建云函数",
@@ -185,8 +184,7 @@ export function registerFunctionTools(server: ExtendedMcpServer) {
185184
);
186185

187186
// updateFunctionCode - 更新函数代码 (cloud-incompatible)
188-
conditionalRegisterTool(
189-
server,
187+
server.registerTool(
190188
"updateFunctionCode",
191189
{
192190
title: "更新云函数代码",

‎mcp/src/tools/hosting.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { z } from "zod";
22
import { getCloudBaseManager } from '../cloudbase-manager.js'
33
import { ExtendedMcpServer } from '../server.js';
4-
import { conditionalRegisterTool } from '../utils/cloud-mode.js';
4+
55

66
// 定义扩展的EnvInfo接口，包含StaticStorages属性
77
interface ExtendedEnvInfo {
@@ -23,8 +23,7 @@ export function registerHostingTools(server: ExtendedMcpServer) {
2323
const getManager = () => getCloudBaseManager({ cloudBaseOptions });
2424

2525
// uploadFiles - 上传文件到静态网站托管 (cloud-incompatible)
26-
conditionalRegisterTool(
27-
server,
26+
server.registerTool(
2827
"uploadFiles",
2928
{
3029
title: "上传静态文件",

‎mcp/src/tools/interactive.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,11 @@ import fs from 'fs/promises';
77
import path from 'path';
88
import os from 'os';
99
import { ExtendedMcpServer } from '../server.js';
10-
import { conditionalRegisterTool } from '../utils/cloud-mode.js';
10+
1111

1212
export function registerInteractiveTools(server: ExtendedMcpServer) {
1313
// 统一的交互式对话工具 (cloud-incompatible)
14-
conditionalRegisterTool(
15-
server,
14+
server.registerTool(
1615
"interactiveDialog",
1716
{
1817
title: "交互式对话",

‎mcp/src/tools/setup.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import * as http from "http";
88
import { execSync } from "child_process";
99
import AdmZip from "adm-zip";
1010
import { ExtendedMcpServer } from '../server.js';
11-
import { conditionalRegisterTool } from '../utils/cloud-mode.js';
11+
1212

1313
// 构建时注入的版本号
1414
// @ts-ignore
@@ -315,8 +315,7 @@ function filterFilesByIDE(files: string[], ide: string): string[] {
315315

316316
export function registerSetupTools(server: ExtendedMcpServer) {
317317
// downloadTemplate - 下载项目模板 (cloud-incompatible)
318-
conditionalRegisterTool(
319-
server,
318+
server.registerTool(
320319
"downloadTemplate",
321320
{
322321
title: "下载项目模板",

‎mcp/src/tools/storage.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { z } from "zod";
22
import { getCloudBaseManager } from '../cloudbase-manager.js'
33
import { ExtendedMcpServer } from '../server.js';
4-
import { conditionalRegisterTool } from '../utils/cloud-mode.js';
4+
55

66
export function registerStorageTools(server: ExtendedMcpServer) {
77
// 获取 cloudBaseOptions，如果没有则为 undefined
@@ -11,8 +11,7 @@ export function registerStorageTools(server: ExtendedMcpServer) {
1111
const getManager = () => getCloudBaseManager({ cloudBaseOptions });
1212

1313
// uploadFile - 上传文件到云存储 (cloud-incompatible)
14-
conditionalRegisterTool(
15-
server,
14+
server.registerTool(
1615
"uploadFile",
1716
{
1817
title: "上传文件到云存储",

‎mcp/src/utils/cloud-mode.ts‎

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ export function shouldRegisterTool(toolName: string): boolean {
6969

7070
// Cloud-incompatible tools that involve local file operations
7171
const cloudIncompatibleTools = [
72+
// Auth tools - local file uploads
73+
'login',
74+
'logout',
75+
7276
// Storage tools - local file uploads
7377
'uploadFile',
7478

@@ -84,6 +88,7 @@ export function shouldRegisterTool(toolName: string): boolean {
8488

8589
// Download tools - local file downloads
8690
'downloadTemplate',
91+
'downloadRemoteFile',
8792

8893
// Setup tools - local config file operations
8994
'setupEnvironmentId',
@@ -102,19 +107,4 @@ export function shouldRegisterTool(toolName: string): boolean {
102107
return shouldRegister;
103108
}
104109

105-
/**
106-
* Conditional tool registration wrapper
107-
* Only registers the tool if it's compatible with current mode
108-
*/
109-
export function conditionalRegisterTool(
110-
server: any,
111-
toolName: string,
112-
toolConfig: any,
113-
handler: any
114-
): void {
115-
if (shouldRegisterTool(toolName)) {
116-
server.registerTool?.(toolName, toolConfig, handler);
117-
} else {
118-
debug(`Skipped registering tool '${toolName}' in cloud mode`);
119-
}
120-
}
110+

‎mcp/src/utils/tool-wrapper.ts‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { ToolAnnotations, Tool } from "@modelcontextprotocol/sdk/types.js";
33
import { reportToolCall } from './telemetry.js';
44
import { debug } from './logger.js';
55
import { CloudBaseOptions } from '../types.js';
6+
import { shouldRegisterTool } from './cloud-mode.js';
67
import os from 'os';
78

89
// 扩展 McpServer 类型以包含 ide
@@ -139,25 +140,31 @@ function createWrappedHandler(name: string, handler: any, server: ExtendedMcpSer
139140
}
140141

141142
/**
142-
* 包装 MCP Server 的 registerTool 方法，添加数据上报功能
143+
* 包装 MCP Server 的 registerTool 方法，添加数据上报功能和条件注册
143144
* @param server MCP Server 实例
144145
*/
145146
export function wrapServerWithTelemetry(server: McpServer): void {
146147
// 保存原始的 registerTool 方法
147148
const originalRegisterTool = server.registerTool.bind(server);
148149

149-
// 重写 registerTool 方法，添加数据上报功能
150+
// Override the registerTool method to add telemetry and conditional registration
150151
server.registerTool = function(toolName: string, toolConfig: any, handler: any) {
151-
152-
// 记录工具注册信息
153-
debug(`注册工具: ${toolName}`, {
152+
// If the tool should not be registered in the current mode, do not register and return undefined
153+
if (!shouldRegisterTool(toolName)) {
154+
debug(`Cloud mode: skipping registration of incompatible tool: ${toolName}`);
155+
// Explicitly return undefined to satisfy the expected type
156+
return undefined as any;
157+
}
158+
159+
// Log tool registration info
160+
debug(`Registering tool: ${toolName}`, {
154161
toolConfig
155162
});
156163

157-
// 使用包装后的处理函数，传递服务器实例
164+
// Use the wrapped handler, passing the server instance
158165
const wrappedHandler = createWrappedHandler(toolName, handler, server as ExtendedMcpServer);
159-
160-
// 调用原始 registerTool 方法
166+
167+
// Call the original registerTool method
161168
return originalRegisterTool(toolName, toolConfig, wrappedHandler);
162169
};
163170
}

0 commit comments

Comments
 (0)