AutoZero/agent/trace_test.go

package agent

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

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

// TestModelCallTrace：每轮模型调用都应记录 StartedAt/DurationMs/Step。
func TestModelCallTrace(t *testing.T) {
	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		time.Sleep(6 * time.Millisecond)
		return llm.Message{Role: llm.RoleAssistant, Content: "done"}, nil
	}}
	a := &Agent{
		Provider: prov,
		Tools:    tools.NewRegistry(),
		Session:  &session.Session{},
		System:   "sys",
		Opt:      Options{Stdout: io.Discard, Stderr: io.Discard},
	}
	if err := a.Run(context.Background(), "hi"); err != nil {
		t.Fatal(err)
	}
	msgs := a.Session.Messages
	last := msgs[len(msgs)-1]
	if last.Role != llm.RoleAssistant {
		t.Fatalf("last message should be assistant, got %s", last.Role)
	}
	if last.StartedAt == 0 {
		t.Error("assistant StartedAt not recorded")
	}
	if last.DurationMs < 5 {
		t.Errorf("assistant DurationMs = %d, want >= 5", last.DurationMs)
	}
	if last.Step != 1 {
		t.Errorf("assistant Step = %d, want 1", last.Step)
	}
}

// slowTool 是一个可计时的假工具，delay 控制耗时。
func slowTool(name string, ro, safe bool, delay time.Duration) *tools.Tool {
	return &tools.Tool{
		Name:            name,
		Parameters:      map[string]any{"type": "object"},
		ReadOnly:        ro,
		ConcurrencySafe: safe,
		Run: func(_ context.Context, _ string) (tools.Result, error) {
			time.Sleep(delay)
			return tools.TextResult("ok"), nil
		},
	}
}

// toolMsgFor 取某次工具调用的回填消息。
func toolMsgFor(sess *session.Session, id string) (llm.Message, bool) {
	for _, m := range sess.Messages {
		if m.Role == llm.RoleTool && m.ToolCallID == id {
			return m, true
		}
	}
	return llm.Message{}, false
}

// TestToolCallTrace：工具耗时与所属轮次写进 tool 消息，且 done 事件带耗时。
func TestToolCallTrace(t *testing.T) {
	a, sess := newTestAgent(t, slowTool("slow", false, false, 6*time.Millisecond))
	a.curStep = 3 // 模拟处于第 3 次模型调用

	var events []ToolEvent
	a.Opt.OnTool = func(ev ToolEvent) { events = append(events, ev) }

	if err := a.execTools(context.Background(), []llm.ToolCall{{ID: "c1", Name: "slow", Input: "{}"}}); err != nil {
		t.Fatal(err)
	}

	msg, ok := toolMsgFor(sess, "c1")
	if !ok {
		t.Fatal("tool result message not appended")
	}
	if msg.StartedAt == 0 || msg.DurationMs < 5 {
		t.Errorf("tool trace missing: %+v", msg)
	}
	if msg.Step != 3 {
		t.Errorf("tool Step = %d, want 3 (parent turn)", msg.Step)
	}

	// done 事件应带耗时（供实时 UI 使用）。
	var done *ToolEvent
	for i := range events {
		if events[i].ID == "c1" && events[i].Done {
			done = &events[i]
		}
	}
	if done == nil {
		t.Fatal("no done event for c1")
	}
	if done.DurationMs < 5 || done.StartedAt == 0 {
		t.Errorf("done event missing trace: %+v", *done)
	}
}

// TestConcurrentToolTraceAttribution：并发批次的耗时按各自索引归属，不串号。
func TestConcurrentToolTraceAttribution(t *testing.T) {
	a, sess := newTestAgent(t,
		slowTool("fast", true, true, 5*time.Millisecond),
		slowTool("slowish", true, true, 30*time.Millisecond),
	)
	calls := []llm.ToolCall{
		{ID: "c1", Name: "fast", Input: "{}"},
		{ID: "c2", Name: "slowish", Input: "{}"},
	}
	if err := a.execTools(context.Background(), calls); err != nil {
		t.Fatal(err)
	}
	fast, _ := toolMsgFor(sess, "c1")
	slowish, _ := toolMsgFor(sess, "c2")
	if fast.DurationMs <= 0 || slowish.DurationMs <= 0 {
		t.Fatalf("durations not recorded: fast=%d slowish=%d", fast.DurationMs, slowish.DurationMs)
	}
	if fast.DurationMs >= slowish.DurationMs {
		t.Errorf("durations attributed to the wrong call: fast=%d slowish=%d", fast.DurationMs, slowish.DurationMs)
	}
}

// TestModelIdentityStamped：assistant 消息落盘要带「协议 + 模型」归因戳，
// 缓存明细也得跟着落进会话文件；没实现 Identifier 的 provider（测试替身）
// 必须留空而不是炸掉。
func TestModelIdentityStamped(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n")
		io.WriteString(w, "data: {\"choices\":[{\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":11,"+
			"\"completion_tokens\":2,\"prompt_tokens_details\":{\"cached_tokens\":9}}}\n\n")
		io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer srv.Close()
	t.Setenv("HOME", t.TempDir())

	a := &Agent{
		Provider: &llm.OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "gpt-test"},
		Tools:    tools.NewRegistry(), Session: &session.Session{}, System: "sys",
		Opt: Options{Stdout: io.Discard, Stderr: io.Discard},
	}
	if err := a.Run(context.Background(), "hi"); err != nil {
		t.Fatal(err)
	}
	last := a.Session.Messages[len(a.Session.Messages)-1]
	if last.Format != llm.FormatChatCompletions || last.Model != "gpt-test" {
		t.Errorf("归因戳 = %q/%q, want chat_completions/gpt-test", last.Format, last.Model)
	}
	if last.Usage == nil || last.Usage.CacheReadTokens != 9 || last.Usage.InputTokens != 11 {
		t.Errorf("明细没落到会话：%+v", last.Usage)
	}

	prov := &scriptedProvider{fn: func(_ int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		return llm.Message{Role: llm.RoleAssistant, Content: "x"}, nil
	}}
	b := &Agent{Provider: prov, Tools: tools.NewRegistry(), Session: &session.Session{}, System: "sys",
		Opt: Options{Stdout: io.Discard, Stderr: io.Discard}}
	if err := b.Run(context.Background(), "hi"); err != nil {
		t.Fatal(err)
	}
	end := b.Session.Messages[len(b.Session.Messages)-1]
	if end.Format != "" || end.Model != "" {
		t.Errorf("未实现 Identifier 时不该有戳：%q/%q", end.Format, end.Model)
	}
}