Agent-Squad/server/server_test.go

package server

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

	"agent-native/exec"
	"agent-squad/squad"
	"agent-squad/store"
)

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

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

func (b *scriptBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, 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 &exec.Result{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 ""
	}
}

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)
	}
}