Agent-System-Backend/internal/runtime/squad_runner_test.go

package runtime

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	agentllm "agent-native/runtime/llm"
	"agent-squad/squad"
	core "agent-system-core"
)

const squadGoodDoc = "# 应先把结论写进标题\n\n一句话：范围已经明确，按会议记录执行。\n\n依据来自附件里的原话。\n"

type fixedRoute struct{ mode squad.RouteMode }

func (r fixedRoute) Route(ctx context.Context, goal squad.Goal) (*squad.RouteDecision, error) {
	return &squad.RouteDecision{Mode: r.mode, Family: "minutes", Reason: "测试"}, nil
}

type failCheck struct{}

func (failCheck) Check(req squad.CheckRequest) *squad.QualityReport {
	return &squad.QualityReport{Passed: false, Score: 10, Violations: []string{"缺结论"}}
}

type entryFake struct {
	ctx      context.Context
	prof     core.AgentProfile
	settings core.EffectiveSettings
	text     string
	inputs   []string
	files    []string
	tools    []core.ToolDesc
	sandbox  string
	events   []core.StreamEventType
	parts    []core.Part
	hitl     int
	wrote    []string
	called   []string
	usage    core.Usage
	cards    int
}

func (f *entryFake) Context() context.Context {
	if f.ctx == nil {
		return context.Background()
	}
	return f.ctx
}
func (f *entryFake) Run() core.Run {
	return core.Run{ID: "run_t", Owner: "owner", ConversationID: "conv"}
}
func (f *entryFake) Profile() core.AgentProfile       { return f.prof }
func (f *entryFake) Settings() core.EffectiveSettings { return f.settings }
func (f *entryFake) Emit(typ core.StreamEventType, payload any) error {
	f.events = append(f.events, typ)
	return nil
}
func (f *entryFake) SetAssistantParts(parts []core.Part) error { f.parts = parts; return nil }
func (f *entryFake) AskHITL(card core.HITLCard) (core.HITLResponse, error) {
	f.hitl++
	return core.HITLResponse{ChoiceID: "approve"}, nil
}
func (f *entryFake) SpawnSubagent(name, task, hint string) (core.AgentResultPayload, error) {
	return core.AgentResultPayload{}, errors.New("no spawn")
}
func (f *entryFake) UpdatePlan(p core.PlanSnapshot) error { return nil }
func (f *entryFake) UpdateTodos(l core.TodoList) error    { return nil }
func (f *entryFake) WriteArtifact(name, kind, mime string, content []byte) (core.Artifact, error) {
	f.wrote = append(f.wrote, name)
	return core.Artifact{ID: "art_aaaaaaaaaaaa", Name: name, Kind: kind}, nil
}
func (f *entryFake) Retrieve(q core.RetrievalQuery) ([]core.RetrievalHit, error) { return nil, nil }
func (f *entryFake) ExecuteTool(call core.ToolCallRequest) core.ToolCallResult {
	f.called = append(f.called, call.Name)
	if call.Name == "register_artifact" {
		return core.ToolCallResult{Status: "ok", Output: "已登记", Artifact: &core.Artifact{ID: "art_bbbbbbbbbbbb", Name: "纪要.docx", Kind: "document"}}
	}
	return core.ToolCallResult{Status: "ok", Output: "ok"}
}
func (f *entryFake) History() []core.Message {
	return []core.Message{{Role: core.RoleUser, Parts: []core.Part{{Type: core.PartText, Text: f.text}}}}
}
func (f *entryFake) Attachments() []core.ParsedDocument {
	return []core.ParsedDocument{{Meta: core.DocMeta{Filename: "notes.docx"}}}
}
func (f *entryFake) Tools() []core.ToolDesc { return f.tools }
func (f *entryFake) Subtask() string        { return "" }
func (f *entryFake) CreateCard(card core.ActivityCard) (core.ActivityCard, error) {
	f.cards++
	card.ID = "card_1"
	return card, nil
}
func (f *entryFake) SetCardStatus(id string, status core.CardStatus, tone, note string) error {
	return nil
}
func (f *entryFake) SubmitCard(card core.ActivityCard) (core.ActivityCard, error) {
	return f.CreateCard(card)
}
func (f *entryFake) BeginStreamCard(card core.ActivityCard) (core.ActivityCard, error) {
	return f.CreateCard(card)
}
func (f *entryFake) AppendStreamCard(text string) error { return nil }
func (f *entryFake) CommitStreamCard() (core.ActivityCard, error) {
	f.cards++
	return core.ActivityCard{ID: "card_1", Title: "流式"}, nil
}
func (f *entryFake) Card(id string) *core.ActivityCard { return nil }
func (f *entryFake) Artifact(id string) (core.Artifact, []byte, bool) {
	return core.Artifact{}, nil, false
}
func (f *entryFake) SetContextObservation(o core.ContextObservation) {}
func (f *entryFake) WorkspaceInputs() []string                       { return f.inputs }
func (f *entryFake) ConversationSandbox() string                     { return f.sandbox }
func (f *entryFake) AddUsage(u core.Usage) {
	f.usage.PromptTokens += u.PromptTokens
	f.usage.CompletionTokens += u.CompletionTokens
}

func writeThenStop(model *string, calls *int) agentllm.StreamFunc {
	var mu sync.Mutex
	n := 0
	body, _ := json.Marshal(map[string]string{"path": "output.md", "content": squadGoodDoc})
	return func(ctx context.Context, m string, messages []agentllm.Message, td []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 4)
		go func() {
			defer close(ch)
			if ctx.Err() != nil {
				ch <- agentllm.Event{Kind: agentllm.KindError, ErrorMessage: ctx.Err().Error()}
				return
			}
			mu.Lock()
			n++
			step := n
			mu.Unlock()
			if model != nil {
				*model = m
			}
			if calls != nil {
				*calls = step
			}
			if step%2 == 1 {
				ch <- agentllm.Event{Kind: agentllm.KindTextDelta, Text: "结论应先写进标题。"}
				ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{
					Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls,
					ToolCalls: []agentllm.ToolCall{{ID: "w", Name: "write_file", Arguments: string(body)}},
				}, Usage: agentllm.Usage{PromptTokens: 4, CompletionTokens: 2}}
				return
			}
			ch <- agentllm.Event{Kind: agentllm.KindTextDelta, Text: "写完了。"}
			ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{
				Role: agentllm.RoleAssistant, Content: "写完了。", StopReason: agentllm.StopEnd,
			}, Usage: agentllm.Usage{PromptTokens: 3, CompletionTokens: 1}}
		}()
		return ch
	}
}

func TestSquadEntryStreamsAndRegisters(t *testing.T) {
	dir := t.TempDir()
	var model string
	rc := &entryFake{
		text:    "整理会议结论",
		inputs:  []string{filepath.Join(dir, "inputs", "notes.docx")},
		sandbox: dir,
		settings: core.EffectiveSettings{
			"storage": map[string]any{"dataDir": dir},
			"models":  map[string]any{"chat": map[string]any{"model": "base", "baseUrl": "http://test", "apiKey": "k"}},
			"ui":      map[string]any{"cardflow": true},
		},
		prof:  core.AgentProfile{Model: core.ModelRoute{Chat: "profile-model"}},
		tools: []core.ToolDesc{{Name: "register_artifact", Description: "登记", Parameters: json.RawMessage(`{"type":"object"}`)}},
	}
	cardBody, _ := json.Marshal(map[string]string{"title": "纪要", "body": "结论应先写。"})
	regBody, _ := json.Marshal(map[string]string{"path": "output.md"})
	runner := &SquadRunner{
		Router:  fixedRoute{mode: squad.RouteDirect},
		Checker: squad.MinimalChecker{},
		Hitl:    squad.HitlPolicy{Mode: squad.HitlDisabled},
		Stream: scriptedCalls(t, &model, [][]agentllm.Event{
			{
				{Kind: agentllm.KindTextDelta, Text: "结论应先写进标题。"},
				{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls, ToolCalls: []agentllm.ToolCall{
					{ID: "c", Name: "submit_card", Arguments: string(cardBody)},
				}}, Usage: agentllm.Usage{PromptTokens: 2, CompletionTokens: 1}},
			},
			{
				{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls, ToolCalls: []agentllm.ToolCall{
					{ID: "r", Name: "register_artifact", Arguments: string(regBody)},
				}}, Usage: agentllm.Usage{PromptTokens: 2, CompletionTokens: 1}},
			},
			{
				{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls, ToolCalls: []agentllm.ToolCall{
					{ID: "w", Name: "write_file", Arguments: mustDocArgs()},
				}}, Usage: agentllm.Usage{PromptTokens: 2, CompletionTokens: 1}},
			},
			{
				{Kind: agentllm.KindTextDelta, Text: "写完了。"},
				{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "写完了。", StopReason: agentllm.StopEnd}, Usage: agentllm.Usage{PromptTokens: 1, CompletionTokens: 1}},
			},
		}),
	}
	if err := runner.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if model != "profile-model" {
		t.Fatalf("model = %s", model)
	}
	if !containsEvent(rc.events, core.EvMessageDelta) || !containsEvent(rc.events, core.EvToolCall) {
		t.Fatalf("events = %v", rc.events)
	}
	if rc.cards == 0 || !containsStr(rc.called, "register_artifact") {
		t.Fatalf("cards=%d tools=%v", rc.cards, rc.called)
	}
	if rc.usage.PromptTokens == 0 {
		t.Fatal("usage not recorded")
	}
	if len(rc.wrote) == 0 {
		t.Fatal("deliverable was not registered")
	}
	task := filepath.Join(dir, "squad-entry", "run_t")
	found := false
	_ = filepath.Walk(task, func(path string, info os.FileInfo, err error) error {
		if err == nil && strings.HasSuffix(path, "TASK.md") {
			b, _ := os.ReadFile(path)
			if strings.Contains(string(b), "inputs") && strings.Contains(string(b), "notes.docx") {
				found = true
			}
		}
		return nil
	})
	if !found {
		t.Fatal("goal did not carry inputs and attachment name")
	}
}

func TestSquadEntryFalseDeliverableAndBudgets(t *testing.T) {
	dir := t.TempDir()
	rc := baseEntry(dir)
	claim := &SquadRunner{
		Router: fixedRoute{mode: squad.RouteDirect}, Checker: squad.MinimalChecker{},
		Hitl: squad.HitlPolicy{Mode: squad.HitlDisabled},
		Stream: scriptedCalls(t, nil, [][]agentllm.Event{
			{{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls, ToolCalls: []agentllm.ToolCall{{ID: "w", Name: "write_file", Arguments: mustDocArgs()}}}}},
			{{Kind: agentllm.KindTextDelta, Text: "已生成 纪要.docx"}, {Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "已生成 纪要.docx", StopReason: agentllm.StopEnd}}},
		}),
	}
	if err := claim.Execute(rc); err == nil || !strings.Contains(err.Error(), "没有登记") {
		t.Fatalf("false claim err = %v", err)
	}

	var calls int
	turns := 1
	rc2 := baseEntry(dir)
	rc2.prof.Budget.MaxTurns = &turns
	limited := &SquadRunner{
		Router: fixedRoute{mode: squad.RouteDirect}, Checker: squad.MinimalChecker{},
		Hitl:   squad.HitlPolicy{Mode: squad.HitlDisabled},
		Stream: writeThenStop(nil, &calls),
	}
	// 第一轮只写了别的文件，第二轮才写 output.md；轮次上限为 1 时交付物缺失。
	limited.Stream = oneTurnOtherFile(&calls)
	if err := limited.Execute(rc2); err == nil {
		t.Fatal("turn cap should stop before the deliverable")
	}
	if calls != 1 {
		t.Fatalf("model calls = %d", calls)
	}

	tools := 1
	rc3 := baseEntry(dir)
	rc3.prof.Budget.MaxToolCalls = &tools
	var toolCalls int
	capped := &SquadRunner{
		Router: fixedRoute{mode: squad.RouteDirect}, Checker: squad.MinimalChecker{},
		Hitl:   squad.HitlPolicy{Mode: squad.HitlDisabled},
		Stream: twoTools(&toolCalls),
	}
	if err := capped.Execute(rc3); err == nil || !strings.Contains(err.Error(), "budget_exceeded") {
		t.Fatalf("tool cap err = %v", err)
	}
}

func TestSquadEntryHITLAndCancel(t *testing.T) {
	dir := t.TempDir()
	rc := baseEntry(dir)
	runner := &SquadRunner{
		Router: fixedRoute{mode: squad.RouteDirect}, Checker: failCheck{},
		Stream: writeThenStop(nil, nil),
	}
	if err := runner.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if rc.hitl == 0 {
		t.Fatal("quality hold did not ask the user")
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	rc2 := baseEntry(dir)
	rc2.ctx = ctx
	if err := (&SquadRunner{}).Execute(rc2); !errors.Is(err, context.Canceled) {
		t.Fatalf("cancel err = %v", err)
	}
}

func TestHostBridgeHonorsAllowlist(t *testing.T) {
	rc := &entryFake{tools: []core.ToolDesc{
		{Name: "register_artifact"},
		{Name: "read_file"},
		{Name: "knowledge_search"},
	}, settings: core.EffectiveSettings{}}
	rc.prof.Tools.Deny = []string{"register_artifact"}
	names := map[string]bool{}
	for _, tool := range bridgeHostTools(rc, nil) {
		names[tool.Name()] = true
	}
	if names["register_artifact"] || names["read_file"] || !names["knowledge_search"] {
		t.Fatalf("bridged = %v", names)
	}
}

func TestParseRoute(t *testing.T) {
	dec := parseRoute("好的 {\"mode\":\"direct\",\"family\":\"report\",\"reason\":\"一句就够\"}")
	if dec.Mode != squad.RouteDirect || dec.Family != "report" {
		t.Fatalf("%+v", dec)
	}
	if parseRoute("不是 json").Mode != squad.RoutePlan {
		t.Fatal("unparsed route should stay on the plan path")
	}
}

func TestSquadToolHiddenWhenEntryIsSquad(t *testing.T) {
	eff := core.EffectiveSettings{"runtime": map[string]any{"entry": "squad", "squad": true}}
	if squadToolEnabled(eff) {
		t.Fatal("squad entry must not expose squad_orchestrate")
	}
	if !squadToolEnabled(core.EffectiveSettings{"runtime": map[string]any{"squad": true}}) {
		t.Fatal("direct entry keeps the tool")
	}
}

func baseEntry(dir string) *entryFake {
	return &entryFake{
		text: "整理会议结论",
		settings: core.EffectiveSettings{
			"storage": map[string]any{"dataDir": dir},
			"models":  map[string]any{"chat": map[string]any{"model": "m", "baseUrl": "http://test", "apiKey": "k"}},
		},
		prof: core.AgentProfile{Model: core.ModelRoute{Chat: "m"}},
	}
}

func mustDocArgs() string {
	b, _ := json.Marshal(map[string]string{"path": "output.md", "content": squadGoodDoc})
	return string(b)
}

func scriptedCalls(t *testing.T, model *string, turns [][]agentllm.Event) agentllm.StreamFunc {
	t.Helper()
	var mu sync.Mutex
	i := 0
	return func(ctx context.Context, m string, messages []agentllm.Message, td []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 8)
		go func() {
			defer close(ch)
			mu.Lock()
			n := i
			if i < len(turns) {
				i++
			}
			mu.Unlock()
			if model != nil {
				*model = m
			}
			if n >= len(turns) {
				ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "done", StopReason: agentllm.StopEnd}}
				return
			}
			for _, ev := range turns[n] {
				ch <- ev
			}
		}()
		return ch
	}
}

func oneTurnOtherFile(calls *int) agentllm.StreamFunc {
	body, _ := json.Marshal(map[string]string{"path": "note.md", "content": "not the deliverable"})
	var once sync.Once
	return func(ctx context.Context, model string, messages []agentllm.Message, td []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 2)
		go func() {
			defer close(ch)
			once.Do(func() {
				if calls != nil {
					*calls = 1
				}
			})
			if calls != nil && *calls == 1 {
				*calls = 1
			}
			ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{
				Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls,
				ToolCalls: []agentllm.ToolCall{{ID: "w", Name: "write_file", Arguments: string(body)}},
			}}
		}()
		return ch
	}
}

func twoTools(calls *int) agentllm.StreamFunc {
	var mu sync.Mutex
	n := 0
	return func(ctx context.Context, model string, messages []agentllm.Message, td []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 2)
		go func() {
			defer close(ch)
			mu.Lock()
			n++
			step := n
			mu.Unlock()
			if calls != nil {
				*calls = step
			}
			name := "write_file"
			args := mustDocArgs()
			if step > 1 {
				body, _ := json.Marshal(map[string]string{"path": "note.md", "content": "x"})
				args = string(body)
			}
			ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{
				Role: agentllm.RoleAssistant, StopReason: agentllm.StopToolCalls,
				ToolCalls: []agentllm.ToolCall{{ID: "w", Name: name, Arguments: args}},
			}}
		}()
		return ch
	}
}

func containsEvent(events []core.StreamEventType, want core.StreamEventType) bool {
	for _, ev := range events {
		if ev == want {
			return true
		}
	}
	return false
}

func containsStr(list []string, want string) bool {
	for _, v := range list {
		if v == want {
			return true
		}
	}
	return false
}