Agent-Squad/squad/router_test.go

package squad

import (
	"context"
	"os"
	"path/filepath"
	"testing"
)

type fixedRouter struct{ dec *RouteDecision }

func (r fixedRouter) Route(ctx context.Context, goal Goal) (*RouteDecision, error) {
	return r.dec, nil
}

func TestRouterDirectRunsAssistantOnly(t *testing.T) {
	dir := t.TempDir()
	reg := NewRoleRegistry()
	if err := reg.Register(RoleSpec{Role: "assistant", Name: "助手", Backend: "script"}); err != nil {
		t.Fatal(err)
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: map[string]map[string]string{
		"01_assistant": {"output.md": goodReport},
	}})
	engine, err := NewEngine(Options{
		DataDir: dir, Registry: reg, Planner: NewMockPlanner(), Backends: backends,
		HitlPolicy: HitlPolicy{Mode: HitlDisabled},
		Router:     fixedRouter{dec: &RouteDecision{Mode: RouteDirect, Family: "minutes", Reason: "短任务"}},
	})
	if err != nil {
		t.Fatal(err)
	}
	res, err := engine.Run(context.Background(), Goal{Text: "写一句纪要"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("status %s err=%s", res.Status, res.Error)
	}
	if res.Plan == nil || len(res.Plan.Steps) != 1 || res.Plan.Steps[0].Role != "assistant" {
		t.Fatalf("plan = %+v", res.Plan)
	}
}

func TestRouterPlanEntersLeaderMode(t *testing.T) {
	dir := t.TempDir()
	reg := NewRoleRegistry()
	_ = reg.Register(RoleSpec{Role: "writer", Name: "撰写", Backend: "script"})
	backends := NewBackendResolver()
	backends.RegisterInstance("script", &scriptBackend{script: map[string]map[string]string{}})
	engine, err := NewEngine(Options{
		DataDir: dir, Registry: reg,
		Planner:    mockScriptPlanner(StepSpec{ID: "w", Role: "writer", Title: "写", Targets: []string{"output.md"}}),
		Backends:   backends,
		HitlPolicy: HitlPolicy{Mode: HitlDisabled},
		Router:     fixedRouter{dec: &RouteDecision{Mode: RoutePlan, Family: "report"}},
	})
	if err != nil {
		t.Fatal(err)
	}
	res, err := engine.Run(context.Background(), Goal{Text: "写完整报告"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunFailed || res.Error == "" {
		t.Fatalf("missing leader should fail closed, got %s %s", res.Status, res.Error)
	}
}

func TestDelegationCarriesSkillsAndScore(t *testing.T) {
	dir := t.TempDir()
	body := []byte(`{"action":"delegate","member":"writer","title":"成文","skills":["docx-author"],"min_score":70}`)
	if err := os.WriteFile(filepath.Join(dir, "delegation.json"), body, 0o644); err != nil {
		t.Fatal(err)
	}
	dec, err := parseDelegation(dir)
	if err != nil {
		t.Fatal(err)
	}
	if len(dec.Skills) != 1 || dec.Skills[0] != "docx-author" || dec.MinScore != 70 {
		t.Fatalf("delegation = %+v", dec)
	}
}

type captureSkillsBackend struct{ got []string }

func (b *captureSkillsBackend) Name() string { return "capture-skills" }

func (b *captureSkillsBackend) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	b.got = append([]string{}, req.Skills...)
	for _, target := range req.Targets {
		path := filepath.Join(req.WorkDir, target)
		_ = os.MkdirAll(filepath.Dir(path), 0o755)
		_ = os.WriteFile(path, []byte(goodReport), 0o644)
	}
	return &ExecResult{Success: true, Patches: map[string][]byte{"output.md": []byte(goodReport)}}, nil
}

func TestStepSkillsOverrideRoleSkills(t *testing.T) {
	dir := t.TempDir()
	reg := NewRoleRegistry()
	_ = reg.Register(RoleSpec{Role: "writer", Name: "撰写", Backend: "capture-skills", Skills: []string{"role-skill"}})
	got := &captureSkillsBackend{}
	backends := NewBackendResolver()
	backends.RegisterInstance("capture-skills", got)
	engine, err := NewEngine(Options{
		DataDir: dir, Registry: reg, Backends: backends,
		HitlPolicy: HitlPolicy{Mode: HitlDisabled},
		Planner: &MockPlanner{Script: []StepSpec{{
			ID: "w", Role: "writer", Title: "写", Targets: []string{"output.md"},
			SubWorkspace: "tasks/01_writer", Skills: []string{"docx-author"}, MinScore: 10,
		}}},
	})
	if err != nil {
		t.Fatal(err)
	}
	res, err := engine.Run(context.Background(), Goal{Text: "写"}, RunOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("status %s err=%s", res.Status, res.Error)
	}
	if len(got.got) != 1 || got.got[0] != "docx-author" {
		t.Fatalf("request skills = %v", got.got)
	}
}