Cyrene_v2/internal/runtime/embedding/budget_test.go

package embedding

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// 本文件把 budget.go 的**不变式**钉成可执行的断言，两态各一条：
//
//   - 「批内耗时 < 超时」⇒ 不降级：DerivedBatchCompletesWithinTimeout；
//   - 「批内耗时 ≥ 超时」⇒ 超时（本次故障形态）：LegacyBatchOf64TimesOut；
//
// 以及**变异自证**：把批上限调回旧值 64（或把推演调松）会让哪条测试红，
// 见 TestBatchBudget_InvariantHolds 与 TestAPIEmbedder_LegacyBatchOf64TimesOut。
//
// 假客户端用的是**实测曲线**（不是随口编的常数）：单次请求耗时
// ≈ measuredOverhead + 字符总量 × measuredPerChar，取值见下面两个常量
// （2026-10-13 本机实测：固定开销 ~50ms、每字符 2.7~3.9ms，此处取上界）。
const (
	measuredOverhead  = 50 * time.Millisecond
	measuredPerChar   = 4 * time.Millisecond
	latencyModelChars = 33 // 活库短正文的平均长度（实测 33 汉字）
)

// latencyTripper 是一个按上面那条曲线**真的睡够时间**的 RoundTripper：
// 它模拟「服务端会算完，不管你客户端还在不在等」。这正是本次故障的形态 ——
// 客户端 5s 放弃、服务端还在算，后面每个请求都排在它后面。
type latencyTripper struct {
	mu       sync.Mutex
	requests int
	maxItems int
	maxChars int
	// abandoned 计数「客户端先放弃、服务端还在算」的次数 —— 本次故障的形态。
	// 修好之后它必须是 0（见 TestAPIEmbedder_OversizedBatchIsSplit 的断言）。
	abandoned int
	// firstSleep 非零时，第一次请求额外多睡这么久（模拟冷启动的模型加载）。
	firstSleep time.Duration
}

func (l *latencyTripper) RoundTrip(r *http.Request) (*http.Response, error) {
	body, err := io.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	var req openAICompatibleRequest
	if err := json.Unmarshal(body, &req); err != nil {
		return nil, err
	}
	chars := 0
	for _, s := range req.Input {
		chars += len([]rune(s))
	}

	l.mu.Lock()
	l.requests++
	if len(req.Input) > l.maxItems {
		l.maxItems = len(req.Input)
	}
	if chars > l.maxChars {
		l.maxChars = chars
	}
	first := l.requests == 1
	extra := l.firstSleep
	l.mu.Unlock()

	wait := measuredOverhead + time.Duration(chars)*measuredPerChar
	if first && extra > 0 {
		wait += extra
	}
	select {
	case <-time.After(wait):
	case <-r.Context().Done():
		// 客户端放弃了：服务端那一侧不会因为客户端走了就把已排队的算力还回来
		// （实测的级联堆积正是从这一步开始的），因此这里明确记一笔。
		l.mu.Lock()
		l.abandoned++
		l.mu.Unlock()
		return nil, r.Context().Err()
	}

	resp := openAICompatibleResponse{Data: make([]openAICompatibleDatum, 0, len(req.Input))}
	for i := range req.Input {
		resp.Data = append(resp.Data, openAICompatibleDatum{Index: i, Embedding: []float32{1, 0}})
	}
	raw, err := json.Marshal(resp)
	if err != nil {
		return nil, err
	}
	return &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(string(raw))),
		Header:     make(http.Header),
	}, nil
}

// shortTexts 造 n 条与活库短正文等长的文本。
func shortTexts(n int) []string {
	out := make([]string, n)
	line := strings.Repeat("河", latencyModelChars)
	for i := range out {
		out[i] = line
	}
	return out
}

// TestBatchBudget_InvariantHolds 是**推演本身**的断言（不连任何网络）：
// 由超时推出来的一对上限，必须使「一次请求的耗时上界 ≤ 超时的一半」成立。
//
// 变异自证：把 DefaultBatchLimit 的推演改成旧口径（条数写回 64）⇒ 第一条断言红；
// 若改成「把超时放大到盖住 64 条」⇒ 最后一条断言红（旧配对必须仍是越界的）。
func TestBatchBudget_InvariantHolds(t *testing.T) {
	t.Parallel()

	if DefaultBatchLimit < 1 || DefaultBatchChars < 1 {
		t.Fatalf("推出来的上限必须 ≥ 1：条数=%d 字符=%d", DefaultBatchLimit, DefaultBatchChars)
	}
	share := DefaultRequestTimeout / callBudgetShare
	byItems := requestOverhead + time.Duration(DefaultBatchLimit)*perItemBudget
	if byItems > share {
		t.Fatalf("条数上限 %d 的耗时上界 %s 超过超时的一半 %s（推演被调松了）",
			DefaultBatchLimit, byItems, share)
	}
	byChars := requestOverhead + time.Duration(DefaultBatchChars)*perCharBudget
	if byChars > share {
		t.Fatalf("字符上限 %d 的耗时上界 %s 超过超时的一半 %s（推演被调松了）",
			DefaultBatchChars, byChars, share)
	}
	// 旧配对（批次 64 + 5s 超时）必须仍被判为越界 —— 那正是 2026-10-13 的故障形态：
	// 实测 64 条 7.5~8.5s > 5s。这里用同一套预算常量复核它仍然出界。
	legacy := requestOverhead + 64*perItemBudget
	if legacy <= DefaultRequestTimeout {
		t.Fatalf("旧配对（64 条）的耗时上界 %s 竟然落在单次超时 %s 之内 —— "+
			"说明超时被放大到迁就批次了，本次修法不允许", legacy, DefaultRequestTimeout)
	}
}

// TestChunkTexts_RespectsBothLimits 钉住切片：条数与字符总量两道界都要守，
// 单条本身越界的文本独占一段且**不被截断**。
func TestChunkTexts_RespectsBothLimits(t *testing.T) {
	t.Parallel()

	got := chunkTexts(shortTexts(9), 2, 100)
	if len(got) != 5 {
		t.Fatalf("2 条/段 切 9 条 = %d 段，want 5", len(got))
	}
	for i, seg := range got {
		if len(seg) > 2 {
			t.Fatalf("第 %d 段有 %d 条，超过条数上限 2", i, len(seg))
		}
	}

	// 字符界绑定：4 条 33 字符 = 132 字符 ⇒ 100 字符/段时必须收口（3 条 = 99 字符满，
	// 余下 1 条自成一段）。
	byChars := chunkTexts(shortTexts(4), 64, 100)
	if len(byChars) != 2 {
		t.Fatalf("100 字符/段 切 4×33 字符 = %d 段，want 2", len(byChars))
	}
	for i, seg := range byChars {
		total := 0
		for _, s := range seg {
			total += len([]rune(s))
		}
		if total > 100 {
			t.Fatalf("第 %d 段 %d 字符，超过字符上限 100", i, total)
		}
	}

	// 超长单条：不切碎、不截断，独占一段。
	huge := strings.Repeat("长", 500)
	segs := chunkTexts([]string{huge, "短"}, 64, 100)
	if len(segs) != 2 || segs[0][0] != huge {
		t.Fatalf("超长单条应独占一段且不改写，got %d 段", len(segs))
	}

	if chunkTexts(nil, 4, 100) != nil {
		t.Fatal("空输入应返回 nil")
	}
}

// TestAPIEmbedder_DerivedBatchCompletesWithinTimeout 钉**不降级态**：
// 按推出来的上限发一次（Embed 内部切成 1 段），在实测曲线下必须在单次超时内拿到向量。
func TestAPIEmbedder_DerivedBatchCompletesWithinTimeout(t *testing.T) {
	t.Parallel()

	rt := &latencyTripper{}
	emb := NewAPI(config.EmbeddingConfig{APIURL: "http://127.0.0.1:11434/v1/embeddings"},
		WithRoundTripper(rt))

	texts := shortTexts(DefaultBatchLimit)
	vectors, err := emb.Embed(context.Background(), texts)
	if err != nil {
		t.Fatalf("上限内的批不该超时：%v", err)
	}
	if len(vectors) != len(texts) {
		t.Fatalf("got %d vectors, want %d", len(vectors), len(texts))
	}
	if rt.requests != 1 {
		t.Fatalf("上限内的批应发 1 次请求，实际 %d 次", rt.requests)
	}
}

// TestAPIEmbedder_OversizedBatchIsSplit 钉**不变式与调用方无关**：
// 调用方给 3 倍上限的切片时，Embed 自己切成多段、每段都在界内，向量顺序与入参一致。
func TestAPIEmbedder_OversizedBatchIsSplit(t *testing.T) {
	t.Parallel()

	rt := &latencyTripper{}
	emb := NewAPI(config.EmbeddingConfig{APIURL: "http://127.0.0.1:11434/v1/embeddings"},
		WithRoundTripper(rt))

	texts := shortTexts(DefaultBatchLimit * 3)
	vectors, err := emb.Embed(context.Background(), texts)
	if err != nil {
		t.Fatalf("超大切片应被切开而不是超时：%v", err)
	}
	if len(vectors) != len(texts) {
		t.Fatalf("got %d vectors, want %d（切开后必须按序拼回）", len(vectors), len(texts))
	}
	if rt.requests < 3 {
		t.Fatalf("3 倍上限应至少切成 3 段，实际 %d 段", rt.requests)
	}
	if rt.maxItems > DefaultBatchLimit {
		t.Fatalf("有请求带了 %d 条，超过推出来的条数上限 %d", rt.maxItems, DefaultBatchLimit)
	}
	if rt.maxChars > DefaultBatchChars {
		t.Fatalf("有请求带了 %d 字符，超过推出来的字符上限 %d", rt.maxChars, DefaultBatchChars)
	}
	// 🔴 本次修复的核心不变式，确定性断言（不依赖机器负载）：
	// 没有任何一次请求发生「客户端先放弃、服务端还在算」——那正是级联堆积的起点。
	if rt.abandoned != 0 {
		t.Fatalf("有 %d 次请求被客户端先放弃（服务端还在算）⇒ 不变式没守住", rt.abandoned)
	}
	// 服务端侧的同一条不变式：任何一次请求的建模耗时（实测曲线）都落在单次超时内。
	if cost := measuredOverhead + time.Duration(rt.maxChars)*measuredPerChar; cost > DefaultRequestTimeout {
		t.Fatalf("最大一次请求的实测曲线耗时 %s 越过单次超时 %s", cost, DefaultRequestTimeout)
	}
}

// TestAPIEmbedder_LegacyBatchOf64TimesOut 是**变异自证**的另一半：同样是实测曲线，
// 旧配对（一次请求 64 条 + 5s 超时）必须**超时**。它证明把上限写回 64 不是「保守」，而是
// 必然把客户端与服务端撕裂（客户端走了、服务端还在算 ⇒ 后续请求全排队）。
func TestAPIEmbedder_LegacyBatchOf64TimesOut(t *testing.T) {
	t.Parallel()

	rt := &latencyTripper{}
	emb := NewAPI(config.EmbeddingConfig{APIURL: "http://127.0.0.1:11434/v1/embeddings"},
		WithRoundTripper(rt))

	// 先把一次性冷启动放宽花掉，否则第一次调用按 warmupTimeout 算、不会超时。
	if _, err := emb.Embed(context.Background(), shortTexts(1)); err != nil {
		t.Fatalf("预热调用失败：%v", err)
	}
	start := time.Now()
	_, err := emb.embedChunk(context.Background(), shortTexts(64))
	elapsed := time.Since(start)
	if err == nil {
		t.Fatalf("旧配对（64 条 + %s 超时）在实测曲线下不该成功", DefaultRequestTimeout)
	}
	if elapsed > DefaultRequestTimeout+2*time.Second {
		t.Fatalf("超时保护没生效：等了 %s", elapsed)
	}
}

// TestAPIEmbedder_WarmupIsOneShot 钉 p1 的核心：冷启动放宽**一次性**，三种失败形态
// 下都必须是「第一次最多等 warmupTimeout，第二次起一律 DefaultRequestTimeout」。
//
// 这是 reviewer 在 round1 打回的那条：旧实现 `warmed` 只在 2xx 之后才翻，于是
// 「服务接受连接但永不回包」时**每一次**调用都等完整放宽预算（30s），把 fail-fast
// 变成 fail-slow —— 一次挂死就能吃光审查的 20s 预算、把索引写入拖成 8×30s。
func TestAPIEmbedder_WarmupIsOneShot(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name string
		rt   func(t *testing.T) http.RoundTripper
	}{
		{
			name: "连接被拒",
			rt: func(t *testing.T) http.RoundTripper {
				return roundTripFunc(func(*http.Request) (*http.Response, error) {
					return nil, errors.New("dial tcp 127.0.0.1:11434: connectex: 目标计算机积极拒绝")
				})
			},
		},
		{
			name: "HTTP 500",
			rt: func(t *testing.T) http.RoundTripper {
				return roundTripFunc(func(*http.Request) (*http.Response, error) {
					return &http.Response{
						StatusCode: http.StatusInternalServerError,
						Body:       io.NopCloser(strings.NewReader("server exploded")),
						Header:     make(http.Header),
					}, nil
				})
			},
		},
		{
			name: "接受连接但永不回包",
			rt: func(t *testing.T) http.RoundTripper {
				return roundTripFunc(func(r *http.Request) (*http.Response, error) {
					<-r.Context().Done() // 只等服务端超时，永不回包
					return nil, r.Context().Err()
				})
			},
		},
	}

	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			emb := NewAPI(config.EmbeddingConfig{APIURL: "http://127.0.0.1:11434/v1/embeddings"},
				WithRoundTripper(tc.rt(t)))

			// 第一次：放宽预算。断言上界 = warmupTimeout（不得再宽），
			// 下界只对会真正等到截止的那一种（挂死）有意义。
			start := time.Now()
			_, err := emb.Embed(context.Background(), shortTexts(1))
			first := time.Since(start)
			if err == nil {
				t.Fatal("这一组失败形态不该成功")
			}
			t.Logf("首次耗时 = %s", first.Round(time.Millisecond))
			if first > warmupTimeout+time.Second {
				t.Fatalf("首次耗时 %s 超过冷启动放宽 %s", first, warmupTimeout)
			}

			// 第二次：**必须**回到常态超时 —— 一次挂死绝不许变成每次都慢。
			start = time.Now()
			if _, err := emb.Embed(context.Background(), shortTexts(1)); err == nil {
				t.Fatal("这一组失败形态不该成功")
			}
			second := time.Since(start)
			t.Logf("第二次耗时 = %s（常态超时 %s）", second.Round(time.Millisecond), DefaultRequestTimeout)
			if second > DefaultRequestTimeout+time.Second {
				t.Fatalf("第二次耗时 %s 超过常态超时 %s ⇒ 放宽成了长期的", second, DefaultRequestTimeout)
			}
			// 放宽许可必须已经被花掉（它是一次性的，任何路径都不许给它续期）。
			if emb.warmupLeft.Load() {
				t.Fatal("第二次请求之后放宽许可仍为真 —— 一次性的约束被绕过了")
			}
		})
	}
}

// TestAPIEmbedder_WarmupCoversColdStartOnce 钉放宽的**正面**作用：第一次请求
// （模型加载 / GPU 缓冲分配）可以比常态超时多等一点，只要它落在 warmupTimeout
// 之内；而这一刻过去之后，请求立刻回到常态超时（deadline 断言）。
func TestAPIEmbedder_WarmupCoversColdStartOnce(t *testing.T) {
	t.Parallel()

	rt := &latencyTripper{firstSleep: DefaultRequestTimeout + time.Second}
	emb := NewAPI(config.EmbeddingConfig{APIURL: "http://127.0.0.1:11434/v1/embeddings"},
		WithRoundTripper(rt))

	if _, err := emb.Embed(context.Background(), shortTexts(1)); err != nil {
		t.Fatalf("冷启动第一次调用应被 warmupTimeout(%s) 兜住：%v", warmupTimeout, err)
	}
	if emb.warmupLeft.Load() {
		t.Fatal("放宽许可应在第一次请求时就被花掉")
	}

	var slack time.Duration
	steady := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		dl, ok := r.Context().Deadline()
		if !ok {
			t.Error("稳态请求没有 deadline")
		} else {
			slack = time.Until(dl)
		}
		return &http.Response{
			StatusCode: http.StatusOK,
			Body:       io.NopCloser(strings.NewReader(`{"data":[{"index":0,"embedding":[1]}]}`)),
			Header:     make(http.Header),
		}, nil
	})
	emb.client = &http.Client{Transport: steady}
	if _, err := emb.Embed(context.Background(), shortTexts(1)); err != nil {
		t.Fatalf("稳态调用失败：%v", err)
	}
	if slack <= 0 || slack > DefaultRequestTimeout {
		t.Fatalf("稳态单次超时余量 = %s, want (0, %s]", slack, DefaultRequestTimeout)
	}
}