Cyrene_v2/internal/runtime/memory/retrieve_fallback_test.go

package memory

import (
	"bytes"
	"context"
	"log/slog"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

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

// 本文件是 2026-10-05 真机性能事故的回归测试：主人只发了一个**不带正文**的
// 附件（空 query），memory.recall 一次花了 39.3 秒 —— 索引因为空 query
// 没被问过，回退路径于是把整库 302 条候选正文逐条向量化。三条断言分别
// 钉住「空 query 不做任何语义工作」「回退不再把整库塞进一次请求且不重复
// 嵌入」「回退被预算拦住、不阻塞回合」。

// recordingEmbedder 记录每次 Embed 收到的文本条数，并返回确定性向量：
// 它是「回退路径到底嵌入了多少条正文」的观测点。
type recordingEmbedder struct {
	dim int

	mu      sync.Mutex
	batches []int
	texts   int
}

func (e *recordingEmbedder) Embed(_ context.Context, texts []string) ([]embedding.Vector, error) {
	e.mu.Lock()
	e.batches = append(e.batches, len(texts))
	e.texts += len(texts)
	e.mu.Unlock()

	out := make([]embedding.Vector, len(texts))
	for i := range texts {
		v := make(embedding.Vector, e.Dim())
		v[0] = 1
		out[i] = v
	}
	return out, nil
}

func (e *recordingEmbedder) Dim() int {
	if e.dim > 0 {
		return e.dim
	}
	return 4
}

// batchSizes 返回到目前为止每次 Embed 的文本条数（副本）。
func (e *recordingEmbedder) batchSizes() []int {
	e.mu.Lock()
	defer e.mu.Unlock()
	return append([]int(nil), e.batches...)
}

// blockingBatchEmbedder 模拟「远端 embedding 接口对一批文本迟迟不回」：
// 单条文本立即返回，一次多条则一直等到 ctx 结束 —— 用来验证回退路径
// 真的会收手，而不是陪着接口一起把回合拖住。
type blockingBatchEmbedder struct{ dim int }

func (e blockingBatchEmbedder) Embed(ctx context.Context, texts []string) ([]embedding.Vector, error) {
	if len(texts) > 1 {
		<-ctx.Done()
		return nil, ctx.Err()
	}
	out := make([]embedding.Vector, len(texts))
	for i := range texts {
		v := make(embedding.Vector, e.dim)
		v[0] = 1
		out[i] = v
	}
	return out, nil
}

func (e blockingBatchEmbedder) Dim() int { return e.dim }

// TestRetriever_EmptyQuery_NoSemanticWork_空查询不做任何语义工作 是事故的
// 主回归：空 query（附件类回合没有正文）既不该去问向量索引，也不该把
// 候选正文送去向量化。旧行为下这里会把每个候选正文都 embed 一遍、
// 真机 302 条即 39.3 秒。
func TestRetriever_EmptyQuery_NoSemanticWork_空查询不做任何语义工作(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	store := newTestStore(t)

	// 候选条数跨过回退的单批上限：一旦这条路径被触发，嵌入条数会明显
	// 大于「一次 query 向量」，断言因此不会因边界而假绿。
	candidates := fallbackEmbedBatch + 8
	for i := 0; i < candidates; i++ {
		addItem(t, store, MemoryItem{
			PersonaID: "persona.a",
			Content:   "记忆条目 " + strconv.Itoa(i),
			Scope:     ScopePersona,
			Kind:      KindWorldEvent,
		})
	}

	emb := &recordingEmbedder{dim: 8}
	// 索引恒「无候选命中」：这正是事故里那条回退路径的入口。
	idx := &countingIndex{Index: &mockVectorIndex{dim: 8}}
	r := NewRetriever(store, WithEmbedder(emb), WithVectorIndex(idx))

	for _, query := range []string{"", " ", "\n\t"} {
		results, err := r.Retrieve(ctx, "persona.a", query, 5)
		if err != nil {
			t.Fatalf("Retrieve(query=%q): %v", query, err)
		}
		if len(results) == 0 {
			t.Fatalf("query=%q：空查询仍应按新近度返回候选，got 0 results", query)
		}
	}
	if emb.texts != 0 {
		t.Fatalf("空 query 向量化了 %d 条文本，want 0（真机事故：302 条候选正文 = 39.3 秒）", emb.texts)
	}
	if idx.searches != 0 {
		t.Fatalf("空 query 查询了向量索引 %d 次，want 0（空 query 必然「无候选命中」，只会把正常检索误报成索引失效）",
			idx.searches)
	}
}

// TestRetriever_Fallback_BoundedAndCached_回退分批且不重复嵌候选 钉住回退
// 路径的两条护栏：候选正文分批送出（不再一个请求塞整库），且同一批
// 正文在第二回合不再重复向量化（embedder 对相同输入确定性）。
func TestRetriever_Fallback_BoundedAndCached_回退分批且不重复嵌候选(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	store := newTestStore(t)

	candidates := fallbackEmbedBatch*2 + 5
	for i := 0; i < candidates; i++ {
		content := "coffee note " + strconv.Itoa(i)
		addItem(t, store, MemoryItem{
			PersonaID: "persona.a", Content: content,
			Scope: ScopePersona, Kind: KindWorldEvent,
		})
	}

	emb := &recordingEmbedder{dim: 8}
	idx := &countingIndex{Index: &mockVectorIndex{dim: 8}}
	r := NewRetriever(store, WithEmbedder(emb), WithVectorIndex(idx))

	results, err := r.Retrieve(ctx, "persona.a", "coffee", 5)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 || results[0].Score <= 0 {
		t.Fatalf("回退没有给出任何语义分：results=%d score=%v", len(results), results[0].Score)
	}

	batches := emb.batchSizes()
	if len(batches) < 2 {
		t.Fatalf("候选正文没有分批：一次 Embed 带了 %v 条，want 多次且每次 ≤ %d", batches, fallbackEmbedBatch)
	}
	for _, n := range batches {
		if n > fallbackEmbedBatch {
			t.Fatalf("单次 Embed 携带 %d 条文本，超过上限 %d（回退不得把整库塞进一个请求）",
				n, fallbackEmbedBatch)
		}
	}
	if got, want := emb.texts, candidates+1; got != want {
		t.Fatalf("首次回退向量化了 %d 条文本，want %d（%d 条候选分批 + 1 条 query）", got, want, candidates)
	}

	// 第二次检索：候选正文全部命中缓存，因此只多一次 query 向量化。
	if _, err := r.Retrieve(ctx, "persona.a", "coffee", 5); err != nil {
		t.Fatalf("second Retrieve: %v", err)
	}
	if got, want := emb.texts, candidates+2; got != want {
		t.Fatalf("第二次回退向量化了 %d 条文本（累计），want %d —— 候选正文被重复嵌入，缓存没生效",
			got, want)
	}
}

// TestRetriever_Fallback_Budget_超预算不阻塞回合 钉住预算护栏：embedding
// 接口对一批文本迟迟不回时，回退用手上已有的向量收手，回合不被拖住，
// 并且留下一条带来源的 Warn（语义分只补了一部分必须可见）。
func TestRetriever_Fallback_Budget_超预算不阻塞回合(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	store := newTestStore(t)

	for i := 0; i < fallbackEmbedBatch+1; i++ {
		addItem(t, store, MemoryItem{
			PersonaID: "persona.a", Content: "记忆条目 " + strconv.Itoa(i),
			Scope: ScopePersona, Kind: KindWorldEvent,
		})
	}

	var buf bytes.Buffer
	logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))
	idx := &countingIndex{Index: &mockVectorIndex{dim: 8}}
	r := NewRetriever(store,
		WithEmbedder(blockingBatchEmbedder{dim: 8}),
		WithVectorIndex(idx),
		WithRetrieverLogger(logger),
		WithFallbackBudget(50*time.Millisecond),
	)

	start := time.Now()
	results, err := r.Retrieve(ctx, "persona.a", "咖啡", 5)
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("回退超预算不应丢掉检索结果（关键词与新近度仍在）")
	}
	if elapsed > 5*time.Second {
		t.Fatalf("回退没有被预算拦住：检索耗时 %v", elapsed)
	}
	if !strings.Contains(buf.String(), "exceeded its embed budget") {
		t.Fatalf("回退超预算没有告警，语义分只补了一部分却无人可见；日志=%q", buf.String())
	}
}