AutoZero/agent/orchestration_test.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/session"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

// quote 把字符串编码为 JSON 字符串字面量（用于拼接工具参数）。
func quote(s string) string {
	b, _ := json.Marshal(s)
	return string(b)
}

func newTestAgent(t *testing.T, extra ...*tools.Tool) (*Agent, *session.Session) {
	t.Helper()
	t.Setenv("HOME", t.TempDir())
	reg := tools.NewRegistry()
	for _, tl := range tools.Builtin() {
		if err := reg.Register(tl); err != nil {
			t.Fatal(err)
		}
	}
	for _, tl := range extra {
		if err := reg.Register(tl); err != nil {
			t.Fatal(err)
		}
	}
	sess, err := session.New()
	if err != nil {
		t.Fatal(err)
	}
	var out bytes.Buffer
	a := &Agent{
		Tools:   reg,
		Session: sess,
		System:  "test",
		Opt:     Options{Stdout: &out, Stderr: &out, ApproveAll: true},
	}
	return a, sess
}

// TestSynthMissingToolResults 验证为未配对的 tool_use 合成错误结果，且幂等。
func TestSynthMissingToolResults(t *testing.T) {
	a, sess := newTestAgent(t)
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: "go"},
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{
			{ID: "c1", Name: "ls"}, {ID: "c2", Name: "grep"},
		}},
		{Role: llm.RoleTool, ToolCallID: "c1", Content: "ok"},
	}

	a.synthMissingToolResults()

	var found bool
	for _, m := range sess.Messages {
		if m.ToolCallID == "c2" {
			found = true
			if m.Role != llm.RoleTool || m.Content == "" {
				t.Fatalf("bad synthesized message: %+v", m)
			}
		}
	}
	if !found {
		t.Fatal("missing tool_result for c2 was not synthesized")
	}

	// 幂等：重复调用不追加
	n := len(sess.Messages)
	a.synthMissingToolResults()
	if len(sess.Messages) != n {
		t.Fatalf("synth is not idempotent: %d -> %d", n, len(sess.Messages))
	}
}

// TestOrchestrationOrderAndConcurrency：两段只读并发、中间夹一个写工具，
// 验证回填顺序与调用顺序一致，且总耗时接近并发而非串行。
func TestOrchestrationOrderAndConcurrency(t *testing.T) {
	const delay = 150 * time.Millisecond
	slow := func(name string) *tools.Tool {
		return &tools.Tool{
			Name:            name,
			Description:     "slow read",
			ReadOnly:        true,
			ConcurrencySafe: true,
			Parameters:      map[string]any{"type": "object"},
			Run: func(_ context.Context, _ string) (tools.Result, error) {
				time.Sleep(delay)
				return tools.TextResult(name + " done"), nil
			},
		}
	}
	a, sess := newTestAgent(t, slow("ra"), slow("rb"), slow("rc"))
	target := filepath.Join(t.TempDir(), "out.txt")

	calls := []llm.ToolCall{
		{ID: "1", Name: "ra"}, {ID: "2", Name: "rb"},
		{ID: "3", Name: "write", Input: `{"path":` + quote(target) + `,"content":"x"}`},
		{ID: "4", Name: "rc"},
	}
	start := time.Now()
	if err := a.execTools(context.Background(), calls); err != nil {
		t.Fatal(err)
	}
	elapsed := time.Since(start)

	// 串行 >= 450ms（3×delay+写开销）；并发批次 ~2×delay。留出余量断言。
	if elapsed >= 4*delay {
		t.Fatalf("calls appear serialized: elapsed=%v", elapsed)
	}

	// 回填必须与调用顺序一致且 ID 正确配对
	results := sess.Messages[len(sess.Messages)-4:]
	wantOrder := []string{"ra", "rb", "wrote", "rc"}
	wantID := []string{"1", "2", "3", "4"}
	for i, m := range results {
		if m.Role != llm.RoleTool {
			t.Fatalf("message %d is not a tool result: %+v", i, m)
		}
		if m.ToolCallID != wantID[i] {
			t.Fatalf("result %d: got ToolCallID %q, want %q", i, m.ToolCallID, wantID[i])
		}
		if !strings.Contains(m.Content, wantOrder[i]) {
			t.Fatalf("result %d content %q does not contain %q", i, m.Content, wantOrder[i])
		}
	}
}

// TestRunInterruptedPairing：assistant 发起多个调用但只回填了部分
// （模拟 Append 中途失败），synth 必须把缺口补齐。
func TestRunInterruptedPairing(t *testing.T) {
	a, sess := newTestAgent(t)
	// 构造：assistant 发起 3 个调用，只回填了 2 个（模拟 Append 中途失败）
	sess.Messages = []llm.Message{
		{Role: llm.RoleUser, Content: "go"},
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{
			{ID: "a", Name: "ls"}, {ID: "b", Name: "grep"}, {ID: "c", Name: "read"},
		}},
		{Role: llm.RoleTool, ToolCallID: "a", Content: "ok"},
		{Role: llm.RoleTool, ToolCallID: "b", Content: "ok"},
	}
	a.synthMissingToolResults()
	for _, id := range []string{"a", "b", "c"} {
		found := false
		for _, m := range sess.Messages {
			if m.ToolCallID == id && m.Role == llm.RoleTool {
				found = true
			}
		}
		if !found {
			t.Fatalf("tool result for %q missing after synth", id)
		}
	}
}