Cyrene_v2/internal/runtime/loop/context_pressure_test.go

package loop

import (
	"bytes"
	"context"
	"encoding/json"
	"log/slog"
	"strings"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// TestReportContextPressureWarnsAtNinetyPercent 钉住输入逼近预算的 WARN
// （design/04 附录 A.3 第 5 条）：达到预算 90% 记一条 WARN、低于不记，且**不阻断**
// 任何流程（它只是日志）。
func TestReportContextPressureWarnsAtNinetyPercent(t *testing.T) {
	t.Parallel()
	const budget = 1000
	cases := []struct {
		name         string
		promptTokens int
		wantWarn     bool
	}{
		{"恰好 90%", 900, true},
		{"超过 90%", 990, true},
		{"低于 90%", 899, false},
		{"提供商未上报用量", 0, false},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			t.Parallel()
			log, buf := captureLoopLog()
			l := New(Options{MaxContextTokens: budget, Log: log})
			l.reportContextPressure(c.promptTokens)

			got := strings.Contains(buf.String(), "near the context budget")
			if got != c.wantWarn {
				t.Fatalf("promptTokens=%d budget=%d: WARN=%v, want %v（日志：%q）",
					c.promptTokens, budget, got, c.wantWarn, buf.String())
			}
			if !c.wantWarn {
				return
			}
			records := loopJSONRecords(t, buf)
			if len(records) != 1 || records[0]["prompt_tokens"] != float64(c.promptTokens) {
				t.Fatalf("告警记录 = %+v, want 恰好一条且带上 prompt_tokens=%d", records, c.promptTokens)
			}
		})
	}
}

// TestReportContextPressureSilentWhenCompactionOff 钉住「压缩关闭时不吵」：
// 关闭压缩（预算 0/负）的部署没有预算可言，绝不能因为预算为 0 而把每条调用都
// 判成「超 90%」。
func TestReportContextPressureSilentWhenCompactionOff(t *testing.T) {
	t.Parallel()
	for _, budget := range []int{0, -1} {
		log, buf := captureLoopLog()
		l := New(Options{MaxContextTokens: budget, Log: log})
		l.reportContextPressure(1 << 20)
		if strings.Contains(buf.String(), "near the context budget") {
			t.Fatalf("预算 %d 时不该记 WARN：%q", budget, buf.String())
		}
	}
}

// TestStreamStepCarriesPromptTokensForPressure 钉住接线：流式主路径把提供商上报的
// 输入 token 带进 stepOutcome，逼近预算的判定因此真的会在这条路径上发生 ——
// 而不是只在 reportContextPressure 这个函数里「理论上可行」。
func TestStreamStepCarriesPromptTokensForPressure(t *testing.T) {
	t.Parallel()
	log, buf := captureLoopLog()
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Model: "m", Content: "好", PromptTokens: 950}}}}
	l := New(Options{Agent: newTestAgent(t, ad, nil), MaxContextTokens: 1000, Log: log,
		OnStream: func(model.StreamEvent) {}})

	if _, err := l.Run(context.Background(), "sess-pressure", []model.Message{model.TextMessage(model.RoleUser, "在吗")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if !strings.Contains(buf.String(), "near the context budget") {
		t.Fatalf("流式路径没有触发预算告警：%q", buf.String())
	}
	records := loopJSONRecords(t, buf)
	found := false
	for _, rec := range records {
		if rec["msg"] == "model input is near the context budget" {
			found = true
			if rec["prompt_tokens"] != float64(950) || rec["budget"] != float64(1000) {
				t.Fatalf("告警字段 = %+v, want prompt_tokens=950 / budget=1000", rec)
			}
		}
	}
	if !found {
		t.Fatalf("日志里没有预算告警记录：%q", buf.String())
	}
}

// captureLoopLog 返回输出 JSON 记录的 logger 与缓冲：断言结构化字段（而不是文本
// 子串）时用它。nil logger 在循环里会被换成丢弃 logger，所以测试必须显式注入。
func captureLoopLog() (*slog.Logger, *bytes.Buffer) {
	buf := &bytes.Buffer{}
	return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug})), buf
}

// loopJSONRecords 把缓冲里的 JSON 记录逐条解出来，便于按字段断言。
func loopJSONRecords(t *testing.T, buf *bytes.Buffer) []map[string]any {
	t.Helper()
	out := []map[string]any{}
	for _, line := range strings.Split(strings.TrimSpace(buf.String()), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		rec := map[string]any{}
		if err := json.Unmarshal([]byte(line), &rec); err != nil {
			t.Fatalf("日志行不是合法 JSON: %v（%q）", err, line)
		}
		out = append(out, rec)
	}
	return out
}