FranklinWH-Toolbox/service/agentmcp/kit/task.go

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)
}