Agent-Squad/squad/process_event_test.go

package squad

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"testing"
)

const toolCallFailure = "grep: permission denied"

// concurrentToolBackend 在一次执行里并发发两次同名工具调用（start + end）：
// CallID 每次执行唯一、一成一败。用于验证过程帧经扇出后仍按 CallID 配对、
// 结局不串（并行工具调用交错到达时不再靠到达顺序猜）。
type concurrentToolBackend struct {
	inner Executor

	mu   sync.Mutex
	exec int
}

func (b *concurrentToolBackend) Name() string { return "concurrent-tool" }

func (b *concurrentToolBackend) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	b.mu.Lock()
	b.exec++
	round := b.exec
	b.mu.Unlock()

	if req.OnEvent != nil {
		var wg sync.WaitGroup
		for _, suffix := range []string{"ok", "err"} {
			wg.Add(1)
			go func(suffix string) {
				defer wg.Done()
				callID := fmt.Sprintf("call_%d_%s", round, suffix)
				req.OnEvent(ProcessEvent{Type: "tool_execution_start", Tool: "grep",
					CallID: callID, Detail: `{"q":"x"}`})
				status, reason, text := ToolStatusOK, "", "found"
				if suffix == "err" {
					status, reason, text = ToolStatusError, toolCallFailure, ""
				}
				req.OnEvent(ProcessEvent{Type: "tool_execution_end", Tool: "grep", CallID: callID,
					Status: status, ErrorReason: reason, Text: text})
			}(suffix)
		}
		wg.Wait() // 后端自己等待：并行工具调用在同一执行内收敛
	}
	if b.inner == nil {
		return &ExecResult{Success: true}, nil
	}
	return b.inner.Execute(ctx, req)
}

// detailFrame 是 detail 事件 Data 的读取视图。
type detailFrame struct {
	processType string
	tool        string
	callID      string
	status      string
	errorReason string
}

func readDetailFrames(t *testing.T, events []Event) []detailFrame {
	t.Helper()
	var out []detailFrame
	for _, ev := range events {
		if ev.Type != EvDetail {
			continue
		}
		var data map[string]string
		if err := json.Unmarshal(ev.Data, &data); err != nil {
			t.Fatalf("detail data must stay a JSON object: %v (%s)", err, ev.Data)
		}
		out = append(out, detailFrame{
			processType: data["process_type"], tool: data["tool"], callID: data["call_id"],
			status: data["status"], errorReason: data["error_reason"],
		})
	}
	return out
}

// assertToolCallsPairedByCallID 断言并发同名工具调用在 detail 事件流里按 CallID
// 配对：每个 call_id 恰好一帧 start、一帧 end，结局与该调用自身一致。
func assertToolCallsPairedByCallID(t *testing.T, frames []detailFrame) {
	t.Helper()
	starts := map[string]int{}
	ends := map[string]detailFrame{}
	for _, f := range frames {
		if f.callID == "" {
			t.Fatalf("frame lost its CallID: %+v", f)
		}
		if f.tool != "grep" {
			t.Fatalf("frame tool = %q, want grep", f.tool)
		}
		switch f.processType {
		case "tool_execution_start":
			starts[f.callID]++
		case "tool_execution_end":
			if _, dup := ends[f.callID]; dup {
				t.Fatalf("duplicate end frame for %s: %+v", f.callID, f)
			}
			ends[f.callID] = f
		}
	}
	if len(ends) == 0 {
		t.Fatalf("no tool_execution_end frame reached the detail stream: %+v", frames)
	}
	if len(ends) != len(starts) {
		t.Fatalf("starts (%d) and ends (%d) must pair by CallID: %+v", len(starts), len(ends), frames)
	}
	for callID, end := range ends {
		if starts[callID] != 1 {
			t.Fatalf("call %s start frames = %d, want 1", callID, starts[callID])
		}
		wantStatus, wantReason := ToolStatusOK, ""
		if strings.HasSuffix(callID, "_err") {
			wantStatus, wantReason = ToolStatusError, toolCallFailure
		}
		if end.status != wantStatus || end.errorReason != wantReason {
			t.Fatalf("call %s outcome = %s/%q, want %s/%q（结局必须属于本次调用）",
				callID, end.status, end.errorReason, wantStatus, wantReason)
		}
	}
}

func TestProcessEventDataLegacyProjectionUnchanged(t *testing.T) {
	// legacy 上游（无 CallID/Status/ErrorReason）的投影必须逐字节不变。
	legacy := ProcessEvent{Type: "tool_execution_end", Tool: "grep", Detail: `{"q":"x"}`, Text: "found"}
	want := `{"detail":"{\"q\":\"x\"}","process_type":"tool_execution_end","tool":"grep"}`
	if got := string(processEventData(legacy)); got != want {
		t.Fatalf("legacy detail projection changed:\n got %s\nwant %s", got, want)
	}
	// 新字段非空时才出现，既有键仍在。
	withID := legacy
	withID.CallID, withID.Status, withID.ErrorReason = "call_1", ToolStatusOK, ""
	var data map[string]string
	if err := json.Unmarshal(processEventData(withID), &data); err != nil {
		t.Fatal(err)
	}
	if data["call_id"] != "call_1" || data["status"] != ToolStatusOK {
		t.Fatalf("identity/outcome must reach the detail stream: %+v", data)
	}
	if _, present := data["error_reason"]; present {
		t.Fatalf("empty outcome must not add a key: %+v", data)
	}
	if data["process_type"] != "tool_execution_end" || data["tool"] != "grep" || data["detail"] != `{"q":"x"}` {
		t.Fatalf("existing projection keys changed: %+v", data)
	}
}

func TestStepDetailStreamCarriesConcurrentToolCallIdentity(t *testing.T) {
	script := map[string]map[string]string{
		"01_researcher": {"notes.md": "# 笔记\n内容"},
		"02_writer":     {"output.md": goodReport},
		"03_reviewer":   {"output.md": goodReport},
	}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, spec := range []RoleSpec{
		{Role: "researcher", Name: "调研", Backend: "concurrent-tool"},
		{Role: "writer", Name: "撰写", Backend: "concurrent-tool"},
		{Role: "reviewer", Name: "审核", Backend: "concurrent-tool"},
	} {
		if err := reg.Register(spec); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("concurrent-tool", &concurrentToolBackend{
		inner: &scriptBackend{script: script},
	})
	collector := &sliceSink{}
	engine, err := NewEngine(Options{DataDir: dir, Registry: reg, Planner: mockScriptPlanner(baseSteps()...),
		Backends: backends, Sink: collector})
	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("want completed, got %s (err=%s)", res.Status, res.Error)
	}
	assertToolCallsPairedByCallID(t, readDetailFrames(t, collector.events))
}

func TestLeaderDetailStreamCarriesConcurrentToolCallIdentity(t *testing.T) {
	inner := &leaderScriptBackend{
		decisions: []string{delegateJSON("writer", "撰写"), closeJSON},
		script:    map[string]map[string]string{"01_writer": {"output.md": goodReport}},
	}
	dir := t.TempDir()
	reg := NewRoleRegistry()
	for _, spec := range []RoleSpec{
		{Role: "leader", Name: "主管", Backend: "concurrent-tool"},
		{Role: "writer", Name: "撰写", Backend: "concurrent-tool"},
	} {
		if err := reg.Register(spec); err != nil {
			t.Fatal(err)
		}
	}
	backends := NewBackendResolver()
	backends.RegisterInstance("concurrent-tool", &concurrentToolBackend{inner: inner})
	collector := &sliceSink{}
	engine, err := NewEngine(Options{DataDir: dir, Registry: reg, Backends: backends, Sink: collector})
	if err != nil {
		t.Fatal(err)
	}
	res, err := engine.Run(context.Background(), Goal{Text: "产出简报"},
		RunOptions{Mode: "leader", LeaderRole: "leader"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Status != RunCompleted {
		t.Fatalf("want completed, got %s err=%s", res.Status, res.Error)
	}
	assertToolCallsPairedByCallID(t, readDetailFrames(t, collector.events))
}