Cyrene_v2/internal/runtime/orchestration/spawntool_test.go

package orchestration

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
)

// countingRunner 是一个调度器 runner，记录有多少次调用互相重叠，使测试能证明该批次
// 仍走调度器的并行闸门。
type countingRunner struct {
	delay  time.Duration
	result *loop.Result

	mu        sync.Mutex
	active    int
	maxActive int
	total     int
}

func (r *countingRunner) Run(ctx context.Context, _ string, _ []model.Message) (*loop.Result, error) {
	r.mu.Lock()
	r.active++
	r.total++
	if r.active > r.maxActive {
		r.maxActive = r.active
	}
	r.mu.Unlock()

	if r.delay > 0 {
		select {
		case <-time.After(r.delay):
		case <-ctx.Done():
		}
	}

	r.mu.Lock()
	r.active--
	r.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if r.result != nil {
		return r.result, nil
	}
	return &loop.Result{Text: "ok", Steps: 1}, nil
}

func (r *countingRunner) peak() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.maxActive
}

// spawnToolFor 在带给定选项与前台窗口的调度器上构建派发工具。
func spawnToolFor(t *testing.T, opts Options, wait time.Duration) (*Scheduler, *SpawnTool) {
	t.Helper()
	s := newScheduler(t, opts)
	return s, NewSpawnTool(s, WithSpawnWait(wait))
}

// decodeSpawnResult 经 JSON 重新读取工具结果，使测试断言的是模型实际收到的形状，而不是
// 内存里的结构体。
func decodeSpawnResult(t *testing.T, raw any) SpawnResult {
	t.Helper()
	b, err := json.Marshal(raw)
	if err != nil {
		t.Fatalf("marshal tool result: %v", err)
	}
	var out SpawnResult
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatalf("unmarshal tool result: %v", err)
	}
	return out
}

// statusOf 返回调度器为 id 报告的状态；未知时返回 ""。
func statusOf(s *Scheduler, id string) Status {
	for _, a := range s.Agents() {
		if a.ID == id {
			return a.Status
		}
	}
	return ""
}

func TestSpawnToolReturnsMergedResult(t *testing.T) {
	t.Parallel()
	s, sp := spawnToolFor(t, Options{MaxParallel: 2}, time.Second)
	s.runner = &mockRunner{result: &loop.Result{Text: "answer", Steps: 3, StopReason: "stop"}}

	raw, err := sp.Execute(context.Background(), map[string]any{
		"goal": "research X", "role": "cyrene", "mode": "engineer", "attach_info": "ctx",
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "done" {
		t.Fatalf("Status = %q, want done", got.Status)
	}
	if len(got.Agents) != 1 {
		t.Fatalf("Agents = %d, want 1", len(got.Agents))
	}
	a := got.Agents[0]
	if a.Status != string(StatusDone) || a.Text != "answer" || a.Steps != 3 {
		t.Fatalf("agent = %+v, want a done agent with text answer/3 steps", a)
	}
	if a.Goal != "research X" || a.Role != "cyrene" || a.Mode != "engineer" {
		t.Fatalf("agent = %+v, want the goal/role/mode to round-trip", a)
	}
	if a.Summary == "" || !strings.Contains(got.Summary, "1 成功") {
		t.Fatalf("summary = %q / %q, want a settled aggregate", got.Summary, a.Summary)
	}
	if statusOf(s, a.ID) != StatusDone {
		t.Fatalf("scheduler status = %q, want done", statusOf(s, a.ID))
	}
}

func TestSpawnToolParallelBatchHonorsParallelLimit(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1})
	r := &countingRunner{delay: 20 * time.Millisecond, result: &loop.Result{Text: "ok", Steps: 1}}
	s.runner = r
	sp := NewSpawnTool(s, WithSpawnWait(5*time.Second))

	raw, err := sp.Execute(context.Background(), map[string]any{
		"tasks": []any{
			map[string]any{"goal": "a"},
			map[string]any{"goal": "b"},
			map[string]any{"goal": "c"},
		},
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "done" || len(got.Agents) != 3 {
		t.Fatalf("status/agents = %q/%d, want done/3", got.Status, len(got.Agents))
	}
	if got.Agents[0].Goal != "a" || got.Agents[1].Goal != "b" || got.Agents[2].Goal != "c" {
		t.Fatalf("agents = %+v, want dispatch order a/b/c", got.Agents)
	}
	for _, a := range got.Agents {
		if a.Status != string(StatusDone) {
			t.Fatalf("agent %s status = %q, want done", a.ID, a.Status)
		}
	}
	// 该批次确实走过了调度器的信号量：MaxParallel=1 时三个子代理绝不可能同时运行。
	if peak := r.peak(); peak != 1 {
		t.Fatalf("max concurrent runs = %d, want 1 (the parallel limit)", peak)
	}
	if !strings.Contains(got.Summary, "3 成功") {
		t.Fatalf("summary = %q, want 3 成功", got.Summary)
	}
}

func TestSpawnToolMovesSlowBatchToBackground(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 2})
	gate := make(chan struct{})
	s.runner = &mockRunner{gate: gate, result: &loop.Result{Text: "late", Steps: 2}}
	sp := NewSpawnTool(s, WithSpawnWait(40*time.Millisecond))

	results := make(chan *SubagentResultEvent, 1)
	unsub, err := s.Bus().Subscribe(context.Background(), SubagentResultSubject,
		func(_ context.Context, e bus.Event) error {
			if res, ok := e.Payload.(*SubagentResultEvent); ok {
				select {
				case results <- res:
				default:
				}
			}
			return nil
		})
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	defer unsub()

	// 派发回合的会话（loop 在 Run 顶部把它播进 ctx）：回填载荷必须带上它，
	// gateway 的中继才能把结果归属到父会话的实时流（design/05 §9）。
	ctx := loop.WithFlowBranch(context.Background(), parentSession)
	start := time.Now()
	raw, err := sp.Execute(ctx, map[string]any{"goal": "slow"})
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "running" {
		t.Fatalf("Status = %q, want running (the batch outlived the foreground window)", got.Status)
	}
	if elapsed > time.Second {
		t.Fatalf("Execute blocked for %v, want it to hand the turn back at ~40ms", elapsed)
	}
	if len(got.Agents) != 1 || got.Agents[0].Status != string(StatusRunning) || got.Agents[0].Goal != "slow" {
		t.Fatalf("agents = %+v, want a running acknowledgement", got.Agents)
	}

	close(gate)
	select {
	case res := <-results:
		if res.SessionID != parentSession {
			t.Fatalf("result sessionId = %q, want the dispatching session %q", res.SessionID, parentSession)
		}
		if res.Status != "done" || len(res.Agents) != 1 {
			t.Fatalf("result = %+v, want one settled agent", res)
		}
		if a := res.Agents[0]; a.Text != "late" || a.Steps != 2 || a.Goal != "slow" || a.Status != string(StatusDone) {
			t.Fatalf("result agent = %+v, want the late/2-steps/done projection", a)
		}
		// 载荷形状是 api-client §3.10 登记的 WS 帧 data：键集锁定，避免
		// 生产者悄悄改键而文档与实现分叉。
		wire, merr := json.Marshal(res)
		if merr != nil {
			t.Fatalf("marshal result payload: %v", merr)
		}
		var keys map[string]json.RawMessage
		if uerr := json.Unmarshal(wire, &keys); uerr != nil {
			t.Fatalf("unmarshal result payload: %v", uerr)
		}
		for _, want := range []string{"sessionId", "status", "summary", "agents"} {
			if _, ok := keys[want]; !ok {
				t.Fatalf("result payload %s lacks key %q", wire, want)
			}
		}
	case <-time.After(waitTimeout):
		t.Fatal("a background batch never published its result payload")
	}
}

func TestSpawnToolKeepsPerAgentTimeout(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: 40 * time.Millisecond})
	s.runner = &mockRunner{gate: make(chan struct{})} // 永不关闭：阻塞到超时
	sp := NewSpawnTool(s, WithSpawnWait(2*time.Second))

	raw, err := sp.Execute(context.Background(), map[string]any{"goal": "hang"})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "done" {
		t.Fatalf("Status = %q, want done (the batch settled by timing out)", got.Status)
	}
	a := got.Agents[0]
	if a.Status != string(StatusAborted) {
		t.Fatalf("agent status = %q, want aborted", a.Status)
	}
	if !strings.Contains(a.Summary, "超时") {
		t.Fatalf("summary = %q, want a timeout notice", a.Summary)
	}
}

func TestSpawnToolAbortStillCancels(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1})
	s.runner = &mockRunner{gate: make(chan struct{})}
	sp := NewSpawnTool(s, WithSpawnWait(20*time.Millisecond))

	raw, err := sp.Execute(context.Background(), map[string]any{"goal": "hang", "background": true})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if got.Status != "running" || len(got.Agents) != 1 {
		t.Fatalf("status/agents = %q/%d, want running/1", got.Status, len(got.Agents))
	}
	id := got.Agents[0].ID

	waitFor(t, func() bool { return statusOf(s, id) == StatusRunning }, waitTimeout)
	if err := s.Abort(context.Background(), id); err != nil {
		t.Fatalf("Abort: %v", err)
	}
	waitFor(t, func() bool { return statusOf(s, id) == StatusAborted }, waitTimeout)
}

func TestSpawnToolRejectsBadArgs(t *testing.T) {
	t.Parallel()
	s, sp := spawnToolFor(t, Options{}, time.Millisecond)
	s.runner = &mockRunner{result: &loop.Result{Text: "x"}}

	tooMany := make([]any, 0, maxSpawnBatch+1)
	for i := 0; i <= maxSpawnBatch; i++ {
		tooMany = append(tooMany, map[string]any{"goal": fmt.Sprintf("g%d", i)})
	}

	cases := []struct {
		name string
		args any
		want error
	}{
		{"args are not an object", "nope", nil},
		{"no goal", map[string]any{}, ErrMissingGoal},
		{"blank goal", map[string]any{"goal": "   "}, ErrMissingGoal},
		{"empty tasks", map[string]any{"tasks": []any{}}, ErrMissingGoal},
		{"task without goal", map[string]any{"tasks": []any{map[string]any{"role": "x"}}}, ErrMissingGoal},
		{"tasks is not an array", map[string]any{"tasks": "x"}, nil},
		{"tasks entry is not an object", map[string]any{"tasks": []any{"x"}}, nil},
		{"batch over the cap", map[string]any{"tasks": tooMany}, nil},
	}
	for _, tc := range cases {
		_, err := sp.Execute(context.Background(), tc.args)
		if err == nil {
			t.Fatalf("%s: Execute succeeded, want an error", tc.name)
		}
		if tc.want != nil && !errors.Is(err, tc.want) {
			t.Fatalf("%s: err = %v, want %v", tc.name, err, tc.want)
		}
	}
	if got := len(s.Agents()); got != 0 {
		t.Fatalf("rejected calls spawned %d sub-agents, want 0", got)
	}
}

func TestSpawnToolWithoutScheduler(t *testing.T) {
	t.Parallel()
	sp := NewSpawnTool(nil)
	_, err := sp.Execute(context.Background(), map[string]any{"goal": "x"})
	if !errors.Is(err, ErrNoScheduler) {
		t.Fatalf("err = %v, want ErrNoScheduler", err)
	}
}

// TestSpawnToolThroughDispatcher 让该工具走真实的 dispatcher，使对外广告的 JSON
// Schema、参数校验与注册表的可达性标记都被演练到 —— 而不只是 Execute。
func TestSpawnToolThroughDispatcher(t *testing.T) {
	t.Parallel()
	s, sp := spawnToolFor(t, Options{MaxParallel: 2}, time.Second)
	s.runner = &mockRunner{result: &loop.Result{Text: "sub answer", Steps: 1}}

	reg := registry.New()
	d, err := tool.New(reg)
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	if err := d.Register(sp); err != nil {
		t.Fatalf("Register: %v", err)
	}

	schemas := d.ToolSchemas()
	if len(schemas) != 1 || schemas[0].Name != SpawnToolName {
		t.Fatalf("schemas = %+v, want the spawn tool advertised", schemas)
	}

	wrongArgs := map[string]any{"role": "cyrene"}
	if _, err := d.Invoke(context.Background(), SpawnToolName, wrongArgs); err == nil {
		t.Fatal("Invoke with neither goal nor tasks succeeded, want a schema/validation error")
	} else if !errors.Is(err, tool.ErrInvalidArgs) {
		t.Fatalf("err = %v, want ErrInvalidArgs", err)
	}

	raw, err := d.Invoke(context.Background(), SpawnToolName, map[string]any{"goal": "do it"})
	if err != nil {
		t.Fatalf("Invoke: %v", err)
	}
	if got := decodeSpawnResult(t, raw); got.Status != "done" || len(got.Agents) != 1 {
		t.Fatalf("result = %+v, want one settled agent", got)
	}

	// 以默认标记（两个都 false）注册：仅主人可用，因此群聊或陌生人会话既看不到也
	// 调不动它（design/11 §3）。
	if !restrict.SetContext(reg, restrict.ActorOwner).Allow(SpawnToolName) {
		t.Fatal("owner was denied the spawn tool")
	}
	if restrict.SetContext(reg, restrict.ActorGroupMember).Allow(SpawnToolName) {
		t.Fatal("group member may invoke the spawn tool; want owner-only")
	}
	if restrict.SetContext(reg, restrict.ActorStranger).Allow(SpawnToolName) {
		t.Fatal("stranger may invoke the spawn tool; want owner-only")
	}
}