Agent-Squad/cmd/agent-squad/server/server_test.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"agent-squad/squad"
	"agent-squad/store"

	"error-code-system/httpx"
	"logging-system/logsys"
)

type scriptBackend struct{ script map[string]map[string]string }

func (b *scriptBackend) Name() string { return "script" }

func (b *scriptBackend) Execute(ctx context.Context, req squad.ExecRequest) (*squad.ExecResult, error) {
	key := filepath.Base(filepath.Clean(req.WorkDir))
	files := b.script[key]
	patches := map[string][]byte{}
	for _, t2 := range req.Targets {
		content := "\n"
		if c, ok := files[t2]; ok {
			content = c
		}
		out := filepath.Join(req.WorkDir, t2)
		_ = os.MkdirAll(filepath.Dir(out), 0o755)
		_ = os.WriteFile(out, []byte(content), 0o644)
		patches[t2] = []byte(content)
	}
	return &squad.ExecResult{Success: true, Patches: patches}, nil
}

const goodReport = "# API 测试交付物\n\n## 概述\n\n脚本后端生成的合格交付物正文，长度超过五十个字符以满足最小质检要求。\n"

func newTestServer(t *testing.T) (*httptest.Server, *squad.SquadEngine) {
	t.Helper()
	dir := t.TempDir()
	st, err := store.New(dir)
	if err != nil {
		t.Fatal(err)
	}
	reg := squad.NewRoleRegistry()
	_ = st.LoadRoles(reg)
	bps := squad.NewBlueprintRegistry()
	_ = st.LoadBlueprints(bps)
	backends := squad.NewBackendResolver()
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "n", "fields.json": `{"k":"v"}`},
		"02_writer":     {"output.md": goodReport},
	}
	backends.RegisterInstance("script", &scriptBackend{script: script})
	sink := squad.NewFanoutSink(1024)
	engine, err := squad.NewEngine(squad.Options{DataDir: dir, Registry: reg, Blueprints: bps,
		Planner: &squad.MockPlanner{Script: []squad.StepSpec{
			{ID: "s1", Role: "researcher", Title: "调研", Targets: []string{"notes.md", "fields.json"}},
			{ID: "s2", Role: "writer", Title: "成文", Targets: []string{"output.md"}},
		}}, Backends: backends, Sink: sink})
	if err != nil {
		t.Fatal(err)
	}
	ts := httptest.NewServer(New(engine, sink, st).Handler())
	t.Cleanup(ts.Close)
	return ts, engine
}

func doJSON(t *testing.T, method, url string, body any) (int, map[string]any) {
	t.Helper()
	var buf bytes.Buffer
	if body != nil {
		_ = json.NewEncoder(&buf).Encode(body)
	}
	req, _ := http.NewRequest(method, url, &buf)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	out := map[string]any{}
	if strings.Contains(resp.Header.Get("Content-Type"), "json") {
		_ = json.NewDecoder(resp.Body).Decode(&out)
	}
	return resp.StatusCode, out
}

func TestRolesCRUDAndPersist(t *testing.T) {
	ts, _ := newTestServer(t)
	code, out := doJSON(t, "POST", ts.URL+"/v1/roles", map[string]any{
		"role": "researcher", "name": "调研专员", "backend": "script", "keywords": []string{"调研"},
	})
	if code != 200 {
		t.Fatalf("create role: %d %v", code, out)
	}
	code, _ = doJSON(t, "GET", ts.URL+"/v1/roles", nil)
	if code != 200 {
		t.Fatalf("list roles: %d", code)
	}
	// list 是数组：直接解码判定。
	resp, _ := http.Get(ts.URL + "/v1/roles")
	var roles []map[string]any
	_ = json.NewDecoder(resp.Body).Decode(&roles)
	resp.Body.Close()
	if len(roles) != 1 || roles[0]["role"] != "researcher" {
		t.Fatalf("role not listed: %v", roles)
	}
	code, _ = doJSON(t, "DELETE", ts.URL+"/v1/roles/researcher", nil)
	if code != 200 {
		t.Fatalf("delete role: %d", code)
	}
	code, _ = doJSON(t, "DELETE", ts.URL+"/v1/roles/researcher", nil)
	if code != 404 {
		t.Fatalf("double delete want 404: %d", code)
	}
}

func TestRunPendingRespondViaAPI(t *testing.T) {
	ts, _ := newTestServer(t)
	_, _ = doJSON(t, "POST", ts.URL+"/v1/roles", map[string]any{"role": "researcher", "backend": "script"})
	_, _ = doJSON(t, "POST", ts.URL+"/v1/roles", map[string]any{"role": "writer", "backend": "script"})

	code, out := doJSON(t, "POST", ts.URL+"/v1/runs", map[string]any{
		"goal": map[string]any{"text": "生成报告"},
	})
	if code != 200 {
		t.Fatalf("create run: %d %v", code, out)
	}
	runID, _ := out["run_id"].(string)
	status, _ := out["status"].(string)
	if runID == "" || status != string(squad.RunCompleted) {
		t.Fatalf("run want completed, got %v", out)
	}
	code, insp := doJSON(t, "GET", ts.URL+"/v1/runs/"+runID, nil)
	if code != 200 || insp["deliverable"] != "output.md" {
		t.Fatalf("inspect: %d %v", code, insp)
	}
	code, bb := doJSON(t, "GET", ts.URL+"/v1/runs/"+runID+"/blackboard", nil)
	if code != 200 || bb["fields"] == nil {
		t.Fatalf("blackboard: %d %v", code, bb)
	}
	code, runs := doJSON(t, "GET", ts.URL+"/v1/runs", nil)
	if code != 200 {
		t.Fatalf("list runs: %d", code)
	}
	_ = runs
	// SSE：应能收到已缓冲事件（回放异步推送，带超时循环读取）。
	req, _ := http.NewRequest("GET", ts.URL+"/v1/runs/"+runID+"/events", nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") {
		t.Fatalf("SSE content type: %s", ct)
	}
	head := collectSSEReadable(t, resp)
	if !strings.Contains(head, "run_completed") && !strings.Contains(head, "run_started") {
		t.Fatalf("SSE replay missing events: %q", head)
	}
}

// collectSSEReadable 持续读取流内容直到出现 run_completed/超时（回放帧异步到达）。
func collectSSEReadable(t *testing.T, resp *http.Response) string {
	t.Helper()
	done := make(chan string, 1)
	go func() {
		var sb strings.Builder
		buf := make([]byte, 4096)
		for {
			n, err := resp.Body.Read(buf)
			if n > 0 {
				sb.Write(buf[:n])
				if s := sb.String(); strings.Contains(s, "run_completed") || strings.Contains(s, "run_failed") {
					done <- s
					return
				}
			}
			if err != nil {
				done <- sb.String()
				return
			}
		}
	}()
	select {
	case s := <-done:
		return s
	case <-time.After(3 * time.Second):
		t.Fatal("SSE stream did not deliver terminal event in time")
		return ""
	}
}

type memSink struct {
	mu      sync.Mutex
	entries []logsys.Entry
}

func (m *memSink) Write(e logsys.Entry) error {
	m.mu.Lock()
	m.entries = append(m.entries, e)
	m.mu.Unlock()
	return nil
}

func (m *memSink) Sync() error  { return nil }
func (m *memSink) Close() error { return nil }

func (m *memSink) has(traceID, code string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, e := range m.entries {
		if e.TraceID == traceID && e.Code == code {
			return true
		}
	}
	return false
}

func TestErrorEnvelopeAndTrace(t *testing.T) {
	dir := t.TempDir()
	st, err := store.New(dir)
	if err != nil {
		t.Fatal(err)
	}
	reg := squad.NewRoleRegistry()
	bps := squad.NewBlueprintRegistry()
	backends := squad.NewBackendResolver()
	sink := squad.NewFanoutSink(16)
	engine, err := squad.NewEngine(squad.Options{DataDir: dir, Registry: reg, Blueprints: bps, Backends: backends})
	if err != nil {
		t.Fatal(err)
	}
	mem := &memSink{}
	logger := logsys.New(logsys.Options{Level: logsys.LevelInfo, Sinks: []logsys.Sink{mem}})
	srv := New(engine, sink, st)
	srv.SetLogger(logger)
	ts := httptest.NewServer(srv.Handler())
	t.Cleanup(ts.Close)

	req, err := http.NewRequest(http.MethodPost, ts.URL+"/v1/runs", strings.NewReader(`{"goal":{}}`))
	if err != nil {
		t.Fatal(err)
	}
	const traceID = "trace-test-1234"
	req.Header.Set(httpx.TraceHeader, traceID)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusBadRequest {
		t.Fatalf("status %d", resp.StatusCode)
	}
	if got := resp.Header.Get(httpx.TraceHeader); got != traceID {
		t.Fatalf("header trace %q", got)
	}
	var body struct {
		Error struct {
			Code    string `json:"code"`
			TraceID string `json:"trace_id"`
		} `json:"error"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
		t.Fatal(err)
	}
	if body.Error.Code != "squad.goal_required" || body.Error.TraceID != traceID {
		t.Fatalf("envelope %+v", body.Error)
	}
	_ = logger.Sync()
	if !mem.has(traceID, "squad.goal_required") {
		t.Fatalf("log missing trace/code: %+v", mem.entries)
	}
}

func TestHealthAndBlueprints(t *testing.T) {
	ts, _ := newTestServer(t)
	code, out := doJSON(t, "GET", ts.URL+"/v1/health", nil)
	if code != 200 || out["ok"] != true {
		t.Fatalf("health: %d %v", code, out)
	}
	code, _ = doJSON(t, "POST", ts.URL+"/v1/blueprints", map[string]any{
		"key": "r", "title": "报告", "steps": []map[string]any{{"role": "writer", "title": "成文"}},
	})
	if code != 201 {
		t.Fatal("blueprint create failed")
	}
	resp, _ := http.Get(ts.URL + "/v1/blueprints")
	var bps []map[string]any
	_ = json.NewDecoder(resp.Body).Decode(&bps)
	resp.Body.Close()
	if len(bps) != 1 {
		t.Fatalf("blueprints: %v", bps)
	}
}

func TestCreateRunAcceptsLeaderModeOptions(t *testing.T) {
	ts, eng := newTestServer(t)
	defer ts.Close()
	// leader 模式先过标准计划校验（规则骨架引用 researcher/writer/reviewer），
	// 再进 leader 循环——四类角色都要注册。
	for _, role := range []string{"leader", "researcher", "writer", "reviewer"} {
		_, _ = doJSON(t, "POST", ts.URL+"/v1/roles", map[string]any{"role": role, "backend": "script"})
	}
	// 直接驱动引擎拿可读错误（HTTP 信封吞 reason）。
	if _, gerr := eng.Run(context.Background(), squad.Goal{Text: "写一份关于主题X的调研报告，输出markdown交付物"},
		squad.RunOptions{Mode: "leader", LeaderRole: "leader", MaxLeaderTurns: 3}); gerr != nil {
		t.Fatalf("engine leader run must return a result not an error, got: %v", gerr)
	}
	code, out := doJSON(t, "POST", ts.URL+"/v1/runs", map[string]any{
		"goal":             map[string]any{"text": "写一份关于主题X的调研报告，输出markdown交付物"},
		"mode":             "leader",
		"leader_role":      "leader",
		"max_leader_turns": 3,
	})
	if code != http.StatusOK {
		t.Fatalf("leader options must be accepted, got http %d: %v", code, out)
	}
	runID, _ := out["run_id"].(string)
	if runID == "" {
		t.Fatalf("run id missing: %v", out)
	}
}