Cyrene_v2/internal/runtime/multimodal/vision/cache_test.go

package vision

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// pngBytesAlt 构造一张与 pngBytes 字节不同的合法 PNG，用于「同一 URL
// 内容变了」的用例。
func pngBytesAlt(t *testing.T) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
	img.Set(0, 0, color.RGBA{R: 0xff, G: 0xee, B: 0xdd, A: 0xff})
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatalf("encode alt png: %v", err)
	}
	return buf.Bytes()
}

// newMainModelResolver 构造一个把图像路由到主模型的解析器（主模型声明
// 支持视觉），工具路径被调用即视为测试失败。
func newMainModelResolver() *MessageResolver {
	return NewMessageResolver(NewRouter(WithVisionTool("vision")), func() ModelInfo {
		return ModelInfo{Name: "vision-model", SupportsVision: true}
	}, func(_ context.Context, _ string, _ map[string]any) (any, error) {
		return nil, ErrNotConfigured
	})
}

// resolveMainImage 走一遍解析器的主模型路径，返回内联后的图像串。
func resolveMainImage(res *MessageResolver, src string) (string, error) {
	msg, err := res.Resolve(context.Background(), model.ImageMessage(model.RoleUser, "x", src))
	if err != nil {
		return "", err
	}
	for _, p := range msg.Content {
		if p.Type == model.ContentTypeImage {
			return p.ImageURL, nil
		}
	}
	return "", errors.New("解析结果里没有图像部件")
}

// TestImageCache_SecondResolveFetchesOnce 是本缺陷的主用例：同一个 URL
// 连续解析两次（两个回合），上游只被取回一次，两次结果逐字节相同。
func TestImageCache_SecondResolveFetchesOnce(t *testing.T) {
	t.Parallel()
	png := pngBytes(t)
	var hits atomic.Int32
	imgSrv := newImageServer(t, png, "image/png", &hits)
	res := newMainModelResolver()
	src := imgSrv.URL + "/cat.png"

	first, err := resolveMainImage(res, src)
	if err != nil {
		t.Fatalf("第一次解析: %v", err)
	}
	second, err := resolveMainImage(res, src)
	if err != nil {
		t.Fatalf("第二次解析: %v", err)
	}
	want := DataURI("image/png", png)
	if first != want || second != want {
		t.Fatalf("内联结果 = (%q, %q)，want %q", first, second, want)
	}
	if got := hits.Load(); got != 1 {
		t.Fatalf("上游被取回 %d 次，want 1（第二个回合必须命中缓存）", got)
	}
	if sources, _, _ := res.cache.stats(); sources != 1 {
		t.Fatalf("缓存来源条数 = %d，want 1", sources)
	}
}

// TestImageCache_RefetchesAfterTTLWhenContentChanges 钉住失效策略：TTL 内
// 返回缓存（旧内容，这是缓存换取的确定性代价），TTL 到期后重新取回，上游
// 换了内容也能自愈 —— 新字节是新内容键，不会被旧条目命中。
func TestImageCache_RefetchesAfterTTLWhenContentChanges(t *testing.T) {
	t.Parallel()
	original := pngBytes(t)
	changed := pngBytesAlt(t)
	var mu sync.Mutex
	payload := original
	var hits atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		hits.Add(1)
		mu.Lock()
		body := payload
		mu.Unlock()
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(body)
	}))
	t.Cleanup(srv.Close)

	res := newMainModelResolver()
	res.cache.ttl = 20 * time.Millisecond
	src := srv.URL + "/cat.png"

	first, err := resolveMainImage(res, src)
	if err != nil {
		t.Fatalf("第一次解析: %v", err)
	}
	if first != DataURI("image/png", original) {
		t.Fatalf("第一次内联结果不对: %q", first)
	}

	mu.Lock()
	payload = changed
	mu.Unlock()

	// TTL 内：仍返回旧图，且没有再取回。
	if got, err := resolveMainImage(res, src); err != nil || got != first {
		t.Fatalf("TTL 内 = (%q, %v)，want 旧缓存 %q", got, err, first)
	}
	if got := hits.Load(); got != 1 {
		t.Fatalf("TTL 内上游被取回 %d 次，want 1", got)
	}

	time.Sleep(60 * time.Millisecond)

	third, err := resolveMainImage(res, src)
	if err != nil {
		t.Fatalf("TTL 到期后解析: %v", err)
	}
	wantChanged := DataURI("image/png", changed)
	if third != wantChanged {
		t.Fatalf("TTL 到期后 = %q，want 新内容 %q（旧映射没有失效）", third, wantChanged)
	}
	if got := hits.Load(); got != 2 {
		t.Fatalf("上游被取回 %d 次，want 2（TTL 到期必须重取）", got)
	}

	// 刷新后的映射必须指向新内容本身：若回填只换了来源映射、内容条目仍是
	// 旧 URI，下一次命中就会把旧图重新喂给模型。
	if got, err := resolveMainImage(res, src); err != nil || got != wantChanged {
		t.Fatalf("刷新后的缓存命中 = (%q, %v)，want %q（回填没有更新内容）", got, err, wantChanged)
	}
	if got := hits.Load(); got != 2 {
		t.Fatalf("刷新后的命中又取回了上游（%d 次），want 2", got)
	}
}

// TestImageCache_FailureIsNotCachedAndDiagnosable 覆盖失败路径：取回失败
// 不写缓存、第二次仍会重试，错误是结构化的 SourceFailure —— 只带来源主机
// 与原因，不带 URL 路径、查询串与凭据，且宿主能按结构断言取到主机（loop
// 的 WARN 因此可诊断）。
func TestImageCache_FailureIsNotCachedAndDiagnosable(t *testing.T) {
	t.Parallel()
	png := pngBytes(t)
	var hits atomic.Int32
	var failing atomic.Bool
	failing.Store(true)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		hits.Add(1)
		if failing.Load() {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(png)
	}))
	t.Cleanup(srv.Close)

	wantHost, err := url.Parse(srv.URL)
	if err != nil {
		t.Fatalf("parse test server url: %v", err)
	}
	res := newMainModelResolver()
	// 带查询串的地址：错误文案里绝不能出现凭据或路径。
	src := srv.URL + "/cat.png?token=SEKRET"

	_, err = resolveMainImage(res, src)
	if !errors.Is(err, ErrFetch) {
		t.Fatalf("第一次错误 = %v，want ErrFetch", err)
	}
	var sf *SourceFailure
	if !errors.As(err, &sf) {
		t.Fatalf("错误不是 *SourceFailure，调用方拿不到来源主机: %v", err)
	}
	if sf.Host != wantHost.Host {
		t.Fatalf("Host = %q，want %q", sf.Host, wantHost.Host)
	}
	var withHost interface{ ImageSourceHost() string }
	if !errors.As(err, &withHost) || withHost.ImageSourceHost() != wantHost.Host {
		t.Fatalf("错误未按结构断言暴露来源主机（loop 的 WARN 会丢 host）: %v", err)
	}
	msg := err.Error()
	for _, leak := range []string{"SEKRET", "token", "cat.png"} {
		if strings.Contains(msg, leak) {
			t.Fatalf("错误文案泄漏了 URL（%q）: %v", leak, err)
		}
	}
	if !strings.Contains(msg, wantHost.Host) {
		t.Fatalf("错误文案没有来源主机: %v", err)
	}
	if sources, contents, bytes := res.cache.stats(); sources != 0 || contents != 0 || bytes != 0 {
		t.Fatalf("失败结果被写进了缓存: sources=%d contents=%d bytes=%d", sources, contents, bytes)
	}

	failing.Store(false)
	got, err := resolveMainImage(res, src)
	if err != nil {
		t.Fatalf("第二次解析（来源已恢复）: %v", err)
	}
	if want := DataURI("image/png", png); got != want {
		t.Fatalf("第二次解析 = %q，want %q", got, want)
	}
	if n := hits.Load(); n != 2 {
		t.Fatalf("上游被取回 %d 次，want 2（失败必须可重试）", n)
	}
}

// TestImageCache_NotImageIsNotCached 是既有断言（非图像 → ErrNotImage）
// 在缓存路径上的延伸：HTML 错误页报错，且不会被缓存住。
func TestImageCache_NotImageIsNotCached(t *testing.T) {
	t.Parallel()
	var hits atomic.Int32
	imgSrv := newImageServer(t, htmlBytes(), "image/png", &hits)
	res := newMainModelResolver()
	src := imgSrv.URL + "/err"

	for attempt := 1; attempt <= 2; attempt++ {
		_, err := resolveMainImage(res, src)
		if !errors.Is(err, ErrNotImage) {
			t.Fatalf("第 %d 次错误 = %v，want ErrNotImage", attempt, err)
		}
	}
	if got := hits.Load(); got != 2 {
		t.Fatalf("上游被取回 %d 次，want 2（非图像不得进缓存）", got)
	}
	if sources, contents, _ := res.cache.stats(); sources != 0 || contents != 0 {
		t.Fatalf("非图像被写进了缓存: sources=%d contents=%d", sources, contents)
	}
}

// TestImageCache_BoundsAreEnforced 钉住边界：来源条数与内容总字节都不会
// 无界增长，最老的条目按 LRU 淘汰（淘汰后不再命中）。
func TestImageCache_BoundsAreEnforced(t *testing.T) {
	t.Parallel()
	c := newImageCache()
	c.maxSrc = 2
	c.maxCtn = 1000 // 条数上限放到很大，这里只让字节上限与来源上限起作用
	uri := strings.Repeat("A", 200)
	c.maxBytes = 3 * len(uri)

	for i := 0; i < 50; i++ {
		c.put(fmt.Sprintf("https://example.com/%d.png", i), "image/png", []byte{byte(i)}, uri)
	}
	sources, contents, bytes := c.stats()
	if sources > c.maxSrc {
		t.Fatalf("来源条数 = %d，超过上限 %d", sources, c.maxSrc)
	}
	if bytes > c.maxBytes {
		t.Fatalf("缓存字节 = %d，超过上限 %d", bytes, c.maxBytes)
	}
	if contents > c.maxCtn {
		t.Fatalf("内容条目数 = %d，超过条数上限 %d", contents, c.maxCtn)
	}
	if maxByBytes := c.maxBytes / len(uri); contents > maxByBytes {
		t.Fatalf("内容条目数 = %d，超过按字节上限换算的 %d", contents, maxByBytes)
	}
	if _, ok := c.get("https://example.com/0.png"); ok {
		t.Fatal("最老的来源仍在缓存里（LRU 淘汰没有生效）")
	}
	if _, ok := c.get("https://example.com/49.png"); !ok {
		t.Fatal("最新的来源不在缓存里")
	}
}

// TestImageCache_ContentCountCapEnforced 单独钉住内容条目数上限：字节上限
// 放得很大时，只有条数上限起作用 —— 极小的图片堆不出极多的条目，字节数
// 也随之收敛，最老的内容被 LRU 淘汰。
func TestImageCache_ContentCountCapEnforced(t *testing.T) {
	t.Parallel()
	c := newImageCache()
	c.maxSrc = 1000
	c.maxCtn = 4
	c.maxBytes = 1 << 30
	uri := strings.Repeat("A", 200)

	for i := 0; i < 50; i++ {
		c.put(fmt.Sprintf("https://example.com/%d.png", i), "image/png", []byte{byte(i)}, uri)
	}
	sources, contents, bytes := c.stats()
	if contents > c.maxCtn {
		t.Fatalf("内容条目数 = %d，超过条数上限 %d", contents, c.maxCtn)
	}
	if want := c.maxCtn * len(uri); bytes != want {
		t.Fatalf("缓存字节 = %d，want %d（字节数应随条数上限收敛）", bytes, want)
	}
	if sources > c.maxSrc {
		t.Fatalf("来源条数 = %d，超过上限 %d", sources, c.maxSrc)
	}
	if _, ok := c.get("https://example.com/0.png"); ok {
		t.Fatal("最老的内容仍在缓存里（LRU 淘汰没有生效）")
	}
	if _, ok := c.get("https://example.com/49.png"); !ok {
		t.Fatal("最新的内容不在缓存里")
	}
}

// TestImageCache_ConcurrentResolveIsSafe 覆盖并发解析同一 URL：结构状态
// 不被破坏（配合 -race），结果一致；并发风暴之后的解析必定命中缓存。
func TestImageCache_ConcurrentResolveIsSafe(t *testing.T) {
	t.Parallel()
	png := pngBytes(t)
	var hits atomic.Int32
	imgSrv := newImageServer(t, png, "image/png", &hits)
	res := newMainModelResolver()
	src := imgSrv.URL + "/cat.png"
	want := DataURI("image/png", png)

	const workers = 8
	results := make([]string, workers)
	errs := make([]error, workers)
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			results[i], errs[i] = resolveMainImage(res, src)
		}(i)
	}
	wg.Wait()

	for i := 0; i < workers; i++ {
		if errs[i] != nil {
			t.Fatalf("并发解析 %d: %v", i, errs[i])
		}
		if results[i] != want {
			t.Fatalf("并发解析 %d = %q，want %q", i, results[i], want)
		}
	}
	if got := hits.Load(); got == 0 || got > workers {
		t.Fatalf("上游被取回 %d 次，want 1..%d", got, workers)
	}

	before := hits.Load()
	if got, err := resolveMainImage(res, src); err != nil || got != want {
		t.Fatalf("并发风暴后 = (%q, %v)，want %q", got, err, want)
	}
	if after := hits.Load(); after != before {
		t.Fatalf("并发风暴后的解析仍取回了上游（%d → %d），缓存没生效", before, after)
	}
	sources, _, _ := res.cache.stats()
	if sources != 1 {
		t.Fatalf("缓存来源条数 = %d，want 1（同一 URL 只应有一条映射）", sources)
	}
}