Files
admin f3a44a901e
Build and Deploy / build-and-push (push) Successful in 54s
refactor(mcp): consolidate MCP implementation into internal package
- Move MCP tools registration to internal/mcp
- cmd/mcp now only contains CLI entry logic
- Support env vars TASKPOOL_URL/TASKPOOL_TOKEN for stdio mode
- Remove duplicate client.go from cmd/mcp
2026-07-27 02:43:01 +08:00

589 lines
21 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package mcp
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
"strings"
"sync"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
const (
serverName = "taskpool"
serverVersion = "1.0.0"
)
var (
mcpServer *server.MCPServer
mcpHTTPServer *server.StreamableHTTPServer
externalClient *Client // 外部客户端(用于 stdio 模式连接远程服务器)
once sync.Once
)
// GetServer 获取 MCP Server 实例(单例)
func GetServer() *server.MCPServer {
once.Do(func() {
mcpServer = server.NewMCPServer(
serverName,
serverVersion,
server.WithToolCapabilities(true),
server.WithResourceCapabilities(false, false),
server.WithPromptCapabilities(true),
server.WithInstructions(instructions),
server.WithRecovery(),
)
registerTools(mcpServer)
registerPrompts(mcpServer)
})
return mcpServer
}
// SetExternalClient 设置外部 OpenAPI 客户端(用于 stdio 模式)
func SetExternalClient(c *Client) {
externalClient = c
}
// getClient 获取当前使用的客户端
func getClient() *Client {
if externalClient != nil {
return externalClient
}
return GetOpenAPIClient()
}
// GetHTTPHandler 获取 MCP HTTP Handler(用于集成到 Gin
func GetHTTPHandler() http.Handler {
s := GetServer()
if mcpHTTPServer == nil {
mcpHTTPServer = server.NewStreamableHTTPServer(s,
server.WithEndpointPath(""),
)
}
return mcpHTTPServer
}
// ServeStdio 启动 stdio 模式的 MCP Server
func ServeStdio(s *server.MCPServer) error {
return server.ServeStdio(s)
}
const instructions = `你是 TaskPool(任务池)运维助手。通过 MCP tools 管理任务调度面板。
核心能力:
- 任务:列表/详情/创建/更新/删除/启停执行
- 脚本:列表/详情/创建/更新/删除
- 环境变量:列表/详情/创建/更新/删除(secret 值为脱敏)
- 日志:列表/详情、最近执行结果
使用原则:
1. 先 list_tasks / get_task 确认目标,再 run_task
2. 删除类操作前向用户确认
3. 不要编造任务 IDID 以 list/get 返回为准
4. 执行后用 list_logs 或 get_log 查看结果
5. 创建任务时 command 必填(普通任务);schedule 使用标准 cron`
func registerTools(s *server.MCPServer) {
// ---- tasks ----
s.AddTool(mcp.NewTool("list_tasks",
mcp.WithDescription("分页列出任务,可按名称/标签/启用状态筛选"),
mcp.WithString("name", mcp.Description("按名称模糊筛选")),
mcp.WithString("tags", mcp.Description("按标签筛选")),
mcp.WithString("enabled", mcp.Description("true/false,是否启用")),
mcp.WithNumber("page", mcp.Description("页码,默认 1")),
mcp.WithNumber("page_size", mcp.Description("每页数量,默认 20")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
q := map[string]string{
"name": strArg(req, "name"),
"tags": strArg(req, "tags"),
"enabled": strArg(req, "enabled"),
"page": numArgStr(req, "page", "1"),
"page_size": numArgStr(req, "page_size", "20"),
}
raw, err := GetOpenAPIClient().Get("/tasks", q)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("get_task",
mcp.WithDescription("获取单个任务详情"),
mcp.WithString("id", mcp.Required(), mcp.Description("任务 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Get("/tasks/"+id, nil)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("create_task",
mcp.WithDescription("创建任务。普通任务必须提供 command;可选 schedule(cron)、tags、timeout、work_dir、remark、enabled 等"),
mcp.WithString("name", mcp.Required(), mcp.Description("任务名称")),
mcp.WithString("command", mcp.Description("执行命令,普通任务必填")),
mcp.WithString("schedule", mcp.Description("Cron 表达式,如 0 0 * * *")),
mcp.WithString("remark", mcp.Description("备注")),
mcp.WithString("tags", mcp.Description("标签,逗号分隔")),
mcp.WithString("work_dir", mcp.Description("工作目录")),
mcp.WithNumber("timeout", mcp.Description("超时秒数")),
mcp.WithString("type", mcp.Description("任务类型,默认 common;仓库同步为 repo")),
mcp.WithString("config", mcp.Description("JSON 字符串配置(repo 任务用)")),
mcp.WithString("envs", mcp.Description("环境变量 JSON 字符串")),
mcp.WithString("trigger_type", mcp.Description("触发类型,如 cron")),
mcp.WithNumber("retry_count", mcp.Description("重试次数")),
mcp.WithNumber("retry_interval", mcp.Description("重试间隔秒")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
body := map[string]any{}
setStr(body, "name", strArg(req, "name"))
setStr(body, "command", strArg(req, "command"))
setStr(body, "schedule", strArg(req, "schedule"))
setStr(body, "remark", strArg(req, "remark"))
setStr(body, "tags", strArg(req, "tags"))
setStr(body, "work_dir", strArg(req, "work_dir"))
setStr(body, "type", strArg(req, "type"))
setStr(body, "config", strArg(req, "config"))
setStr(body, "envs", strArg(req, "envs"))
setStr(body, "trigger_type", strArg(req, "trigger_type"))
if n, ok := numArg(req, "timeout"); ok {
body["timeout"] = n
}
if n, ok := numArg(req, "retry_count"); ok {
body["retry_count"] = n
}
if n, ok := numArg(req, "retry_interval"); ok {
body["retry_interval"] = n
}
if body["name"] == nil || body["name"] == "" {
return mcp.NewToolResultError("name 必填"), nil
}
raw, err := GetOpenAPIClient().Post("/tasks", body)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("update_task",
mcp.WithDescription("更新任务字段(只传需要改的字段)"),
mcp.WithString("id", mcp.Required(), mcp.Description("任务 ID")),
mcp.WithString("name", mcp.Description("任务名称")),
mcp.WithString("command", mcp.Description("命令")),
mcp.WithString("schedule", mcp.Description("Cron")),
mcp.WithString("remark", mcp.Description("备注")),
mcp.WithString("tags", mcp.Description("标签")),
mcp.WithString("work_dir", mcp.Description("工作目录")),
mcp.WithString("enabled", mcp.Description("true/false")),
mcp.WithNumber("timeout", mcp.Description("超时秒数")),
mcp.WithString("envs", mcp.Description("环境变量 JSON 字符串")),
mcp.WithString("type", mcp.Description("任务类型")),
mcp.WithString("config", mcp.Description("配置 JSON 字符串")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
body := map[string]any{}
setStr(body, "name", strArg(req, "name"))
setStr(body, "command", strArg(req, "command"))
setStr(body, "schedule", strArg(req, "schedule"))
setStr(body, "remark", strArg(req, "remark"))
setStr(body, "tags", strArg(req, "tags"))
setStr(body, "work_dir", strArg(req, "work_dir"))
setStr(body, "envs", strArg(req, "envs"))
setStr(body, "type", strArg(req, "type"))
setStr(body, "config", strArg(req, "config"))
if e := strArg(req, "enabled"); e != "" {
body["enabled"] = strings.EqualFold(e, "true") || e == "1"
}
if n, ok := numArg(req, "timeout"); ok {
body["timeout"] = n
}
raw, err := GetOpenAPIClient().Put("/tasks/"+id, body)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("delete_task",
mcp.WithDescription("删除任务(危险操作,执行前请确认)"),
mcp.WithString("id", mcp.Required(), mcp.Description("任务 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Delete("/tasks/" + id)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("list_task_tags",
mcp.WithDescription("获取所有任务标签"),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
raw, err := GetOpenAPIClient().Get("/tasks/tags", nil)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("run_task",
mcp.WithDescription("立即执行指定任务。可选传入临时环境变量 envs_json"),
mcp.WithString("id", mcp.Required(), mcp.Description("任务 ID")),
mcp.WithString("envs_json", mcp.Description("临时环境变量 JSON 对象字符串,如 {\"FOO\":\"bar\"}")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
var body any
if ej := strArg(req, "envs_json"); ej != "" {
var envs map[string]string
if err := json.Unmarshal([]byte(ej), &envs); err != nil {
return mcp.NewToolResultError("envs_json 不是合法 JSON 对象: " + err.Error()), nil
}
body = map[string]any{"envs": envs}
} else {
body = map[string]any{}
}
raw, err := GetOpenAPIClient().Post("/execute/task/"+id, body)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("stop_task",
mcp.WithDescription("停止正在运行的任务(需要执行日志 log_id)"),
mcp.WithString("log_id", mcp.Required(), mcp.Description("运行日志 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
logID, err := req.RequireString("log_id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Post("/tasks/stop/"+logID, map[string]any{})
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("get_last_results",
mcp.WithDescription("获取最近若干次执行结果摘要"),
mcp.WithNumber("count", mcp.Description("数量,默认 10")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
q := map[string]string{"count": numArgStr(req, "count", "10")}
raw, err := GetOpenAPIClient().Get("/execute/results", q)
return toolResult(raw, err)
})
// ---- logs ----
s.AddTool(mcp.NewTool("list_logs",
mcp.WithDescription("分页查询任务执行日志"),
mcp.WithString("task_id", mcp.Description("任务 ID")),
mcp.WithString("task_name", mcp.Description("任务名称模糊")),
mcp.WithString("status", mcp.Description("状态筛选")),
mcp.WithNumber("page", mcp.Description("页码,默认 1")),
mcp.WithNumber("page_size", mcp.Description("每页数量,默认 20")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
q := map[string]string{
"task_id": strArg(req, "task_id"),
"task_name": strArg(req, "task_name"),
"status": strArg(req, "status"),
"page": numArgStr(req, "page", "1"),
"page_size": numArgStr(req, "page_size", "20"),
}
raw, err := GetOpenAPIClient().Get("/logs", q)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("get_log",
mcp.WithDescription("获取单条执行日志详情(含输出)"),
mcp.WithString("id", mcp.Required(), mcp.Description("日志 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Get("/logs/"+id, nil)
return toolResult(raw, err)
})
// ---- scripts ----
s.AddTool(mcp.NewTool("list_scripts",
mcp.WithDescription("列出脚本(不含内容)"),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
raw, err := GetOpenAPIClient().Get("/scripts", nil)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("get_script",
mcp.WithDescription("获取脚本详情(含内容)"),
mcp.WithString("id", mcp.Required(), mcp.Description("脚本 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Get("/scripts/"+id, nil)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("create_script",
mcp.WithDescription("创建脚本"),
mcp.WithString("name", mcp.Required(), mcp.Description("脚本名称/路径")),
mcp.WithString("content", mcp.Required(), mcp.Description("脚本内容")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
name, err := req.RequireString("name")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
content, err := req.RequireString("content")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Post("/scripts", map[string]any{"name": name, "content": content})
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("update_script",
mcp.WithDescription("更新脚本"),
mcp.WithString("id", mcp.Required(), mcp.Description("脚本 ID")),
mcp.WithString("name", mcp.Description("名称")),
mcp.WithString("content", mcp.Description("内容")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
body := map[string]any{}
setStr(body, "name", strArg(req, "name"))
setStr(body, "content", strArg(req, "content"))
raw, err := GetOpenAPIClient().Put("/scripts/"+id, body)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("delete_script",
mcp.WithDescription("删除脚本(危险操作)"),
mcp.WithString("id", mcp.Required(), mcp.Description("脚本 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Delete("/scripts/" + id)
return toolResult(raw, err)
})
// ---- env ----
s.AddTool(mcp.NewTool("list_envs",
mcp.WithDescription("分页列出环境变量(secret 值已脱敏)"),
mcp.WithString("name", mcp.Description("按名称筛选")),
mcp.WithString("type", mcp.Description("类型筛选")),
mcp.WithString("tags", mcp.Description("标签筛选")),
mcp.WithNumber("page", mcp.Description("页码")),
mcp.WithNumber("page_size", mcp.Description("每页数量")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
q := map[string]string{
"name": strArg(req, "name"),
"type": strArg(req, "type"),
"tags": strArg(req, "tags"),
"page": numArgStr(req, "page", "1"),
"page_size": numArgStr(req, "page_size", "20"),
}
raw, err := GetOpenAPIClient().Get("/env", q)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("list_all_envs",
mcp.WithDescription("获取全部环境变量(不分页,secret 脱敏)"),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
raw, err := GetOpenAPIClient().Get("/env/all", nil)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("get_env",
mcp.WithDescription("获取环境变量详情"),
mcp.WithString("id", mcp.Required(), mcp.Description("环境变量 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Get("/env/"+id, nil)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("create_env",
mcp.WithDescription("创建环境变量"),
mcp.WithString("name", mcp.Required(), mcp.Description("变量名")),
mcp.WithString("value", mcp.Required(), mcp.Description("变量值")),
mcp.WithString("type", mcp.Description("类型,如 text/secret")),
mcp.WithString("remark", mcp.Description("备注")),
mcp.WithString("tags", mcp.Description("标签")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
name, err := req.RequireString("name")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
value, err := req.RequireString("value")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
body := map[string]any{"name": name, "value": value}
setStr(body, "type", strArg(req, "type"))
setStr(body, "remark", strArg(req, "remark"))
setStr(body, "tags", strArg(req, "tags"))
raw, err := GetOpenAPIClient().Post("/env", body)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("update_env",
mcp.WithDescription("更新环境变量"),
mcp.WithString("id", mcp.Required(), mcp.Description("环境变量 ID")),
mcp.WithString("name", mcp.Description("变量名")),
mcp.WithString("value", mcp.Description("变量值")),
mcp.WithString("type", mcp.Description("类型")),
mcp.WithString("remark", mcp.Description("备注")),
mcp.WithString("tags", mcp.Description("标签")),
mcp.WithString("enabled", mcp.Description("true/false")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
body := map[string]any{}
setStr(body, "name", strArg(req, "name"))
setStr(body, "value", strArg(req, "value"))
setStr(body, "type", strArg(req, "type"))
setStr(body, "remark", strArg(req, "remark"))
setStr(body, "tags", strArg(req, "tags"))
if e := strArg(req, "enabled"); e != "" {
body["enabled"] = strings.EqualFold(e, "true") || e == "1"
}
raw, err := GetOpenAPIClient().Put("/env/"+id, body)
return toolResult(raw, err)
})
s.AddTool(mcp.NewTool("delete_env",
mcp.WithDescription("删除环境变量(危险操作)"),
mcp.WithString("id", mcp.Required(), mcp.Description("环境变量 ID")),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
id, err := req.RequireString("id")
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
raw, err := GetOpenAPIClient().Delete("/env/" + id)
return toolResult(raw, err)
})
// ---- meta ----
s.AddTool(mcp.NewTool("ping_api",
mcp.WithDescription("检查 OpenAPI 连通性与鉴权是否有效(通过 list_tasks page_size=1"),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
raw, err := GetOpenAPIClient().Get("/tasks", map[string]string{"page": "1", "page_size": "1"})
if err != nil {
return mcp.NewToolResultError("连通失败: " + err.Error()), nil
}
return mcp.NewToolResultText("OpenAPI 连通正常\n" + prettyRaw(raw)), nil
})
}
func registerPrompts(s *server.MCPServer) {
s.AddPrompt(mcp.NewPrompt("ops_overview",
mcp.WithPromptDescription("巡检任务池:列出任务与最近失败日志"),
), func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
return mcp.NewGetPromptResult(
"TaskPool 巡检",
[]mcp.PromptMessage{
mcp.NewPromptMessage(mcp.RoleUser, mcp.NewTextContent(
"请使用 TaskPool MCP1) list_tasks 查看任务概况 2) list_logs 筛选失败状态 3) 对重要失败 get_log 总结原因并给出是否建议 run_task 重跑。",
)),
},
), nil
})
s.AddPrompt(mcp.NewPrompt("run_and_check",
mcp.WithPromptDescription("执行指定任务并检查日志"),
mcp.WithArgument("task_id", mcp.RequiredArgument(), mcp.ArgumentDescription("任务 ID")),
), func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
taskID := ""
if request.Params.Arguments != nil {
taskID = request.Params.Arguments["task_id"]
}
text := fmt.Sprintf("请对 TaskPool 任务 %s1) get_task 确认 2) run_task 执行 3) list_logs(task_id) 取最新日志 4) get_log 汇总输出与状态。", taskID)
return mcp.NewGetPromptResult(
"执行并检查",
[]mcp.PromptMessage{
mcp.NewPromptMessage(mcp.RoleUser, mcp.NewTextContent(text)),
},
), nil
})
}
func toolResult(raw json.RawMessage, err error) (*mcp.CallToolResult, error) {
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText(prettyRaw(raw)), nil
}
func strArg(req mcp.CallToolRequest, key string) string {
v, err := req.RequireString(key)
if err != nil {
if args, ok := req.Params.Arguments.(map[string]any); ok {
if x, ok := args[key]; ok {
return strings.TrimSpace(fmt.Sprint(x))
}
}
return ""
}
return strings.TrimSpace(v)
}
func numArg(req mcp.CallToolRequest, key string) (float64, bool) {
if args, ok := req.Params.Arguments.(map[string]any); ok {
if x, ok := args[key]; ok {
switch t := x.(type) {
case float64:
return t, true
case int:
return float64(t), true
case string:
if t == "" {
return 0, false
}
f, err := strconv.ParseFloat(t, 64)
if err != nil {
return 0, false
}
return f, true
}
}
}
return 0, false
}
func numArgStr(req mcp.CallToolRequest, key, def string) string {
if n, ok := numArg(req, key); ok {
return strconv.FormatInt(int64(n), 10)
}
if s := strArg(req, key); s != "" {
return s
}
return def
}
func setStr(m map[string]any, k, v string) {
if strings.TrimSpace(v) != "" {
m[k] = v
}
}
func prettyJSON(v any) string {
b, err := json.MarshalIndent(v, "", " ")
if err != nil {
return fmt.Sprintf("%v", v)
}
return string(b)
}
func prettyRaw(raw json.RawMessage) string {
if len(raw) == 0 {
return "null"
}
var v any
if err := json.Unmarshal(raw, &v); err != nil {
return string(raw)
}
return prettyJSON(v)
}