fix encapsulation bypass on registry internals

pkg/tools/registry.go

@@ -88,6 +88,45 @@ func (r *ToolRegistry) TickTTL() {
 	}
 }
 
+// Version returns the current registry version (atomically).
+func (r *ToolRegistry) Version() uint64 {
+	return r.version.Load()
+}
+
+// HiddenToolSnapshot holds a consistent snapshot of hidden tools and the
+// registry version at which it was taken. Used by BM25SearchTool cache.
+type HiddenToolSnapshot struct {
+	Docs    []HiddenToolDoc
+	Version uint64
+}
+
+// HiddenToolDoc is a lightweight representation of a hidden tool for search indexing.
+type HiddenToolDoc struct {
+	Name        string
+	Description string
+}
+
+// SnapshotHiddenTools returns all non-core tools and the current registry
+// version under a single read-lock, guaranteeing consistency between the
+// two values.
+func (r *ToolRegistry) SnapshotHiddenTools() HiddenToolSnapshot {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	docs := make([]HiddenToolDoc, 0, len(r.tools))
+	for name, entry := range r.tools {
+		if !entry.IsCore {
+			docs = append(docs, HiddenToolDoc{
+				Name:        name,
+				Description: entry.Tool.Description(),
+			})
+		}
+	}
+	return HiddenToolSnapshot{
+		Docs:    docs,
+		Version: r.version.Load(),
+	}
+}
+
 func (r *ToolRegistry) Get(name string) (Tool, bool) {
 	r.mu.RLock()
 	defer r.mu.RUnlock()

pkg/tools/search_tool.go

@@ -214,85 +214,69 @@ type bm25CachedEngine struct {
 	engine *utils.BM25Engine[searchDoc]
 }
 
+// snapshotToSearchDocs converts a HiddenToolSnapshot to BM25 searchDoc slice.
+func snapshotToSearchDocs(snap HiddenToolSnapshot) []searchDoc {
+	docs := make([]searchDoc, len(snap.Docs))
+	for i, d := range snap.Docs {
+		docs[i] = searchDoc{Name: d.Name, Description: d.Description}
+	}
+	return docs
+}
+
+// buildBM25Engine creates a BM25Engine from a slice of searchDocs.
+func buildBM25Engine(docs []searchDoc) *utils.BM25Engine[searchDoc] {
+	return utils.NewBM25Engine(
+		docs,
+		func(doc searchDoc) string {
+			return doc.Name + " " + doc.Description
+		},
+	)
+}
+
 // getOrBuildEngine returns a cached BM25 engine, rebuilding it only when
 // the registry version has changed (new tools registered).
 func (t *BM25SearchTool) getOrBuildEngine() *bm25CachedEngine {
-	// Fast path: optimistic check without locking the registry.
+	// Fast path: optimistic check without locking.
-	// If the version hasn't changed, the cache is still valid.
+	if t.cachedEngine != nil && t.cacheVersion == t.registry.Version() {
-	if t.cachedEngine != nil && t.cacheVersion == t.registry.version.Load() {
 		return t.cachedEngine
 	}
 
 	t.cacheMu.Lock()
 	defer t.cacheMu.Unlock()
 
-	// Read version inside the registry RLock so the snapshot and version
+	// Snapshot + version are read under a single registry RLock,
-	// are guaranteed to be consistent (no TOCTOU between Load and RLock).
+	// guaranteeing consistency (no TOCTOU).
-	t.registry.mu.RLock()
+	snap := t.registry.SnapshotHiddenTools()
-	snapshotVersion := t.registry.version.Load()
-	snapshot := make([]searchDoc, 0, len(t.registry.tools))
-	for name, entry := range t.registry.tools {
-		if !entry.IsCore {
-			snapshot = append(snapshot, searchDoc{
-				Name:        name,
-				Description: entry.Tool.Description(),
-			})
-		}
-	}
-	t.registry.mu.RUnlock()
 
 	// Re-check: another goroutine may have rebuilt while we waited for cacheMu.
-	if t.cachedEngine != nil && t.cacheVersion == snapshotVersion {
+	if t.cachedEngine != nil && t.cacheVersion == snap.Version {
 		return t.cachedEngine
 	}
 
-	if len(snapshot) == 0 {
+	docs := snapshotToSearchDocs(snap)
+	if len(docs) == 0 {
 		t.cachedEngine = nil
-		t.cacheVersion = snapshotVersion
+		t.cacheVersion = snap.Version
 		return nil
 	}
 
-	engine := utils.NewBM25Engine(
+	cached := &bm25CachedEngine{engine: buildBM25Engine(docs)}
-		snapshot,
-		func(doc searchDoc) string {
-			return doc.Name + " " + doc.Description
-		},
-	)
-
-	cached := &bm25CachedEngine{engine: engine}
 	t.cachedEngine = cached
-	t.cacheVersion = snapshotVersion
+	t.cacheVersion = snap.Version
 	return cached
 }
 
 // SearchBM25 ranks hidden tools against query using BM25 via utils.BM25Engine.
-// The corpus snapshot is built under the registry read-lock, then released
+// This non-cached variant rebuilds the engine on every call. Used by tests
-// before scoring so the lock is not held during CPU-intensive work.
+// and any code that doesn't hold a BM25SearchTool instance.
 func (r *ToolRegistry) SearchBM25(query string, maxSearchResults int) []ToolSearchResult {
-	r.mu.RLock()
+	snap := r.SnapshotHiddenTools()
-	snapshot := make([]searchDoc, 0, len(r.tools))
+	docs := snapshotToSearchDocs(snap)
-	for name, entry := range r.tools {
+	if len(docs) == 0 {
-		if !entry.IsCore {
-			snapshot = append(snapshot, searchDoc{
-				Name:        name,
-				Description: entry.Tool.Description(),
-			})
-		}
-	}
-	r.mu.RUnlock()
-
-	if len(snapshot) == 0 {
 		return nil
 	}
 
-	engine := utils.NewBM25Engine(
+	ranked := buildBM25Engine(docs).Search(query, maxSearchResults)
-		snapshot,
-		func(doc searchDoc) string {
-			return doc.Name + " " + doc.Description
-		},
-	)
-
-	ranked := engine.Search(query, maxSearchResults)
 	if len(ranked) == 0 {
 		return nil
 	}