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

package runtime

import (
	"context"
	"strings"
	"testing"

	agentllm "agent-native/runtime/llm"
	"agent-system-backend/internal/prompts"
	core "agent-system-core"
)

type subRC struct {
	fakeRunCtx
	id    string
	hint  string
	task  string
	tools []core.ToolDesc
}

func (s *subRC) SubagentID() string     { return s.id }
func (s *subRC) Subhint() string        { return s.hint }
func (s *subRC) Subtask() string        { return s.task }
func (s *subRC) Tools() []core.ToolDesc { return s.tools }

func subProfile() core.AgentProfile {
	return core.AgentProfile{Model: core.ModelRoute{Chat: "test-model", Lightweight: "lite-model"}}
}

func subTools() []core.ToolDesc {
	names := []string{
		"knowledge_search", "web_search", "delegate_task", "squad_orchestrate",
		"run_python", "register_artifact", "write_artifact", "quality_review",
		"fs_read_file", "fs_list_dir",
	}
	out := make([]core.ToolDesc, 0, len(names))
	for _, name := range names {
		out = append(out, core.ToolDesc{Name: name, Description: name, Parameters: []byte(`{"type":"object"}`)})
	}
	return out
}

func mustPrompts(t *testing.T) *prompts.Service {
	t.Helper()
	lib, err := prompts.DefaultLibrary()
	if err != nil {
		t.Fatal(err)
	}
	svc, err := prompts.NewService(nil, lib)
	if err != nil {
		t.Fatal(err)
	}
	return svc
}

func TestSubagentLoopRunsToolThenFinal(t *testing.T) {
	turns := [][]agentllm.Event{
		textDeltaTurn("先查一下。", toolCall("knowledge_search", `{"q":"x"}`)),
		textDeltaTurn("查完了。"),
	}
	var decls []agentllm.ToolDecl
	var user string
	base := scriptStream(turns, nil)
	rc := &subRC{
		fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()},
		id:         "agent.sub.explore",
		hint:       "只要路径",
		task:       "找出配置文件",
		tools:      subTools(),
	}
	r := &SubLLMRunner{
		Prompts: mustPrompts(t),
		Stream: func(ctx context.Context, model string, msgs []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
			if model != "lite-model" {
				t.Errorf("explore model = %q", model)
			}
			if decls == nil {
				decls = tools
			}
			for _, m := range msgs {
				if m.Role == agentllm.RoleUser {
					user = m.Content
				}
			}
			return base(ctx, model, msgs, tools, opts)
		},
	}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if rc.toolRuns != 1 {
		t.Fatalf("tool runs = %d", rc.toolRuns)
	}
	if !strings.Contains(user, "找出配置文件") || !strings.Contains(user, "执行提示：只要路径") {
		t.Fatalf("user message = %q", user)
	}
	names := map[string]bool{}
	for _, d := range decls {
		names[d.Name] = true
	}
	if !names["knowledge_search"] || names["delegate_task"] || names["squad_orchestrate"] || names["run_python"] {
		t.Fatalf("decls = %v", names)
	}
	var sawFinal, sawProcess bool
	for _, p := range rc.assistant {
		if p.Type == core.PartText && p.Text == "查完了。" && !p.Process {
			sawFinal = true
		}
		if p.Type == core.PartToolCall {
			sawProcess = true
		}
	}
	if !sawFinal || !sawProcess {
		t.Fatalf("parts = %+v", rc.assistant)
	}
}

func TestSubagentWriterUsesChatModel(t *testing.T) {
	var model string
	base := scriptStream([][]agentllm.Event{textDeltaTurn("写好了。")}, nil)
	rc := &subRC{
		fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()},
		id:         "agent.sub.writer",
		task:       "写一段",
		tools:      subTools(),
	}
	r := &SubLLMRunner{
		Prompts: mustPrompts(t),
		Stream: func(ctx context.Context, m string, msgs []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
			model = m
			return base(ctx, m, msgs, tools, opts)
		},
	}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if model != "test-model" {
		t.Fatalf("writer model = %q", model)
	}
}

func TestSubagentMaxTurns(t *testing.T) {
	turns := make([][]agentllm.Event, 7)
	for i := range turns {
		turns[i] = textDeltaTurn("", toolCall("knowledge_search", `{}`))
	}
	calls := 0
	base := scriptStream(turns, nil)
	rc := &subRC{
		fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()},
		id:         "agent.sub.plan",
		task:       "列步骤",
		tools:      subTools(),
	}
	r := &SubLLMRunner{
		Prompts: mustPrompts(t),
		Stream: func(ctx context.Context, model string, msgs []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
			calls++
			return base(ctx, model, msgs, tools, opts)
		},
	}
	err := r.Execute(rc)
	if err == nil || !strings.Contains(err.Error(), "max turns") {
		t.Fatalf("err = %v", err)
	}
	if calls != 6 {
		t.Fatalf("stream calls = %d", calls)
	}
}

func TestSubagentArtifactPart(t *testing.T) {
	rc := &subRC{
		fakeRunCtx: fakeRunCtx{
			ctx:          context.Background(),
			prof:         subProfile(),
			execArtifact: &core.Artifact{ID: "art_0123456789ab", Name: "note.md", Kind: "document"},
		},
		id:    "agent.sub.general",
		task:  "记一笔",
		tools: subTools(),
	}
	r := &SubLLMRunner{
		Prompts: mustPrompts(t),
		Stream: scriptStream([][]agentllm.Event{
			textDeltaTurn("", toolCall("write_artifact", `{}`)),
			textDeltaTurn("记下了。"),
		}, nil),
	}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	var saw bool
	for _, p := range rc.assistant {
		if p.Type == core.PartArtifactRef && p.Artifact != nil && p.Artifact.ArtifactID == "art_0123456789ab" && p.Artifact.Name == "note.md" {
			saw = true
		}
	}
	if !saw {
		t.Fatalf("parts = %+v", rc.assistant)
	}
}

type spawnRC struct {
	fakeRunCtx
	agent  string
	called int
}

func (s *spawnRC) SpawnAgent(agentID, name, task, hint string) (core.AgentResultPayload, error) {
	s.called++
	s.agent = agentID
	return core.AgentResultPayload{Summary: "写好了\n工件：art_0123456789ab 说明.md；工件：art_ffffffffffff 没有"}, nil
}

func TestDelegatePinsExistingArtifact(t *testing.T) {
	for _, tc := range []struct {
		name  string
		args  string
		agent string
	}{
		{name: "writer", args: `{"agent":"writer","name":"写手","task":"写说明"}`, agent: "agent.sub.writer"},
		{name: "default", args: `{"name":"助手","task":"看一下"}`, agent: "agent.sub.general"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			rc := &spawnRC{fakeRunCtx: fakeRunCtx{
				ctx:          context.Background(),
				prof:         subProfile(),
				execArtifact: &core.Artifact{ID: "art_0123456789ab", Name: "说明.md", Kind: "document"},
			}}
			r := &LLMRunner{Stream: scriptStream([][]agentllm.Event{
				textDeltaTurn("", agentllm.ToolCall{ID: "c-del", Name: "delegate_task", Arguments: tc.args}),
				textDeltaTurn("完成。"),
			}, nil)}
			if err := r.Execute(rc); err != nil {
				t.Fatal(err)
			}
			if rc.called != 1 || rc.agent != tc.agent {
				t.Fatalf("spawn agent=%q called=%d", rc.agent, rc.called)
			}
			var pinned []string
			for _, p := range rc.assistant {
				if p.Type == core.PartArtifactRef && p.Artifact != nil {
					pinned = append(pinned, p.Artifact.ArtifactID)
				}
			}
			if len(pinned) != 1 || pinned[0] != "art_0123456789ab" {
				t.Fatalf("pinned = %v", pinned)
			}
		})
	}
}