|
|
@@ -0,0 +1,372 @@
|
|
|
+package kit
|
|
|
+
|
|
|
+// 持久化长任务（MCP Tasks，规范 2025-11-25；AgentServer 接入指南 §5）。
|
|
|
+//
|
|
|
+// mcp-go 的任务只保存在进程内存里，fwh-agent-mcp 重启后平台查询会得到“任务不存在”而按失败通知用户，
|
|
|
+// 而底层作业（如主服务的站点信息同步）其实还在执行。所以长任务工具不使用 mcp-go 的任务表：
|
|
|
+// - 任务状态由业务数据给出（如 Mongo 中的同步记录），MCP 进程不保存任何任务状态，重启后照常回答查询；
|
|
|
+// - 任务 id 携带工具名、业务键与标记，并带 HMAC 签名，不可伪造、不可猜测；
|
|
|
+// - TaskRouter 在 mcp-go 之前处理“以任务方式调用工具”与 tasks/get、tasks/result、tasks/cancel，其余请求原样交给 mcp-go。
|
|
|
+// 普通调用（不带 task 参数）照常由 mcp-go 按 Def.Handler 同步返回。
|
|
|
+
|
|
|
+import (
|
|
|
+ "bytes"
|
|
|
+ "context"
|
|
|
+ "crypto/hmac"
|
|
|
+ "crypto/sha256"
|
|
|
+ "encoding/hex"
|
|
|
+ "encoding/json"
|
|
|
+ "errors"
|
|
|
+ "fmt"
|
|
|
+ "io"
|
|
|
+ "net/http"
|
|
|
+ "strings"
|
|
|
+ "time"
|
|
|
+
|
|
|
+ "github.com/mark3labs/mcp-go/mcp"
|
|
|
+ "github.com/mark3labs/mcp-go/server"
|
|
|
+ "go.uber.org/zap"
|
|
|
+)
|
|
|
+
|
|
|
+const (
|
|
|
+ taskIDPrefix = "fwhtask"
|
|
|
+ immediateMeta = "io.modelcontextprotocol/model-immediate-response"
|
|
|
+ relatedMeta = "io.modelcontextprotocol/related-task"
|
|
|
+ // defaultTaskTTL 调用方未指定 ttl 时声明的保留时间（毫秒）。业务数据本身长期保存，这里只是告知调用方
|
|
|
+ defaultTaskTTL = int64(24 * time.Hour / time.Millisecond)
|
|
|
+ // resultWaitStep tasks/result 在任务未结束时的查询间隔（规范要求阻塞到任务结束）
|
|
|
+ resultWaitStep = 5 * time.Second
|
|
|
+ maxRPCBody = 1 << 20
|
|
|
+)
|
|
|
+
|
|
|
+// ErrTaskNotFound 业务键对应的数据不存在（返回给调用方“找不到任务”）
|
|
|
+var ErrTaskNotFound = errors.New("task not found")
|
|
|
+
|
|
|
+// TaskRef 任务 id 中携带的信息
|
|
|
+type TaskRef struct {
|
|
|
+ Tool string // 工具名
|
|
|
+ Key string // 业务键，如同步记录 id（只含字母数字）
|
|
|
+ Flag string // 工具自定义的标记，如本次调用是否触发了作业（只含小写字母、数字、下划线）
|
|
|
+}
|
|
|
+
|
|
|
+// TaskStart 以任务方式调用时 Start 的返回
|
|
|
+type TaskStart struct {
|
|
|
+ Key string
|
|
|
+ Flag string
|
|
|
+ Note string // 立即交给模型的说明（model-immediate-response），如“已开始同步，预计约 1 小时”
|
|
|
+}
|
|
|
+
|
|
|
+// TaskState 由业务数据计算出的任务状态
|
|
|
+type TaskState struct {
|
|
|
+ Status mcp.TaskStatus
|
|
|
+ Message string // statusMessage，显示在平台的后台调用栏
|
|
|
+ Result *mcp.CallToolResult // 终态时的结果；failed 时应为 isError
|
|
|
+ CreatedAt time.Time
|
|
|
+}
|
|
|
+
|
|
|
+// TaskTool 持久化长任务工具
|
|
|
+type TaskTool interface {
|
|
|
+ Registrar
|
|
|
+ PollInterval() time.Duration
|
|
|
+ // StartTask 立即返回（几秒内）：启动或挂到已有作业上。参数或业务错误时返回 reject（作为错误返回给调用方）
|
|
|
+ StartTask(ctx context.Context, args json.RawMessage) (start TaskStart, reject *mcp.CallToolResult)
|
|
|
+ // TaskStatus 按业务数据计算任务状态；数据不存在时返回 ErrTaskNotFound
|
|
|
+ TaskStatus(ctx context.Context, ref TaskRef) (TaskState, error)
|
|
|
+ // CancelTask 取消任务；由工具按 ref.Flag 决定是否真的停止作业（只是跟踪他人触发的作业时不停止）
|
|
|
+ CancelTask(ctx context.Context, ref TaskRef) error
|
|
|
+}
|
|
|
+
|
|
|
+// TaskDef 类型化的持久化长任务工具：Def 用于普通调用，其余回调用于任务方式调用
|
|
|
+type TaskDef[T any] struct {
|
|
|
+ Def[T]
|
|
|
+ Poll time.Duration
|
|
|
+ Start func(ctx context.Context, args T) (TaskStart, *mcp.CallToolResult)
|
|
|
+ Status func(ctx context.Context, ref TaskRef) (TaskState, error)
|
|
|
+ Cancel func(ctx context.Context, ref TaskRef) error
|
|
|
+}
|
|
|
+
|
|
|
+func (d TaskDef[T]) Tool() mcp.Tool {
|
|
|
+ t := d.Def.Tool()
|
|
|
+ mcp.WithTaskSupport(mcp.TaskSupportOptional)(&t)
|
|
|
+ return t
|
|
|
+}
|
|
|
+
|
|
|
+func (d TaskDef[T]) Register(s *server.MCPServer) {
|
|
|
+ s.AddTool(d.Tool(), mcp.NewTypedToolHandler(func(ctx context.Context, _ mcp.CallToolRequest, args T) (*mcp.CallToolResult, error) {
|
|
|
+ return d.Handler(ctx, args)
|
|
|
+ }))
|
|
|
+}
|
|
|
+
|
|
|
+func (d TaskDef[T]) PollInterval() time.Duration { return d.Poll }
|
|
|
+
|
|
|
+func (d TaskDef[T]) StartTask(ctx context.Context, raw json.RawMessage) (TaskStart, *mcp.CallToolResult) {
|
|
|
+ var args T
|
|
|
+ if len(raw) > 0 && string(raw) != "null" {
|
|
|
+ if err := json.Unmarshal(raw, &args); err != nil {
|
|
|
+ return TaskStart{}, Failf("参数格式不正确：%v", err)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return d.Start(ctx, args)
|
|
|
+}
|
|
|
+
|
|
|
+func (d TaskDef[T]) TaskStatus(ctx context.Context, ref TaskRef) (TaskState, error) {
|
|
|
+ return d.Status(ctx, ref)
|
|
|
+}
|
|
|
+
|
|
|
+func (d TaskDef[T]) CancelTask(ctx context.Context, ref TaskRef) error { return d.Cancel(ctx, ref) }
|
|
|
+
|
|
|
+// ---- 任务 id ----
|
|
|
+
|
|
|
+type taskIDCodec struct{ secret []byte }
|
|
|
+
|
|
|
+func (c taskIDCodec) encode(ref TaskRef) string {
|
|
|
+ body := strings.Join([]string{taskIDPrefix, ref.Tool, ref.Key, ref.Flag, errorRef()}, ".")
|
|
|
+ return body + "." + c.sign(body)
|
|
|
+}
|
|
|
+
|
|
|
+// decode 返回 ok=false 表示不是本服务签发的任务 id（交给 mcp-go）；err 表示前缀匹配但校验失败
|
|
|
+func (c taskIDCodec) decode(id string) (ref TaskRef, ok bool, err error) {
|
|
|
+ parts := strings.Split(id, ".")
|
|
|
+ if len(parts) == 0 || parts[0] != taskIDPrefix {
|
|
|
+ return TaskRef{}, false, nil
|
|
|
+ }
|
|
|
+ if len(parts) != 6 || !hmac.Equal([]byte(c.sign(strings.Join(parts[:5], "."))), []byte(parts[5])) {
|
|
|
+ return TaskRef{}, true, ErrTaskNotFound
|
|
|
+ }
|
|
|
+ return TaskRef{Tool: parts[1], Key: parts[2], Flag: parts[3]}, true, nil
|
|
|
+}
|
|
|
+
|
|
|
+func (c taskIDCodec) sign(body string) string {
|
|
|
+ m := hmac.New(sha256.New, c.secret)
|
|
|
+ m.Write([]byte(body))
|
|
|
+ return hex.EncodeToString(m.Sum(nil))[:32]
|
|
|
+}
|
|
|
+
|
|
|
+// ---- 协议处理 ----
|
|
|
+
|
|
|
+// TaskRouter 在 mcp-go 之前处理持久化长任务的协议请求
|
|
|
+type TaskRouter struct {
|
|
|
+ next http.Handler
|
|
|
+ tools map[string]TaskTool
|
|
|
+ ids taskIDCodec
|
|
|
+ log *zap.Logger
|
|
|
+}
|
|
|
+
|
|
|
+// NewTaskRouter secret 用于任务 id 签名（使用 MCP 访问令牌即可，只有持有令牌的调用方能访问本服务）
|
|
|
+func NewTaskRouter(next http.Handler, tools []TaskTool, secret string, log *zap.Logger) *TaskRouter {
|
|
|
+ m := make(map[string]TaskTool, len(tools))
|
|
|
+ for _, t := range tools {
|
|
|
+ m[t.Meta().Name] = t
|
|
|
+ }
|
|
|
+ return &TaskRouter{next: next, tools: m, ids: taskIDCodec{secret: []byte(secret)}, log: log}
|
|
|
+}
|
|
|
+
|
|
|
+type rpcRequest struct {
|
|
|
+ JSONRPC string `json:"jsonrpc"`
|
|
|
+ ID json.RawMessage `json:"id"`
|
|
|
+ Method string `json:"method"`
|
|
|
+ Params json.RawMessage `json:"params"`
|
|
|
+}
|
|
|
+
|
|
|
+func (r *TaskRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
|
|
+ if req.Method != http.MethodPost || len(r.tools) == 0 {
|
|
|
+ r.next.ServeHTTP(w, req)
|
|
|
+ return
|
|
|
+ }
|
|
|
+ body, err := io.ReadAll(io.LimitReader(req.Body, maxRPCBody))
|
|
|
+ if err != nil {
|
|
|
+ http.Error(w, "read body failed", http.StatusBadRequest)
|
|
|
+ return
|
|
|
+ }
|
|
|
+ req.Body = io.NopCloser(bytes.NewReader(body))
|
|
|
+
|
|
|
+ var rpc rpcRequest
|
|
|
+ if json.Unmarshal(body, &rpc) != nil || len(rpc.ID) == 0 || !r.handle(w, req.Context(), rpc) {
|
|
|
+ r.next.ServeHTTP(w, req)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// handle 处理本服务负责的请求，返回 false 表示交给 mcp-go
|
|
|
+func (r *TaskRouter) handle(w http.ResponseWriter, ctx context.Context, rpc rpcRequest) bool {
|
|
|
+ switch rpc.Method {
|
|
|
+ case "tools/call":
|
|
|
+ var p struct {
|
|
|
+ Name string `json:"name"`
|
|
|
+ Arguments json.RawMessage `json:"arguments"`
|
|
|
+ Task *mcp.TaskParams `json:"task"`
|
|
|
+ }
|
|
|
+ if json.Unmarshal(rpc.Params, &p) != nil || p.Task == nil {
|
|
|
+ return false
|
|
|
+ }
|
|
|
+ tool, ok := r.tools[p.Name]
|
|
|
+ if !ok {
|
|
|
+ return false
|
|
|
+ }
|
|
|
+ r.create(w, ctx, rpc.ID, tool, p.Arguments, p.Task.TTL)
|
|
|
+ return true
|
|
|
+ case "tasks/get", "tasks/result", "tasks/cancel":
|
|
|
+ var p struct {
|
|
|
+ TaskID string `json:"taskId"`
|
|
|
+ }
|
|
|
+ if json.Unmarshal(rpc.Params, &p) != nil {
|
|
|
+ return false
|
|
|
+ }
|
|
|
+ ref, ours, err := r.ids.decode(p.TaskID)
|
|
|
+ if !ours {
|
|
|
+ return false
|
|
|
+ }
|
|
|
+ tool, known := r.tools[ref.Tool]
|
|
|
+ if err != nil || !known {
|
|
|
+ writeRPCError(w, rpc.ID, mcp.INVALID_PARAMS, "task not found")
|
|
|
+ return true
|
|
|
+ }
|
|
|
+ switch rpc.Method {
|
|
|
+ case "tasks/get":
|
|
|
+ r.get(w, ctx, rpc.ID, tool, ref, p.TaskID)
|
|
|
+ case "tasks/result":
|
|
|
+ r.result(w, ctx, rpc.ID, tool, ref, p.TaskID)
|
|
|
+ default:
|
|
|
+ r.cancel(w, ctx, rpc.ID, tool, ref, p.TaskID)
|
|
|
+ }
|
|
|
+ return true
|
|
|
+ }
|
|
|
+ return false
|
|
|
+}
|
|
|
+
|
|
|
+func (r *TaskRouter) create(w http.ResponseWriter, ctx context.Context, id json.RawMessage, tool TaskTool, args json.RawMessage, ttl *int64) {
|
|
|
+ start, reject := tool.StartTask(ctx, args)
|
|
|
+ if reject != nil {
|
|
|
+ writeRPCError(w, id, mcp.INVALID_PARAMS, resultText(reject))
|
|
|
+ return
|
|
|
+ }
|
|
|
+ ref := TaskRef{Tool: tool.Meta().Name, Key: start.Key, Flag: start.Flag}
|
|
|
+ taskID := r.ids.encode(ref)
|
|
|
+ state, err := tool.TaskStatus(ctx, ref)
|
|
|
+ if err != nil {
|
|
|
+ r.log.Warn("[agent-mcp] 任务创建后查询状态失败", zap.String("taskId", taskID), zap.Error(err))
|
|
|
+ writeRPCError(w, id, mcp.INTERNAL_ERROR, "任务已创建，但查询状态失败，请稍后重试")
|
|
|
+ return
|
|
|
+ }
|
|
|
+ if ttl == nil {
|
|
|
+ t := defaultTaskTTL
|
|
|
+ ttl = &t
|
|
|
+ }
|
|
|
+ res := map[string]any{"task": r.task(taskID, tool, state, ttl)}
|
|
|
+ if start.Note != "" {
|
|
|
+ res["_meta"] = map[string]any{immediateMeta: start.Note}
|
|
|
+ }
|
|
|
+ r.log.Info("[agent-mcp] 任务已创建", zap.String("tool", ref.Tool), zap.String("taskId", taskID), zap.String("flag", ref.Flag), zap.String("status", string(state.Status)))
|
|
|
+ writeRPCResult(w, id, res)
|
|
|
+}
|
|
|
+
|
|
|
+func (r *TaskRouter) get(w http.ResponseWriter, ctx context.Context, id json.RawMessage, tool TaskTool, ref TaskRef, taskID string) {
|
|
|
+ state, ok := r.status(w, ctx, id, tool, ref)
|
|
|
+ if ok {
|
|
|
+ writeRPCResult(w, id, r.task(taskID, tool, state, nil))
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// result 阻塞到任务结束（规范要求），然后返回工具结果
|
|
|
+func (r *TaskRouter) result(w http.ResponseWriter, ctx context.Context, id json.RawMessage, tool TaskTool, ref TaskRef, taskID string) {
|
|
|
+ for {
|
|
|
+ state, ok := r.status(w, ctx, id, tool, ref)
|
|
|
+ if !ok {
|
|
|
+ return
|
|
|
+ }
|
|
|
+ if state.Status.IsTerminal() {
|
|
|
+ res := state.Result
|
|
|
+ if res == nil {
|
|
|
+ res = mcp.NewToolResultText(state.Message)
|
|
|
+ res.IsError = state.Status != mcp.TaskStatusCompleted
|
|
|
+ }
|
|
|
+ out := *res
|
|
|
+ out.Meta = &mcp.Meta{AdditionalFields: map[string]any{relatedMeta: map[string]any{"taskId": taskID}}}
|
|
|
+ writeRPCResult(w, id, out)
|
|
|
+ return
|
|
|
+ }
|
|
|
+ select {
|
|
|
+ case <-ctx.Done():
|
|
|
+ writeRPCError(w, id, mcp.INTERNAL_ERROR, "等待任务结束时连接已断开")
|
|
|
+ return
|
|
|
+ case <-time.After(resultWaitStep):
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func (r *TaskRouter) cancel(w http.ResponseWriter, ctx context.Context, id json.RawMessage, tool TaskTool, ref TaskRef, taskID string) {
|
|
|
+ state, ok := r.status(w, ctx, id, tool, ref)
|
|
|
+ if !ok {
|
|
|
+ return
|
|
|
+ }
|
|
|
+ if state.Status.IsTerminal() {
|
|
|
+ writeRPCError(w, id, mcp.INVALID_PARAMS, fmt.Sprintf("cannot cancel task in terminal status: %s", state.Status))
|
|
|
+ return
|
|
|
+ }
|
|
|
+ if err := tool.CancelTask(ctx, ref); err != nil {
|
|
|
+ r.log.Warn("[agent-mcp] 取消任务失败", zap.String("taskId", taskID), zap.Error(err))
|
|
|
+ writeRPCError(w, id, mcp.INTERNAL_ERROR, "取消任务失败，请稍后重试")
|
|
|
+ return
|
|
|
+ }
|
|
|
+ r.log.Info("[agent-mcp] 任务已取消", zap.String("taskId", taskID), zap.String("flag", ref.Flag))
|
|
|
+ state.Status, state.Message = mcp.TaskStatusCancelled, "已取消"
|
|
|
+ writeRPCResult(w, id, r.task(taskID, tool, state, nil))
|
|
|
+}
|
|
|
+
|
|
|
+func (r *TaskRouter) status(w http.ResponseWriter, ctx context.Context, id json.RawMessage, tool TaskTool, ref TaskRef) (TaskState, bool) {
|
|
|
+ state, err := tool.TaskStatus(ctx, ref)
|
|
|
+ switch {
|
|
|
+ case errors.Is(err, ErrTaskNotFound):
|
|
|
+ writeRPCError(w, id, mcp.INVALID_PARAMS, "task not found")
|
|
|
+ return state, false
|
|
|
+ case err != nil:
|
|
|
+ // 暂时性错误：平台下次再查（连续失败 30 分钟才按失败处理）
|
|
|
+ r.log.Warn("[agent-mcp] 查询任务状态失败", zap.String("tool", ref.Tool), zap.String("key", ref.Key), zap.Error(err))
|
|
|
+ writeRPCError(w, id, mcp.INTERNAL_ERROR, "查询任务状态失败，请稍后重试")
|
|
|
+ return state, false
|
|
|
+ }
|
|
|
+ return state, true
|
|
|
+}
|
|
|
+
|
|
|
+func (r *TaskRouter) task(taskID string, tool TaskTool, s TaskState, ttl *int64) mcp.Task {
|
|
|
+ if ttl == nil {
|
|
|
+ t := defaultTaskTTL
|
|
|
+ ttl = &t
|
|
|
+ }
|
|
|
+ poll := int64(tool.PollInterval() / time.Millisecond)
|
|
|
+ created := s.CreatedAt
|
|
|
+ if created.IsZero() {
|
|
|
+ created = time.Now()
|
|
|
+ }
|
|
|
+ return mcp.Task{
|
|
|
+ TaskId: taskID,
|
|
|
+ Status: s.Status,
|
|
|
+ StatusMessage: s.Message,
|
|
|
+ CreatedAt: created.UTC().Format(time.RFC3339),
|
|
|
+ LastUpdatedAt: time.Now().UTC().Format(time.RFC3339),
|
|
|
+ TTL: ttl,
|
|
|
+ PollInterval: &poll,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+func resultText(r *mcp.CallToolResult) string {
|
|
|
+ var parts []string
|
|
|
+ for _, c := range r.Content {
|
|
|
+ if t, ok := c.(mcp.TextContent); ok {
|
|
|
+ parts = append(parts, t.Text)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return strings.Join(parts, "\n")
|
|
|
+}
|
|
|
+
|
|
|
+func writeRPCResult(w http.ResponseWriter, id json.RawMessage, result any) {
|
|
|
+ writeRPC(w, map[string]any{"jsonrpc": "2.0", "id": id, "result": result})
|
|
|
+}
|
|
|
+
|
|
|
+func writeRPCError(w http.ResponseWriter, id json.RawMessage, code int, msg string) {
|
|
|
+ writeRPC(w, map[string]any{"jsonrpc": "2.0", "id": id, "error": map[string]any{"code": code, "message": msg}})
|
|
|
+}
|
|
|
+
|
|
|
+func writeRPC(w http.ResponseWriter, v any) {
|
|
|
+ w.Header().Set("Content-Type", "application/json")
|
|
|
+ _ = json.NewEncoder(w).Encode(v)
|
|
|
+}
|
Jerry Liang