From f3a44a901ec34e54cddd8afc31aa9806ed888f28 Mon Sep 17 00:00:00 2001 From: TaskPool Date: Mon, 27 Jul 2026 02:43:01 +0800 Subject: [PATCH] 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 --- cmd/mcp/client.go | 157 ---------- cmd/mcp/server.go | 686 +++++------------------------------------ internal/mcp/client.go | 35 ++- internal/mcp/server.go | 25 +- 4 files changed, 130 insertions(+), 773 deletions(-) delete mode 100644 cmd/mcp/client.go diff --git a/cmd/mcp/client.go b/cmd/mcp/client.go deleted file mode 100644 index 9c1842a..0000000 --- a/cmd/mcp/client.go +++ /dev/null @@ -1,157 +0,0 @@ -package mcp - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "net/http" - "net/url" - "strings" - "time" -) - -// Client 调用 TaskPool OpenAPI (/open2api/v1) -type Client struct { - BaseURL string - Token string - HTTPClient *http.Client -} - -func NewClient(baseURL, token string) *Client { - return &Client{ - BaseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"), - Token: strings.TrimSpace(token), - HTTPClient: &http.Client{ - Timeout: 60 * time.Second, - }, - } -} - -func (c *Client) openBase() string { - return c.BaseURL + "/open2api/v1" -} - -func (c *Client) do(method, path string, query map[string]string, body any) (json.RawMessage, error) { - if c.BaseURL == "" { - return nil, fmt.Errorf("TASKPOOL_URL 未配置") - } - if c.Token == "" { - return nil, fmt.Errorf("TASKPOOL_TOKEN 未配置") - } - - u, err := url.Parse(c.openBase() + path) - if err != nil { - return nil, err - } - if len(query) > 0 { - q := u.Query() - for k, v := range query { - if strings.TrimSpace(v) != "" { - q.Set(k, v) - } - } - u.RawQuery = q.Encode() - } - - var reader io.Reader - if body != nil { - b, err := json.Marshal(body) - if err != nil { - return nil, err - } - reader = bytes.NewReader(b) - } - - req, err := http.NewRequest(method, u.String(), reader) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+c.Token) - req.Header.Set("Accept", "application/json") - if body != nil { - req.Header.Set("Content-Type", "application/json") - } - - resp, err := c.HTTPClient.Do(req) - if err != nil { - return nil, fmt.Errorf("请求失败: %w", err) - } - defer resp.Body.Close() - - raw, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - var envelope struct { - Code int `json:"code"` - Msg string `json:"msg"` - Data json.RawMessage `json:"data"` - } - if err := json.Unmarshal(raw, &envelope); err != nil { - if resp.StatusCode >= 400 { - return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(raw), 500)) - } - return raw, nil - } - - // TaskPool 统一返回 HTTP 200，用 JSON body 中 code 区分成功/失败 - // code=200 → 成功；其他 → 错误 - if envelope.Code != 200 { - msg := envelope.Msg - if msg == "" { - msg = truncate(string(raw), 500) - } - return nil, fmt.Errorf("[%d] %s", envelope.Code, msg) - } - if len(envelope.Data) > 0 && string(envelope.Data) != "null" { - return envelope.Data, nil - } - // data 为空但请求成功（如 delete/stop 只返回 msg） - result := map[string]any{"code": envelope.Code, "msg": envelope.Msg} - b, _ := json.Marshal(result) - return b, nil -} - -func (c *Client) Get(path string, query map[string]string) (json.RawMessage, error) { - return c.do(http.MethodGet, path, query, nil) -} - -func (c *Client) Post(path string, body any) (json.RawMessage, error) { - return c.do(http.MethodPost, path, nil, body) -} - -func (c *Client) Put(path string, body any) (json.RawMessage, error) { - return c.do(http.MethodPut, path, nil, body) -} - -func (c *Client) Delete(path string) (json.RawMessage, error) { - return c.do(http.MethodDelete, path, nil, nil) -} - -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) -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] + "..." -} diff --git a/cmd/mcp/server.go b/cmd/mcp/server.go index c9bf154..0e66a9d 100644 --- a/cmd/mcp/server.go +++ b/cmd/mcp/server.go @@ -1,21 +1,13 @@ package mcp import ( - "context" - "encoding/json" "fmt" "net/http" "os" - "strconv" "strings" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" -) - -const ( - serverName = "taskpool" - serverVersion = "1.0.0" + "github.com/engigu/taskpool/internal/mcp" + imc "github.com/engigu/taskpool/internal/mcp" ) // Run 启动 TaskPool MCP Server @@ -28,14 +20,17 @@ const ( // TASKPOOL_TOKEN 设置页 OpenAPI Token // MCP_HTTP_ADDR HTTP 模式监听地址（默认 :8053） func Run(args []string) { - mode := "stdio" - httpAddr := ":8053" - - for i, a := range args { + for _, a := range args { if a == "-h" || a == "--help" { printHelp() return } + } + + mode := "stdio" + httpAddr := ":8053" + + for i, a := range args { if a == "--http" { mode = "http" if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") { @@ -44,627 +39,60 @@ func Run(args []string) { } } - // HTTP 模式优先从环境变量读取地址 if envAddr := os.Getenv("MCP_HTTP_ADDR"); envAddr != "" { httpAddr = envAddr } - 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) - if mode == "http" { - runHTTP(s, httpAddr, client) + runHTTP(httpAddr) } else { - runStdio(s, client) + runStdio() } } -func runStdio(s *server.MCPServer, client *Client) { - // 日志只能走 stderr，避免污染 stdio JSON-RPC +func runStdio() { + // 使用内部 OpenAPI 客户端（自动获取服务器地址和 Token） + client := imc.GetOpenAPIClient() fmt.Fprintf(os.Stderr, "[taskpool-mcp] mode=stdio url=%s token=%s\n", client.BaseURL, maskToken(client.Token)) - if err := server.ServeStdio(s); err != nil { + s := mcp.GetServer() + // 设置外部客户端（用于 stdio 模式连接远程服务器） + if client.Token != "" { + mcp.SetExternalClient(client) + } + + if err := mcp.ServeStdio(s); err != nil { fmt.Fprintf(os.Stderr, "[taskpool-mcp] server error: %v\n", err) os.Exit(1) } } -func runHTTP(s *server.MCPServer, addr string, client *Client) { +func runHTTP(addr string) { + client := imc.GetOpenAPIClient() fmt.Fprintf(os.Stderr, "[taskpool-mcp] mode=http addr=%s url=%s token=%s\n", addr, client.BaseURL, maskToken(client.Token)) - httpServer := server.NewStreamableHTTPServer(s, - server.WithEndpointPath("/mcp"), - ) + // 设置外部客户端 + if client.Token != "" { + mcp.SetExternalClient(client) + } + + handler := mcp.GetHTTPHandler() // 添加健康检查端点 - http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { + mux := http.NewServeMux() + mux.Handle("/mcp/", http.StripPrefix("/mcp", handler)) + mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) w.Write([]byte(`{"status":"ok"}`)) }) fmt.Fprintf(os.Stderr, "[taskpool-mcp] MCP endpoint: http://%s/mcp\n", strings.TrimPrefix(addr, ":")) - if err := httpServer.Start(addr); err != nil { + if err := http.ListenAndServe(addr, mux); err != nil { fmt.Fprintf(os.Stderr, "[taskpool-mcp] http 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 # stdio 模式（默认），适用于本地 Agent - taskpool mcp --http # HTTP 模式，监听 :8053 - taskpool mcp --http :9000 # HTTP 模式，监听指定端口 - -环境变量: - TASKPOOL_URL 面板地址 (默认 http://127.0.0.1:8052) - TASKPOOL_TOKEN 设置页生成的 OpenAPI Token（必填） - MCP_HTTP_ADDR HTTP 模式监听地址（默认 :8053） - -模式说明: - stdio 通过 stdin/stdout 通信，需要本地安装 taskpool 二进制 - http 通过 HTTP/SSE 通信，支持远程访问，无需本地安装 - ---- stdio 模式配置示例 --- -Hermes / OpenClaw / Cursor 配置: - { - "mcpServers": { - "taskpool": { - "command": "taskpool", - "args": ["mcp"], - "env": { - "TASKPOOL_URL": "http://127.0.0.1:8052", - "TASKPOOL_TOKEN": "你的 OpenAPI Token" - } - } - } - } - ---- HTTP 模式配置示例 --- -1. 启动 MCP HTTP Server: - taskpool mcp --http :8053 - -2. Hermes / OpenClaw 配置（使用 URL）: - { - "mcpServers": { - "taskpool": { - "url": "http://your-server:8053/mcp" - } - } - } - -注意: HTTP 模式建议配置反向代理（如 Nginx）添加 HTTPS 和认证 - -前置条件: - 1. 面板「系统设置」启用 OpenAPI 并生成 Token - 2. 后端服务已启动 (taskpool server) -`) -} - -const instructions = `你是 TaskPool（任务池）运维助手。通过 MCP tools 管理任务调度面板。 - -核心能力： -- 任务：列表/详情/创建/更新/删除/启停执行 -- 脚本：列表/详情/创建/更新/删除 -- 环境变量：列表/详情/创建/更新/删除（secret 值为脱敏） -- 日志：列表/详情、最近执行结果 - -使用原则： -1. 先 list_tasks / get_task 确认目标，再 run_task -2. 删除类操作前向用户确认 -3. 不要编造任务 ID；ID 以 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 MCP：1) 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 任务 %s：1) 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 == "" { @@ -675,3 +103,51 @@ func maskToken(t string) string { } return t[:4] + "****" + t[len(t)-4:] } + +func printHelp() { + fmt.Fprintf(os.Stderr, ` +TaskPool MCP Server（内置模式） + +MCP Server 已集成到后端服务中，推荐直接使用内置端点： + http://your-server:8052/mcp + +此命令用于： + - stdio 模式：本地 Agent（Claude Desktop / Cursor） + - 独立 HTTP 服务：需要单独端口时 + +用法: + taskpool mcp # stdio 模式 + taskpool mcp --http # HTTP 模式，监听 :8053 + taskpool mcp --http :9000 # HTTP 模式，指定端口 + +环境变量（stdio 模式可选）: + TASKPOOL_URL 面板地址（默认本机） + TASKPOOL_TOKEN OpenAPI Token（默认从系统设置读取） + +Agent 配置示例: + HTTP 模式（推荐）: + { + "mcpServers": { + "taskpool": { + "url": "http://your-server:8052/mcp", + "headers": { + "Authorization": "Bearer 你的Token" + } + } + } + } + + stdio 模式: + { + "mcpServers": { + "taskpool": { + "command": "taskpool", + "args": ["mcp"] + } + } + } + +前置条件: + 系统设置 → 启用 OpenAPI → 生成 Token +`) +} diff --git a/internal/mcp/client.go b/internal/mcp/client.go index c03f63c..f2b83d6 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "net/url" + "os" "strings" "sync" "time" @@ -21,18 +22,36 @@ var ( openAPIClientOnce sync.Once ) -// GetOpenAPIClient 获取内部 OpenAPI 客户端（单例） -// 使用当前服务器的地址和系统设置中的 OpenAPI Token +// GetOpenAPIClient 获取 OpenAPI 客户端（单例） +// 优先使用环境变量 TASKPOOL_URL 和 TASKPOOL_TOKEN +// 否则使用本地服务器地址和系统设置中的 OpenAPI Token func GetOpenAPIClient() *Client { openAPIClientOnce.Do(func() { - cfg := services.GetConfig() - baseURL := fmt.Sprintf("http://127.0.0.1:%d", cfg.Server.Port) - if cfg.Server.URLPrefix != "" { - baseURL += cfg.Server.URLPrefix + // 优先从环境变量读取（支持 stdio 模式连接远程服务器） + envURL := strings.TrimSpace(os.Getenv("TASKPOOL_URL")) + envToken := strings.TrimSpace(os.Getenv("TASKPOOL_TOKEN")) + if envURL == "" { + envURL = strings.TrimSpace(os.Getenv("BAIHU_URL")) + } + if envToken == "" { + envToken = strings.TrimSpace(os.Getenv("BAIHU_TOKEN")) } - settingsSvc := services.NewSettingsService() - token := getOpenAPIToken(settingsSvc) + var baseURL, token string + if envURL != "" && envToken != "" { + // 使用环境变量配置 + baseURL = envURL + token = envToken + } else { + // 使用本地服务器配置 + cfg := services.GetConfig() + baseURL = fmt.Sprintf("http://127.0.0.1:%d", cfg.Server.Port) + if cfg.Server.URLPrefix != "" { + baseURL += cfg.Server.URLPrefix + } + settingsSvc := services.NewSettingsService() + token = getOpenAPIToken(settingsSvc) + } openAPIClient = NewClient(baseURL, token) }) diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 52ac1b4..f215150 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -19,9 +19,10 @@ const ( ) var ( - mcpServer *server.MCPServer - mcpHTTPServer *server.StreamableHTTPServer - once sync.Once + mcpServer *server.MCPServer + mcpHTTPServer *server.StreamableHTTPServer + externalClient *Client // 外部客户端（用于 stdio 模式连接远程服务器） + once sync.Once ) // GetServer 获取 MCP Server 实例（单例） @@ -42,6 +43,19 @@ func GetServer() *server.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() @@ -53,6 +67,11 @@ func GetHTTPHandler() http.Handler { return mcpHTTPServer } +// ServeStdio 启动 stdio 模式的 MCP Server +func ServeStdio(s *server.MCPServer) error { + return server.ServeStdio(s) +} + const instructions = `你是 TaskPool（任务池）运维助手。通过 MCP tools 管理任务调度面板。 核心能力：