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

package kit

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"go.uber.org/zap"
)

type fakeTaskArgs struct {
	Name string `json:"name"`
}

// 业务数据：key -> 状态；cancelled 记录被真正取消的 key
type fakeJobs struct {
	status    map[string]mcp.TaskStatus
	cancelled []string
}

func newFakeTaskTool(jobs *fakeJobs) TaskDef[fakeTaskArgs] {
	return TaskDef[fakeTaskArgs]{
		Def: Def[fakeTaskArgs]{Name: "job_run", Title: "运行作业", Description: "测试", Handler: func(context.Context, fakeTaskArgs) (*mcp.CallToolResult, error) {
			return mcp.NewToolResultText("sync"), nil
		}},
		Poll: time.Minute,
		Start: func(_ context.Context, a fakeTaskArgs) (TaskStart, *mcp.CallToolResult) {
			if a.Name == "" {
				return TaskStart{}, Failf("参数 name 必填")
			}
			jobs.status[a.Name] = mcp.TaskStatusWorking
			return TaskStart{Key: a.Name, Flag: "owner", Note: "已开始"}, nil
		},
		Status: func(_ context.Context, ref TaskRef) (TaskState, error) {
			st, ok := jobs.status[ref.Key]
			if !ok {
				return TaskState{}, ErrTaskNotFound
			}
			s := TaskState{Status: st, Message: string(st)}
			if st == mcp.TaskStatusCompleted {
				s.Result = mcp.NewToolResultText("done " + ref.Key)
			}
			if st == mcp.TaskStatusFailed {
				s.Result = Failf("failed %s", ref.Key)
			}
			return s, nil
		},
		Cancel: func(_ context.Context, ref TaskRef) error {
			if ref.Flag == "owner" {
				jobs.cancelled = append(jobs.cancelled, ref.Key)
			}
			return nil
		},
	}
}

type rpcReply struct {
	Result json.RawMessage `json:"result"`
	Error  *struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
	} `json:"error"`
}

func TestTaskRouter(t *testing.T) {
	jobs := &fakeJobs{status: map[string]mcp.TaskStatus{}}
	passed := 0
	next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		passed++
		_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`))
	})
	router := NewTaskRouter(next, []TaskTool{newFakeTaskTool(jobs)}, "secret", zap.NewNop())

	call := func(method string, params any) rpcReply {
		t.Helper()
		body, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 7, "method": method, "params": params})
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)))
		var r rpcReply
		if err := json.Unmarshal(rec.Body.Bytes(), &r); err != nil {
			t.Fatalf("%s: 响应不是 JSON-RPC：%s", method, rec.Body.String())
		}
		return r
	}

	// 以任务方式调用：立即返回任务与即时说明
	r := call("tools/call", map[string]any{"name": "job_run", "arguments": map[string]any{"name": "a"}, "task": map[string]any{"ttl": 1000}})
	var created struct {
		Task mcp.Task       `json:"task"`
		Meta map[string]any `json:"_meta"`
	}
	if r.Error != nil || json.Unmarshal(r.Result, &created) != nil {
		t.Fatalf("创建任务失败：%+v", r)
	}
	id := created.Task.TaskId
	if !strings.HasPrefix(id, taskIDPrefix+".job_run.a.owner.") || created.Task.Status != mcp.TaskStatusWorking ||
		*created.Task.TTL != 1000 || *created.Task.PollInterval != 60000 || created.Meta[immediateMeta] != "已开始" {
		t.Errorf("创建任务的返回不符合预期：%+v", created)
	}

	// 参数错误：返回 JSON-RPC 错误
	if r := call("tools/call", map[string]any{"name": "job_run", "arguments": map[string]any{}, "task": map[string]any{}}); r.Error == nil || !strings.Contains(r.Error.Message, "name 必填") {
		t.Errorf("参数错误应返回错误：%+v", r)
	}

	// 查询状态
	var got mcp.Task
	if r := call("tasks/get", map[string]any{"taskId": id}); r.Error != nil || json.Unmarshal(r.Result, &got) != nil || got.Status != mcp.TaskStatusWorking {
		t.Errorf("tasks/get 应为 working：%+v", r)
	}

	// 篡改任务 id：找不到任务
	tampered := strings.Replace(id, ".a.", ".b.", 1)
	if r := call("tasks/get", map[string]any{"taskId": tampered}); r.Error == nil || r.Error.Code != mcp.INVALID_PARAMS {
		t.Errorf("篡改的任务 id 应返回 INVALID_PARAMS：%+v", r)
	}

	// 完成后取结果：带 related-task
	jobs.status["a"] = mcp.TaskStatusCompleted
	var res struct {
		Content []mcp.TextContent `json:"content"`
		IsError bool              `json:"isError"`
		Meta    map[string]any    `json:"_meta"`
	}
	if r := call("tasks/result", map[string]any{"taskId": id}); r.Error != nil || json.Unmarshal(r.Result, &res) != nil ||
		res.Content[0].Text != "done a" || res.IsError || res.Meta[relatedMeta] == nil {
		t.Errorf("tasks/result 不符合预期：%+v %+v", r, res)
	}

	// 已结束的任务不能取消
	if r := call("tasks/cancel", map[string]any{"taskId": id}); r.Error == nil {
		t.Errorf("已完成的任务取消应报错：%+v", r)
	}

	// 取消进行中的任务：调用业务取消
	r = call("tools/call", map[string]any{"name": "job_run", "arguments": map[string]any{"name": "b"}, "task": map[string]any{}})
	_ = json.Unmarshal(r.Result, &created)
	if r := call("tasks/cancel", map[string]any{"taskId": created.Task.TaskId}); r.Error != nil || json.Unmarshal(r.Result, &got) != nil || got.Status != mcp.TaskStatusCancelled {
		t.Errorf("取消应返回 cancelled：%+v", r)
	}
	if len(jobs.cancelled) != 1 || jobs.cancelled[0] != "b" {
		t.Errorf("应取消作业 b，实际 %v", jobs.cancelled)
	}

	// 失败的任务：结果为 isError
	jobs.status["b"] = mcp.TaskStatusFailed
	if r := call("tasks/result", map[string]any{"taskId": created.Task.TaskId}); json.Unmarshal(r.Result, &res) != nil || !res.IsError {
		t.Errorf("失败任务的结果应为 isError：%+v", r)
	}

	// 不归本服务处理的请求交给 mcp-go：普通调用、别的任务 id、其他方法
	passed = 0
	call("tools/call", map[string]any{"name": "job_run", "arguments": map[string]any{"name": "c"}})
	call("tasks/get", map[string]any{"taskId": "other-task"})
	call("tools/list", map[string]any{})
	if passed != 3 {
		t.Errorf("应有 3 个请求交给 mcp-go，实际 %d", passed)
	}
}

func TestTaskDefDeclaresTaskSupport(t *testing.T) {
	tool := newFakeTaskTool(&fakeJobs{}).Tool()
	if tool.Execution == nil || tool.Execution.TaskSupport != mcp.TaskSupportOptional {
		t.Errorf("TaskDef 应声明 taskSupport=optional：%+v", tool.Execution)
	}
}