AutoZero/agent/context_test_helpers_test.go

package agent

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

// TestEvictOldImages：只保留最近 K 张截图，驱逐发生在视图层、transcript 不动。
func TestEvictOldImages(t *testing.T) {
	img := func(i int) llm.Image { return llm.Image{Path: string(rune('a'+i)) + ".png"} }
	msgs := []llm.Message{
		{Role: llm.RoleTool, ToolCallID: "1", Content: "shot1", Images: []llm.Image{img(0)}},
		{Role: llm.RoleTool, ToolCallID: "2", Content: "shot2", Images: []llm.Image{img(1)}},
		{Role: llm.RoleTool, ToolCallID: "3", Content: "shot3", Images: []llm.Image{img(2)}},
	}
	view := evictOldImages(msgs, 2, 0)

	// transcript 不动
	for i, m := range msgs {
		if len(m.Images) != 1 {
			t.Fatalf("transcript message %d mutated: %v", i, m.Images)
		}
	}
	// 视图：最旧的被驱逐并留说明，其余保留
	if len(view[0].Images) != 0 || !strings.Contains(view[0].Content, "evicted") {
		t.Fatalf("oldest image not evicted: %+v", view[0])
	}
	if len(view[1].Images) != 1 || len(view[2].Images) != 1 {
		t.Fatalf("recent images should be kept: %+v %+v", view[1].Images, view[2].Images)
	}
	// 数量充足时零拷贝
	same := evictOldImages(msgs, 10, 0)
	if &same[0] != &msgs[0] {
		t.Fatal("expected original slice when under budget")
	}
}

// TestUsageTotals：Usage 累计到 agent 汇总。
func TestUsageTotals(t *testing.T) {
	prov := &scriptedProvider{fn: func(call int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		u := &llm.Usage{InputTokens: 10, OutputTokens: 5}
		if call == 2 {
			u = &llm.Usage{InputTokens: 20, OutputTokens: 8}
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "ok", Usage: u}, nil
	}}
	a, _ := newTestAgent(t)
	a.Provider = prov
	if err := a.Run(context.Background(), "hi"); err != nil {
		t.Fatal(err)
	}
	// 第二次 Run 再走一轮，验证跨 turn 累计
	if err := a.Run(context.Background(), "again"); err != nil {
		t.Fatal(err)
	}
	if a.usedIn != 30 || a.usedOut != 13 {
		t.Fatalf("usage totals wrong: in=%d out=%d", a.usedIn, a.usedOut)
	}
}

// TestSafeRunLenSideEffectContract：错误标注（既 ConcurrencySafe 又
// SideEffect）的工具不得进入并发批次——运行时契约兜底。
func TestSafeRunLenSideEffectContract(t *testing.T) {
	a, _ := newTestAgent(t)
	if err := a.Tools.Register(&tools.Tool{
		Name: "risky", ConcurrencySafe: true, SideEffect: true,
		Parameters: map[string]any{"type": "object"},
		Run:        func(context.Context, string) (tools.Result, error) { return tools.TextResult("ok"), nil },
	}); err != nil {
		t.Fatal(err)
	}
	calls := []llm.ToolCall{{Name: "risky"}, {Name: "read"}}
	if n := a.safeRunLen(calls, 0); n != 0 {
		t.Fatalf("side-effect tool must not be concurrent-batched, got n=%d", n)
	}
}

// TestEvictOldImagesByteBudget：字节预算触顶时只保留预算内最新的图片。
func TestEvictOldImagesByteBudget(t *testing.T) {
	mk := func(dir, name string, size int) string {
		p := filepath.Join(dir, name)
		os.WriteFile(p, make([]byte, size), 0o644)
		return p
	}
	dir := t.TempDir()
	// 三张图：最旧 2MB、中间 1MB、最新 1MB；预算 2.5MB → 只留最新两张（估算 4/3 膨胀后 1.33+1.33=2.67 > 2.5 → 只留最新一张）
	imgs := []llm.Image{
		{Path: mk(dir, "old.png", 2_000_000)},
		{Path: mk(dir, "mid.png", 1_000_000)},
		{Path: mk(dir, "new.png", 1_000_000)},
	}
	msgs := []llm.Message{
		{Role: llm.RoleTool, ToolCallID: "1", Content: "a", Images: imgs[:1]},
		{Role: llm.RoleTool, ToolCallID: "2", Content: "b", Images: imgs[1:2]},
		{Role: llm.RoleTool, ToolCallID: "3", Content: "c", Images: imgs[2:]},
	}
	view := evictOldImages(msgs, 10, 2_500_000)
	if len(view[2].Images) != 1 {
		t.Fatal("newest image must be kept")
	}
	if len(view[0].Images) != 0 || len(view[1].Images) != 0 {
		t.Fatalf("older images must be evicted: %d %d", len(view[0].Images), len(view[1].Images))
	}
	if !strings.Contains(view[0].Content, "evicted") {
		t.Fatal("eviction note missing")
	}
	// transcript 不动
	if len(msgs[0].Images) != 1 {
		t.Fatal("transcript mutated")
	}
}