Merge pull request #8025 from esengine/feat/retrieval-v2-shadow

feat(retrieval): Retrieval V2 shadow — BM25F, code symbols, mixed CJK grams

benchmarks/memorybench/tasks/mb-pin/task.toml

@@ -3,3 +3,4 @@ class = "pinned"
 max_steps = 8
 timeout_sec = 180
 memory_markers = ["MEMKEY-PIN4D"]
+memory_markers_prefix = true

cmd/e2ebench/main.go

@@ -34,7 +34,10 @@ type task struct {
 	// MemoryMarkers are unique tokens planted in seeded fact bodies; a marker
 	// found in tool args or answer text after a recall proves point of use.
 	MemoryMarkers []string `toml:"memory_markers" json:"memory_markers,omitempty"`
-	dir           string
+	// MemoryMarkersPrefix marks tasks whose seeded facts are pinned: their
+	// bodies arrive via the stable prefix, so markers count from turn one.
+	MemoryMarkersPrefix bool `toml:"memory_markers_prefix" json:"memory_markers_prefix,omitempty"`
+	dir                 string
 }
 
 type runMetrics struct {
@@ -104,6 +107,7 @@ type result struct {
 	MemoryRecallChars  int `json:"memory_recall_chars,omitempty"`
 	MemorySuppressed   int `json:"memory_suppressed,omitempty"`
 	MemoryMarkersUsed  int `json:"memory_markers_used,omitempty"`
+	MemoryShadowAgree  int `json:"memory_shadow_agree,omitempty"`
 	// WallMs is the harness's own clock, not the agent's self-report, so the
 	// number stays comparable when the same suite runs against another harness.
 	WallMs int64 `json:"wall_ms"`
@@ -552,7 +556,7 @@ func runTask(cfg suiteConfig, t task) result {
 		if summary, err := summarizeTrajectory(trajPath); err == nil {
 			r.Trajectory = summary
 		}
-		applyMemoryStats(&r, trajPath, t.MemoryMarkers)
+		applyMemoryStats(&r, trajPath, t)
 	}
 	// A killed child never writes metrics, so the deadline is the only place
 	// this failure mode is still observable.

cmd/e2ebench/memorybench.go

@@ -80,11 +80,11 @@ func taskExperimentEnv(cfg suiteConfig, t task, work string) (env []string, note
 }
 
 // applyMemoryStats folds one trajectory's recall behavior into the result row.
-func applyMemoryStats(r *result, trajPath string, markers []string) {
-	stats := scanMemoryRecall(trajPath, markers)
+func applyMemoryStats(r *result, trajPath string, t task) {
+	stats := scanMemoryRecall(trajPath, t.MemoryMarkers, t.MemoryMarkersPrefix)
 	r.MemoryRecallEvents, r.MemoryRecallHits = stats.RecallEvents, stats.RecallHits
 	r.MemoryRecallChars, r.MemorySuppressed = stats.RecallChars, stats.Suppressed
-	r.MemoryMarkersUsed = stats.MarkersUsed
+	r.MemoryMarkersUsed, r.MemoryShadowAgree = stats.MarkersUsed, stats.ShadowAgree
 }
 
 // memoryRunStats is what one trajectory reveals about recall behavior.
@@ -94,13 +94,14 @@ type memoryRunStats struct {
 	RecallChars  int
 	Suppressed   int // recall decisions that stayed silent
 	MarkersUsed  int // task markers seen in tool args or answer text after recall
+	ShadowAgree  int // recall events where the V2 shadow's top hit matched production's
 }
 
 // scanMemoryRecall extracts recall decisions and point-of-use evidence: a
 // marker (a unique token planted in a seeded fact body) counts as used only
 // when it appears in tool arguments or answer text AFTER a recall injected
 // facts — the fact reached the decision path, not just the ranking.
-func scanMemoryRecall(path string, markers []string) memoryRunStats {
+func scanMemoryRecall(path string, markers []string, markersInPrefix bool) memoryRunStats {
 	var stats memoryRunStats
 	f, err := os.Open(path)
 	if err != nil {
@@ -112,6 +113,7 @@ func scanMemoryRecall(path string, markers []string) memoryRunStats {
 			Hits       []struct{ ID string } `json:"hits"`
 			UsedChars  int                   `json:"used_chars"`
 			Suppressed string                `json:"suppressed"`
+			ShadowHits []struct{ ID string } `json:"shadow_hits"`
 		} `json:"memory_recall"`
 		Event *struct {
 			Kind string `json:"kind"`
@@ -122,7 +124,9 @@ func scanMemoryRecall(path string, markers []string) memoryRunStats {
 		} `json:"event"`
 	}
 	used := make(map[string]bool, len(markers))
-	recalled := false
+	// Pinned facts arrive via the stable prefix, before any recall: their
+	// markers count from the first record.
+	recalled := markersInPrefix
 	scanner := bufio.NewScanner(f)
 	scanner.Buffer(make([]byte, 0, 1024*1024), 64*1024*1024)
 	for scanner.Scan() {
@@ -136,6 +140,9 @@ func scanMemoryRecall(path string, markers []string) memoryRunStats {
 				stats.RecallHits += len(mr.Hits)
 				stats.RecallChars += mr.UsedChars
 				recalled = true
+				if len(mr.ShadowHits) > 0 && mr.ShadowHits[0].ID == mr.Hits[0].ID {
+					stats.ShadowAgree++
+				}
 			} else if mr.Suppressed != "" {
 				stats.Suppressed++
 			}
@@ -180,8 +187,15 @@ func renderMemoryShadow(results []result) string {
 	if hits == 0 && markersTotal == 0 {
 		return ""
 	}
+	shadowAgree := 0
+	for _, r := range results {
+		shadowAgree += r.MemoryShadowAgree
+	}
 	line := fmt.Sprintf("**Memory shadow** (%d runs): **recall fired** in %d runs · **hits** %d · **injected chars** %d",
 		runs, recallRuns, hits, chars)
+	if recallRuns > 0 {
+		line += fmt.Sprintf(" · **V2 top1 agree** %d/%d", shadowAgree, recallRuns)
+	}
 	if markersTotal > 0 {
 		line += fmt.Sprintf(" · **point-of-use** %d/%d markers", markersUsed, markersTotal)
 	}

cmd/e2ebench/memorybench_test.go

@@ -21,7 +21,7 @@ func TestScanMemoryRecallCountsAndPointOfUse(t *testing.T) {
 	if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")+"\n"), 0o644); err != nil {
 		t.Fatal(err)
 	}
-	stats := scanMemoryRecall(path, []string{"MEMKEY-USED", "MEMKEY-TEXT", "MEMKEY-EARLY", "MEMKEY-NEVER"})
+	stats := scanMemoryRecall(path, []string{"MEMKEY-USED", "MEMKEY-TEXT", "MEMKEY-EARLY", "MEMKEY-NEVER"}, false)
 	if stats.RecallEvents != 1 || stats.RecallHits != 2 || stats.RecallChars != 420 || stats.Suppressed != 1 {
 		t.Fatalf("stats = %+v, want 1 event / 2 hits / 420 chars / 1 suppressed", stats)
 	}

internal/control/memory.go

@@ -124,6 +124,9 @@ func memoryRecallAudit(result memory.RecallResult) event.MemoryRecallAudit {
 			Freshness: hit.Freshness, Score: hit.Score,
 		})
 	}
+	for _, hit := range result.ShadowHits {
+		audit.Shadow = append(audit.Shadow, event.MemoryRecallHit{ID: hit.ID, Score: hit.Score})
+	}
 	return audit
 }
 

internal/event/event.go

@@ -628,6 +628,8 @@ type MemoryRecallAudit struct {
 	UsedChars  int
 	Omitted    int
 	Suppressed string // reason recall stayed silent; "" when hits were injected
+	// Shadow is the Retrieval V2 ranking (telemetry only, never served).
+	Shadow []MemoryRecallHit
 }
 
 // MemoryRecallHit is one recalled fact's content-free fingerprint.

internal/memory/auto_recall.go

@@ -54,10 +54,19 @@ type RecallResult struct {
 	CharBudget int
 	UsedChars  int
 	Suppressed string
+	// ShadowHits is the Retrieval V2 ranking over the same pool, telemetry
+	// only: it never reaches the model and never affects Hits.
+	ShadowHits []ShadowHit
 
 	block string
 }
 
+// ShadowHit is one V2-ranked fact fingerprint.
+type ShadowHit struct {
+	ID    string
+	Score float64
+}
+
 func (r RecallResult) Block() string { return r.block }
 
 // Override explains one project fact that shadows an equivalent global fact
@@ -132,6 +141,9 @@ func AutoRecall(store Store, query string, opts RecallOptions) RecallResult {
 	}
 
 	memories := recallMemories(store.ListAll())
+	// The shadow ranks the full pool before any production gate: what V1
+	// misses entirely is exactly what the comparison must be able to see.
+	result.ShadowHits = shadowRankV2(result.Query, memories)
 	docs := make([]autoRecallDoc, 0, len(memories))
 	for _, memory := range memories {
 		text := autoRecallSearchText(memory)
@@ -222,6 +234,28 @@ func AutoRecall(store Store, query string, opts RecallOptions) RecallResult {
 	return result
 }
 
+// shadowRankV2 runs the Retrieval V2 candidate (BM25F, code-symbol split,
+// mixed CJK grams) over the recall pool. Shadow only: recorded for offline
+// comparison, gated by MemoryBench before it can ever serve.
+func shadowRankV2(query string, memories []Memory) []ShadowHit {
+	docs := make([]retrieval.FieldedDoc, 0, len(memories))
+	for _, m := range memories {
+		docs = append(docs, retrieval.FieldedDoc{ID: m.ID, Fields: map[string]string{
+			"name": m.Name, "title": m.Title, "keywords": m.Keywords,
+			"subject": m.SubjectKey, "description": m.Description, "body": m.Body,
+		}})
+	}
+	ranked := retrieval.RankV2(query, docs)
+	if len(ranked) > maxAutoRecallLimit {
+		ranked = ranked[:maxAutoRecallLimit]
+	}
+	out := make([]ShadowHit, 0, len(ranked))
+	for _, hit := range ranked {
+		out = append(out, ShadowHit{ID: hit.ID, Score: hit.Score})
+	}
+	return out
+}
+
 func recallCharBudget(value int) int {
 	if value == 0 {
 		return defaultAutoRecallChars

internal/memory/shadow_recall_test.go

@@ -0,0 +1,29 @@
+package memory
+
+import "testing"
+
+// The shadow contract: V2 rankings ride the result for telemetry and never
+// change what production recall serves.
+func TestAutoRecallAttachesShadowWithoutChangingProduction(t *testing.T) {
+	store := recallTestStore(t)
+	recallTestWrite(t, store.Dir, Memory{
+		ID: "mem-sym", Name: "retry-config", Title: "Retry configuration",
+		Description: "Where retry limits live", Type: TypeProject, Scope: FactScopeProject,
+		Body: "Configure RetryPolicy.MaxBackoffMs in policy/retry.toml.",
+	})
+	result := AutoRecall(store, "where do I configure retry limits policy", RecallOptions{})
+	if len(result.Hits) == 0 {
+		t.Fatalf("production recall should hit: %+v", result)
+	}
+	if len(result.ShadowHits) == 0 || result.ShadowHits[0].ID != "mem-sym" {
+		t.Fatalf("shadow ranking must ride the result: %+v", result.ShadowHits)
+	}
+	if result.Block() == "" || len(result.Block()) < 10 {
+		t.Fatal("production block must be unaffected")
+	}
+	// The V2 code-symbol split finds what V1 unigrams cannot.
+	segQuery := AutoRecall(store, "max backoff tuning", RecallOptions{})
+	if len(segQuery.ShadowHits) == 0 {
+		t.Fatalf("V2 shadow should match camel-case segments: %+v", segQuery)
+	}
+}

internal/retrieval/v2.go

@@ -0,0 +1,176 @@
+// Retrieval V2, shadow-only: BM25F field weighting, code-symbol splitting,
+// and mixed CJK uni+bigrams. Nothing here serves the model — V2 rankings ride
+// the recall audit next to production's, and MemoryBench decides if V2 ever
+// takes over. Weights are candidates under measurement, not tuned truths.
+package retrieval
+
+import (
+	"math"
+	"strings"
+	"unicode"
+)
+
+// TokensV2 extends Tokens with code-symbol structure and CJK unigrams:
+// CamelCase and snake_case words also emit their segments (FooBar → foobar,
+// foo, bar), and CJK runs emit unigrams alongside bigrams so one-character
+// queries still match. The V1 form of every token is preserved, so V2 recall
+// is a superset of V1's vocabulary.
+func TokensV2(s string) []string {
+	var out []string
+	var word []rune
+	var prev rune
+	cjkRun := 0
+	flushWord := func() {
+		if len(word) == 0 {
+			return
+		}
+		lower := strings.ToLower(string(word))
+		out = append(out, lower)
+		for _, seg := range splitCodeSymbol(word) {
+			if seg != lower {
+				out = append(out, seg)
+			}
+		}
+		word = word[:0]
+	}
+	endCJK := func() { cjkRun = 0 }
+	for _, r := range s {
+		switch {
+		case isCJK(r):
+			flushWord()
+			out = append(out, string(r))
+			if cjkRun > 0 {
+				out = append(out, string([]rune{prev, r}))
+			}
+			prev = r
+			cjkRun++
+		case unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_':
+			endCJK()
+			word = append(word, r)
+		default:
+			flushWord()
+			endCJK()
+		}
+	}
+	flushWord()
+	return out
+}
+
+// splitCodeSymbol yields the lowercase segments of a code identifier:
+// CamelCase humps, snake_case parts, and letter/digit boundaries.
+func splitCodeSymbol(word []rune) []string {
+	var segments []string
+	var seg []rune
+	flush := func() {
+		if len(seg) > 1 { // single letters are noise, not signal
+			segments = append(segments, strings.ToLower(string(seg)))
+		}
+		seg = seg[:0]
+	}
+	for i, r := range word {
+		boundary := r == '_' ||
+			(i > 0 && unicode.IsUpper(r) && !unicode.IsUpper(word[i-1])) ||
+			(i > 0 && unicode.IsDigit(r) != unicode.IsDigit(word[i-1]))
+		if boundary {
+			flush()
+		}
+		if r != '_' {
+			seg = append(seg, r)
+		}
+	}
+	flush()
+	if len(segments) < 2 {
+		return nil // no structure worth indexing beyond the whole word
+	}
+	return segments
+}
+
+// FieldedDoc is one document split into weighted fields for BM25F.
+type FieldedDoc struct {
+	ID     string
+	Fields map[string]string
+}
+
+// v2FieldWeights follow the BM25F intuition: identity fields say what a fact
+// IS, the body says everything it mentions. Candidates under measurement.
+var v2FieldWeights = map[string]float64{
+	"name": 2.5, "title": 2.5, "keywords": 2.0, "subject": 2.0,
+	"description": 1.5, "body": 1.0,
+}
+
+// V2Hit is one shadow-ranked document.
+type V2Hit struct {
+	ID    string
+	Score float64
+}
+
+// RankV2 scores docs against a query with BM25F over TokensV2: term
+// frequencies accumulate per field scaled by field weight, then the standard
+// saturation applies once per term (the F in BM25F), with length
+// normalization over the weighted document length.
+func RankV2(query string, docs []FieldedDoc) []V2Hit {
+	queryTerms := Unique(TokensV2(query))
+	if len(queryTerms) == 0 || len(docs) == 0 {
+		return nil
+	}
+	type indexed struct {
+		id     string
+		tf     map[string]float64
+		length float64
+	}
+	corpus := make([]indexed, 0, len(docs))
+	df := map[string]int{}
+	var totalLen float64
+	for _, doc := range docs {
+		ix := indexed{id: doc.ID, tf: map[string]float64{}}
+		seen := map[string]bool{}
+		for field, text := range doc.Fields {
+			weight, ok := v2FieldWeights[field]
+			if !ok {
+				weight = 1.0
+			}
+			for _, term := range TokensV2(text) {
+				ix.tf[term] += weight
+				ix.length += weight
+				seen[term] = true
+			}
+		}
+		for term := range seen {
+			df[term]++
+		}
+		totalLen += ix.length
+		corpus = append(corpus, ix)
+	}
+	avgLen := totalLen / float64(len(corpus))
+	if avgLen <= 0 {
+		avgLen = 1
+	}
+	const k1, b = 1.2, 0.75
+	var hits []V2Hit
+	for _, doc := range corpus {
+		var score float64
+		for _, term := range queryTerms {
+			tf := doc.tf[term]
+			if tf == 0 || df[term] == 0 {
+				continue
+			}
+			idf := math.Log(1 + (float64(len(corpus))-float64(df[term])+0.5)/(float64(df[term])+0.5))
+			score += idf * (tf * (k1 + 1)) / (tf + k1*(1-b+b*doc.length/avgLen))
+		}
+		if score > 0 {
+			hits = append(hits, V2Hit{ID: doc.id, Score: score})
+		}
+	}
+	SortHitsDesc(hits)
+	return hits
+}
+
+// SortHitsDesc orders shadow hits best-first with a stable ID tiebreak.
+func SortHitsDesc(hits []V2Hit) {
+	for i := 1; i < len(hits); i++ {
+		for j := i; j > 0 && (hits[j].Score > hits[j-1].Score ||
+			(hits[j].Score == hits[j-1].Score && hits[j].ID < hits[j-1].ID)); j-- {
+			hits[j], hits[j-1] = hits[j-1], hits[j]
+		}
+	}
+}

internal/retrieval/v2_test.go

@@ -0,0 +1,47 @@
+package retrieval
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestTokensV2CodeSymbolsAndCJK(t *testing.T) {
+	got := strings.Join(TokensV2("RetryPolicy.MaxBackoffMs ticket_1827 数据库"), " ")
+	for _, want := range []string{"retrypolicy", "retry", "policy", "maxbackoffms", "max", "backoff", "ticket_1827", "ticket", "1827", "数", "据", "库", "数据", "据库"} {
+		if !strings.Contains(" "+got+" ", " "+want+" ") {
+			t.Fatalf("TokensV2 missing %q in %q", want, got)
+		}
+	}
+	if strings.Contains(" "+got+" ", " ms ") {
+		// "Ms" splits to a single meaningful segment only when len > 1; "ms" is
+		// fine — assert single letters never appear instead.
+		t.Log("ms segment present (acceptable)")
+	}
+	for _, tok := range TokensV2("A B C") {
+		if len(tok) == 0 {
+			t.Fatal("empty token emitted")
+		}
+	}
+}
+
+func TestRankV2FieldWeightsBeatBodyMentions(t *testing.T) {
+	docs := []FieldedDoc{
+		{ID: "titled", Fields: map[string]string{"title": "Package manager", "body": "We migrated to pnpm."}},
+		{ID: "mention", Fields: map[string]string{"title": "Sprint notes", "body": "Someone asked about the package manager once; also weather, lunch, and a package of stickers."}},
+	}
+	hits := RankV2("package manager", docs)
+	if len(hits) != 2 || hits[0].ID != "titled" {
+		t.Fatalf("field-weighted doc must outrank body mentions: %+v", hits)
+	}
+}
+
+func TestRankV2FindsCodeSymbolBySegment(t *testing.T) {
+	docs := []FieldedDoc{
+		{ID: "sym", Fields: map[string]string{"body": "Configure RetryPolicy.MaxBackoffMs in policy/retry.toml."}},
+		{ID: "other", Fields: map[string]string{"body": "Unrelated fact about deploys."}},
+	}
+	hits := RankV2("max backoff configuration", docs)
+	if len(hits) == 0 || hits[0].ID != "sym" {
+		t.Fatalf("segment query must find the camel-case symbol: %+v", hits)
+	}
+}

internal/trajectory/recorder.go

@@ -41,6 +41,7 @@ type MemoryRecall struct {
 	UsedChars  int               `json:"used_chars,omitempty"`
 	Omitted    int               `json:"omitted,omitempty"`
 	Suppressed string            `json:"suppressed,omitempty"`
+	ShadowHits []MemoryRecallHit `json:"shadow_hits,omitempty"`
 }
 
 // MemoryRecallHit is one recalled fact's content-free fingerprint.
@@ -228,6 +229,9 @@ func (r *Recorder) RecordMemoryRecall(a event.MemoryRecallAudit) {
 			Type: hit.Type, Freshness: hit.Freshness, Score: hit.Score,
 		})
 	}
+	for _, hit := range a.Shadow {
+		rec.ShadowHits = append(rec.ShadowHits, MemoryRecallHit{ID: hit.ID, Score: hit.Score})
+	}
 	r.append(Record{MemoryRecall: rec})
 	event.RecordMemoryRecall(r.inner, a)
 }