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 is a scheduler runner that records how many invocations overlap,
// so a test can prove the batch still goes through the scheduler's parallel gate.
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 builds a dispatch tool over a scheduler with the given options and
// foreground window.
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 re-reads the tool result through JSON, so a test asserts the
// shape the model actually receives rather than the in-memory struct.
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 returns the status the scheduler reports for id, or "" when unknown.
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)
		}
	}
	// The batch went through the scheduler's semaphore: with MaxParallel=1 the
	// three sub-agents can never run at the same time.
	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 *MergedResult, 1)
	unsub, err := s.Bus().Subscribe(context.Background(), SubagentResultSubject,
		func(_ context.Context, e bus.Event) error {
			if res, ok := e.Payload.(*MergedResult); ok {
				select {
				case results <- res:
				default:
				}
			}
			return nil
		})
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	defer unsub()

	start := time.Now()
	raw, err := sp.Execute(context.Background(), 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.Total != 1 || res.Succeeded != 1 {
			t.Fatalf("background merge = %d total / %d ok, want 1/1", res.Total, res.Succeeded)
		}
		if res.Agents[0].Text != "late" {
			t.Fatalf("background result text = %q, want late", res.Agents[0].Text)
		}
	case <-time.After(waitTimeout):
		t.Fatal("a background batch never published its merged result")
	}
}

func TestSpawnToolKeepsPerAgentTimeout(t *testing.T) {
	t.Parallel()
	s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: 40 * time.Millisecond})
	s.runner = &mockRunner{gate: make(chan struct{})} // never closed: block until the timeout
	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 drives the tool through the real dispatcher so the
// advertised JSON Schema, the argument validation and the registry's reachability
// markers are all exercised — not just 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)
	}

	// Registered with the default (both false) markers: owner-only, so a group or
	// stranger session can neither see nor invoke it (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")
	}
}