Files
TaskPool/cmd/mcp/server.go
T

607 lines
21 KiB
Go
Raw 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"
"os"
"strconv"
"strings"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
const (
serverName = "taskpool"
serverVersion = "1.0.0"
)
// Run 启动 TaskPool MCP Serverstdio
// 环境变量:
// TASKPOOL_URL 面板地址,如 http://127.0.0.1:8052 或 https://panel.example.com
// TASKPOOL_TOKEN 设置页 OpenAPI Token
func Run(args []string) {
for _, a := range args {
if a == "-h" || a == "--help" {
printHelp()
return
}
}
baseURL := firstNonEmpty(os.Getenv("TASKPOOL_URL"), os.Getenv("BAIHU_URL"), "http://127.0.0.1:8052")
token := firstNonEmpty(os.Getenv("TASKPOOL_TOKEN"), os.Getenv("BAIHU_TOKEN"), os.Getenv("OPENAPI_TOKEN"))
client := NewClient(baseURL, token)
s := server.NewMCPServer(
serverName,
serverVersion,
server.WithToolCapabilities(true),
server.WithResourceCapabilities(false, false),
server.WithPromptCapabilities(true),
server.WithInstructions(instructions),
server.WithRecovery(),
)
registerTools(s, client)
registerPrompts(s)
// 日志只能走 stderr,避免污染 stdio JSON-RPC
fmt.Fprintf(os.Stderr, "[taskpool-mcp] url=%s token=%s\n", client.BaseURL, maskToken(client.Token))
if err := server.ServeStdio(s); err != nil {
fmt.Fprintf(os.Stderr, "[taskpool-mcp] server error: %v\n", err)
os.Exit(1)
}
}
func printHelp() {
fmt.Fprintf(os.Stderr, `
TaskPool MCP Server
通过 MCP 协议把 TaskPool OpenAPI 暴露给 Hermes / OpenClaw / Cursor 等 Agent。
用法:
taskpool mcp
环境变量:
TASKPOOL_URL 面板地址 (默认 http://127.0.0.1:8052)
TASKPOOL_TOKEN 设置页生成的 OpenAPI Token(必填)
Hermes / OpenClaw 配置示例:
{
"mcpServers": {
"taskpool": {
"command": "taskpool",
"args": ["mcp"],
"env": {
"TASKPOOL_URL": "http://127.0.0.1:8052",
"TASKPOOL_TOKEN": "你的 OpenAPI Token"
}
}
}
}
前置条件:
1. 面板「系统设置」启用 OpenAPI 并生成 Token
2. 后端服务已启动 (taskpool server)
`)
}
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, c *Client) {
// ---- 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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 := c.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 {
// optional
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 firstNonEmpty(vals ...string) string {
for _, v := range vals {
if strings.TrimSpace(v) != "" {
return strings.TrimSpace(v)
}
}
return ""
}
func maskToken(t string) string {
t = strings.TrimSpace(t)
if t == "" {
return "(empty)"
}
if len(t) <= 8 {
return "****"
}
return t[:4] + "****" + t[len(t)-4:]
}