AutoZero/agent/oversized_test.go

package agent

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"strings"
	"sync"
	"testing"

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

// 请求过大自救。网关的 token 计数与我们的 bytes/3 估算必然有偏差，真撞墙时拿到的
// 往往是一句 400（param=messages）或一个空补全——按字面意思处理就是把错误丢回给用户，
// 而正确动作是先把上下文压下来，再重发这一次。

// sizeProvider 记录摘要调用与真实调用各发生了几次。
type sizeProvider struct {
	mu       sync.Mutex
	real     int
	compact  int
	sizes    []int // 每次真实调用发出去的消息字节数（断言重发用的是压缩后的视图）
	failReal int   // 前 N 次真实调用返回这个错误
	err      error
	reply    string
}

func (p *sizeProvider) Stream(_ context.Context, system string, msgs []llm.Message, _ []llm.ToolDef,
	_ llm.StreamOptions, _ func(llm.StreamDelta)) (llm.Message, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if strings.HasPrefix(system, "You are a context compaction") {
		p.compact++
		return llm.Message{Role: llm.RoleAssistant, Content: "<summary>压缩后的交接</summary>"}, nil
	}
	p.real++
	if b, err := json.Marshal(msgs); err == nil {
		p.sizes = append(p.sizes, len(b))
	}
	if p.err != nil && p.real <= p.failReal {
		return llm.Message{}, p.err
	}
	return llm.Message{Role: llm.RoleAssistant, Content: p.reply}, nil
}

const badRequest = "openai: HTTP 400: {\"error\":{\"message\":\"inference request is invalid\"," +
	"\"type\":\"invalid_request_error\",\"param\":\"messages\",\"code\":\"400002\"}}"

// oversizedAgent 造一个"上下文已经不小、但还没到自动压缩阈值"的回合：
// 估算 ~7000 tokens，阈值 8500，窗口 10000（六成线 = 6000）。
// 这样常规压缩不会触发，只有"请求过大自救"这条路会去压。
func oversizedAgent(t *testing.T, p *sizeProvider) (*Agent, int) {
	t.Helper()
	a, sess := newTestAgent(t)
	a.Provider = p
	a.Opt.Stderr, a.Opt.Stdout = io.Discard, io.Discard
	a.ContextWindow, a.CompactAt, a.MinCompactTokens = 10000, 0.85, 10000
	sess.Append(llm.Message{Role: llm.RoleUser, Content: "把这件事做完"})
	sess.Append(llm.Message{Role: llm.RoleAssistant, Content: strings.Repeat("a", 21000)})
	view, err := a.contextView(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if got := estimateTokens(view, a.sysNow()); got >= 8500 || got < 6000 {
		t.Fatalf("夹具的上下文占比不对：est=%d，want 6000..8499", got)
	}
	return a, len(view)
}

// TestOversizedRequestCompactsAndResends：400 之后压一次、重发一次，回合正常收尾。
func TestOversizedRequestCompactsAndResends(t *testing.T) {
	p := &sizeProvider{failReal: 1, err: errors.New(badRequest), reply: "压完就能发出去了"}
	a, _ := oversizedAgent(t, p)
	if err := a.Run(context.Background(), ""); err != nil {
		t.Fatalf("该自救成功而不是失败：%v", err)
	}
	if p.compact != 1 {
		t.Errorf("摘要调用 %d 次，want 1", p.compact)
	}
	if p.real != 2 {
		t.Errorf("真实调用 %d 次，want 2（失败一次 + 压缩后重发一次）", p.real)
	}
	last := a.Session.Messages[len(a.Session.Messages)-1]
	if last.Content != "压完就能发出去了" {
		t.Errorf("最终回复不对：%q", last.Content)
	}
	// 重发必须用压缩后的视图：端到端实测里漏了这一步，等于把同一个 1.4MB 的请求
	// 原样再撞一次墙（第一次单测是绿的，因为假 provider 不看消息内容）。
	if len(p.sizes) != 2 {
		t.Fatalf("真实调用记录 %v，want 2 次", p.sizes)
	}
	if p.sizes[1] >= p.sizes[0] {
		t.Errorf("重发的请求没变小：%d → %d（视图没重取）", p.sizes[0], p.sizes[1])
	}
}

// TestOversizedRetryOnlyOncePerStep：压完还是 400 时不无限自救——一个 step 一次，
// 之后把原错误如实交出去。
func TestOversizedRetryOnlyOncePerStep(t *testing.T) {
	p := &sizeProvider{failReal: 99, err: errors.New(badRequest)}
	a, _ := oversizedAgent(t, p)
	err := a.Run(context.Background(), "")
	if err == nil {
		t.Fatal("一直 400 时该失败")
	}
	if !strings.Contains(err.Error(), "HTTP 400") {
		t.Errorf("错误没原样交出：%v", err)
	}
	if p.compact != 1 {
		t.Errorf("摘要调用 %d 次，want 1（同一 step 不该反复自救）", p.compact)
	}
	if p.real != 2 {
		t.Errorf("真实调用 %d 次，want 2", p.real)
	}
}

// TestRequestLooksOversizedNeedsBigContext：判定必须带"上下文本来就大"这个前提。
// 小上下文下的 400 是请求本身写错了，空补全是网关抖动——去压缩只是白费一次调用。
func TestRequestLooksOversizedNeedsBigContext(t *testing.T) {
	a, _ := newTestAgent(t)
	a.Opt.Stderr = io.Discard
	a.ContextWindow, a.CompactAt = 10000, 0.85
	big := []llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("a", 21000)}}
	small := []llm.Message{{Role: llm.RoleUser, Content: "hi"}}
	medium := []llm.Message{{Role: llm.RoleUser, Content: strings.Repeat("a", 12600)}}

	cases := []struct {
		name string
		err  error
		view []llm.Message
		want bool
	}{
		{"大上下文 + 400 param=messages", errors.New(badRequest), big, true},
		{"大上下文 + 空补全", llm.ErrEmptyCompletion, big, true},
		{"大上下文 + 坏流", llm.ErrTruncatedStream, big, true},
		{"小上下文 + 400", errors.New(badRequest), small, false},
		{"小上下文 + 空补全", llm.ErrEmptyCompletion, small, false},
		// 把 context_window 按真实模型配大（1M 的模型本该如此）之后，"占窗口六成"
		// 这条线高得够不着；400 是网关在点名说请求本身有问题，必须仍然能自救。
		{"大窗口 + 400（六成线够不着）", errors.New(badRequest), big, true},
		{"刚过地板线 + 400", errors.New(badRequest), medium, true},
		{"大窗口 + 空补全（弱证据要占比）", llm.ErrEmptyCompletion, medium, false},
		{"大上下文 + 权限错误", errors.New("permission denied"), big, false},
		{"大上下文 + 超时", errors.New("context deadline exceeded"), big, false},
		{"没配窗口", errors.New(badRequest), big, false},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			switch c.name {
			case "没配窗口":
				a.ContextWindow = 0
			case "大窗口 + 400（六成线够不着）", "大窗口 + 空补全（弱证据要占比）":
				a.ContextWindow = 500000
			default:
				a.ContextWindow = 10000
			}
			if got := a.requestLooksOversized(c.err, c.view); got != c.want {
				t.Errorf("requestLooksOversized = %v, want %v", got, c.want)
			}
		})
	}
}

// TestOversizedRetryClearsGaveUpFlag：上一 step 压缩失败留下的"本回合放弃"不该挡住
// 这次的自救——那时还没有"请求过大"这个新证据。
func TestOversizedRetryWorksAfterGaveUp(t *testing.T) {
	p := &sizeProvider{failReal: 1, err: errors.New(badRequest), reply: "压完就能发出去了"}
	a, _ := oversizedAgent(t, p)
	a.compactGaveUp = true // 模拟本回合早前压缩失败过
	if err := a.Run(context.Background(), ""); err != nil {
		t.Fatalf("放弃标记不该挡住新证据下的自救：%v", err)
	}
	if p.compact != 1 {
		t.Errorf("摘要调用 %d 次，want 1", p.compact)
	}
}