feat(storage): restore HNSW graph from embeddings_hnsw across restarts

storage/hnsw.go

@@ -3,6 +3,7 @@ package storage
 import (
 	"container/heap"
 	"context"
+	"database/sql"
 	"encoding/json"
 	"math"
 	"math/rand"
@@ -29,11 +30,6 @@ type HNSWIndex struct {
 	rng *rand.Rand
 }
 
-type HNSWNode struct {
-	ID    string `json:"id"`
-	Level int    `json:"level"`
-}
-
 // NewHNSWIndex creates a new HNSW index with the given parameters.
 func NewHNSWIndex(efConstruction, m, efSearch int) *HNSWIndex {
 	return &HNSWIndex{
@@ -103,7 +99,12 @@ func (h *HNSWIndex) Build(ctx context.Context, s *Store, model string) error {
 	if err != nil {
 		return err
 	}
+	return h.buildFrom(ctx, s, model, embeddings)
+}
 
+// buildFrom constructs the graph from an in-memory embedding set and
+// persists it.
+func (h *HNSWIndex) buildFrom(ctx context.Context, s *Store, model string, embeddings map[string][]float32) error {
 	h.mu.Lock()
 	h.vectors = embeddings
 	h.graph = make(map[string]map[int][]string)
@@ -124,6 +125,108 @@ func (h *HNSWIndex) Build(ctx context.Context, s *Store, model string) error {
 	return h.persist(ctx, s, model)
 }
 
+// neighborsPayload is the JSON shape persisted in embeddings_hnsw.neighbors.
+// Parameters are recorded so a Restore can reject graphs built with a
+// different configuration (link lists depend on m and efConstruction).
+type neighborsPayload struct {
+	Version        int                 `json:"version"`
+	M              int                 `json:"m"`
+	EfConstruction int                 `json:"efConstruction"`
+	Levels         map[string][]string `json:"levels"`
+}
+
+const neighborsPayloadVersion = 1
+
+// Restore reconstructs the index from the neighbors graph persisted by a
+// previous Build/Upsert/Remove cycle, without recomputing distances. It
+// returns true when a stored graph was loaded. It returns false (and leaves
+// the index empty) when the stored graph is absent, structurally
+// incompatible (different parameters, old payload shape), or out of sync
+// with the embeddings table — callers should then fall back to Build.
+func (h *HNSWIndex) Restore(ctx context.Context, s *Store, model string) (bool, error) {
+	embeddings, err := s.AllEmbeddings(ctx, model)
+	if err != nil {
+		return false, err
+	}
+	if len(embeddings) == 0 {
+		return false, nil // nothing indexed; nothing to restore
+	}
+
+	qctx, cancel := s.withTimeout(ctx)
+	defer cancel()
+	rows, err := s.db.QueryContext(qctx,
+		`SELECT node_id, neighbors FROM embeddings_hnsw WHERE model=?`, model)
+	if err != nil {
+		return false, err
+	}
+	defer func() { _ = rows.Close() }()
+
+	graph := make(map[string]map[string][]string)
+	for rows.Next() {
+		var nodeID string
+		var neighbors sql.NullString
+		if err := rows.Scan(&nodeID, &neighbors); err != nil {
+			return false, err
+		}
+		var payload neighborsPayload
+		var parsed map[string][]string
+		if !neighbors.Valid || neighbors.String == "" {
+			return false, nil // legacy or empty payload — must rebuild
+		}
+		if err := json.Unmarshal([]byte(neighbors.String), &payload); err != nil {
+			return false, nil // unrecognized shape — must rebuild
+		}
+		if payload.Version != neighborsPayloadVersion ||
+			payload.M != h.m || payload.EfConstruction != h.efConstruction {
+			return false, nil // incompatible parameters — must rebuild
+		}
+		if parsed = payload.Levels; parsed == nil {
+			return false, nil
+		}
+		graph[nodeID] = parsed
+	}
+	if err := rows.Err(); err != nil {
+		return false, err
+	}
+
+	// Exact membership match: the graph is only reusable when it covers the
+	// current embedding set for this model.
+	if len(graph) != len(embeddings) {
+		return false, nil
+	}
+	for nodeID := range graph {
+		if _, ok := embeddings[nodeID]; !ok {
+			return false, nil
+		}
+	}
+
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.vectors = embeddings
+	h.graph = make(map[string]map[int][]string, len(graph))
+	for nodeID, rawLevels := range graph {
+		levels := make(map[int][]string, len(rawLevels))
+		for key, links := range rawLevels {
+			l, err := strconv.Atoi(key)
+			if err != nil {
+				return false, nil // corrupted payload — must rebuild
+			}
+			cleaned := make([]string, 0, len(links))
+			for _, id := range links {
+				if _, ok := embeddings[id]; ok {
+					cleaned = append(cleaned, id)
+				}
+			}
+			levels[l] = cleaned
+		}
+		h.graph[nodeID] = levels
+	}
+	entry, maxLevel := h.topNodeLocked()
+	h.entryPoint = entry
+	h.maxLevel = maxLevel
+	return true, nil
+}
+
 // insert adds a single vector to the graph. Caller must hold h.mu.
 func (h *HNSWIndex) insert(nodeID string, vec []float32) {
 	if len(h.graph) == 0 {
@@ -318,15 +421,16 @@ func (h *HNSWIndex) persistNodes(ctx context.Context, s *Store, model string, id
 			continue
 		}
 		vec := h.vectors[nodeID]
-		byLevel := make(map[string][]HNSWNode, len(levels))
+		byLevel := make(map[string][]string, len(levels))
 		for l, links := range levels {
-			nodes := make([]HNSWNode, len(links))
-			for i, id := range links {
-				nodes[i] = HNSWNode{ID: id, Level: l}
-			}
-			byLevel[strconv.Itoa(l)] = nodes
+			byLevel[strconv.Itoa(l)] = links
 		}
-		neighborsJSON, err := json.Marshal(byLevel)
+		payload, err := json.Marshal(neighborsPayload{
+			Version:        neighborsPayloadVersion,
+			M:              h.m,
+			EfConstruction: h.efConstruction,
+			Levels:         byLevel,
+		})
 		if err != nil {
 			return err
 		}
@@ -335,7 +439,7 @@ func (h *HNSWIndex) persistNodes(ctx context.Context, s *Store, model string, id
 			VALUES (?, ?, ?, ?, ?)
 			ON CONFLICT(node_id) DO UPDATE SET
 				vector=excluded.vector, neighbors=excluded.neighbors, updated_at=excluded.updated_at
-		`, nodeID, EncodeVector(vec), model, string(neighborsJSON), time.Now().UTC())
+		`, nodeID, EncodeVector(vec), model, string(payload), time.Now().UTC())
 		if err != nil {
 			return err
 		}

storage/hnsw_restore_test.go

@@ -0,0 +1,181 @@
+package storage
+
+import (
+	"context"
+	"fmt"
+	"testing"
+)
+
+// TestHNSWRestoreAcrossRestart verifies a fresh store on the same database
+// restores the persisted HNSW graph instead of recomputing it.
+func TestHNSWRestoreAcrossRestart(t *testing.T) {
+	dir := t.TempDir()
+	dbPath := dir + "/restart.db"
+	ctx := context.Background()
+	model := "restart-model"
+
+	s1, err := NewStore(dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 10; i++ {
+		id := fmt.Sprintf("r-%d", i)
+		if err := s1.CreateNode(ctx, &Node{
+			ID: id, Type: "convention", Content: id,
+			ContentHash: "h-" + id, Scope: "project", Project: "test",
+		}); err != nil {
+			t.Fatal(err)
+		}
+		vec := make([]float32, 8)
+		vec[0] = 1.0
+		vec[1] = float32(i) * 0.1
+		if err := s1.SaveEmbedding(ctx, id, model, vec); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := s1.BuildHNSWIndex(ctx, model); err != nil {
+		t.Fatal(err)
+	}
+	wantIDs, _, err := s1.SearchHNSW(ctx, model, dirVec(t, 0.5), 3, 50)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Close releases the process lock; the second store simulates a restart.
+	if err := s1.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	s2, err := NewStore(dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = s2.Close() }()
+
+	// The index cache is empty in a new process; Restore must load the
+	// persisted graph without a fallback build.
+	idx := NewHNSWIndex(200, 16, 50)
+	restored, err := idx.Restore(ctx, s2, model)
+	if err != nil {
+		t.Fatalf("Restore: %v", err)
+	}
+	if !restored {
+		t.Fatal("expected persisted graph to be restorable, got false")
+	}
+
+	// And the full search path must return identical results after restart.
+	gotIDs, _, err := s2.SearchHNSW(ctx, model, dirVec(t, 0.5), 3, 50)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(gotIDs) != len(wantIDs) {
+		t.Fatalf("result count changed across restart: %v vs %v", gotIDs, wantIDs)
+	}
+	for i := range wantIDs {
+		if gotIDs[i] != wantIDs[i] {
+			t.Errorf("result %d: got %s want %s", i, gotIDs[i], wantIDs[i])
+		}
+	}
+}
+
+// TestHNSWRestoreFallbackOnStaleGraph verifies that a persisted graph which
+// no longer matches the embeddings table is rejected and a rebuild takes
+// over (via the SearchHNSW fallback path).
+func TestHNSWRestoreFallbackOnStaleGraph(t *testing.T) {
+	dir := t.TempDir()
+	dbPath := dir + "/stale.db"
+	ctx := context.Background()
+	model := "stale-model"
+
+	s1, err := NewStore(dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 10; i++ {
+		id := fmt.Sprintf("s-%d", i)
+		if err := s1.CreateNode(ctx, &Node{
+			ID: id, Type: "convention", Content: id,
+			ContentHash: "h-" + id, Scope: "project", Project: "test",
+		}); err != nil {
+			t.Fatal(err)
+		}
+		vec := make([]float32, 8)
+		vec[0] = 1.0
+		vec[1] = float32(i) * 0.1
+		if err := s1.SaveEmbedding(ctx, id, model, vec); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := s1.BuildHNSWIndex(ctx, model); err != nil {
+		t.Fatal(err)
+	}
+	if err := s1.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	s2, err := NewStore(dbPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = s2.Close() }()
+
+	// Simulate an out-of-band embedding write that bypassed the index API:
+	// a row in embeddings with no corresponding graph row.
+	if err := s2.CreateNode(ctx, &Node{
+		ID: "s-ghost", Type: "convention", Content: "s-ghost",
+		ContentHash: "h-s-ghost", Scope: "project", Project: "test",
+	}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := s2.DB().ExecContext(ctx,
+		`INSERT INTO embeddings(node_id, vector, model) VALUES('s-ghost', X'0000803F0000803F0000803F0000803F0000803F0000803F0000803F0000803F', ?)`, model); err != nil {
+		t.Fatal(err)
+	}
+
+	idx := NewHNSWIndex(200, 16, 50)
+	restored, err := idx.Restore(ctx, s2, model)
+	if err != nil {
+		t.Fatalf("Restore: %v", err)
+	}
+	if restored {
+		t.Fatal("expected stale graph to be rejected, but Restore returned true")
+	}
+
+	// SearchHNSW must fall back to a full build and still return results.
+	ids, _, err := s2.SearchHNSW(ctx, model, dirVec(t, 0.9), 1, 50)
+	if err != nil {
+		t.Fatalf("SearchHNSW after stale graph: %v", err)
+	}
+	if len(ids) == 0 {
+		t.Fatal("expected results after rebuild fallback")
+	}
+
+	// The fallback build must have re-persisted the graph so the next
+	// process can restore it again.
+	idx2 := NewHNSWIndex(200, 16, 50)
+	restored2, err := idx2.Restore(ctx, s2, model)
+	if err != nil {
+		t.Fatalf("Restore after rebuild: %v", err)
+	}
+	if !restored2 {
+		t.Fatal("expected rebuilt graph to be restorable")
+	}
+}
+
+// TestHNSWRestoreRejectsForeignParameters verifies graphs persisted by an
+// index with different construction parameters are rejected.
+func TestHNSWRestoreRejectsForeignParameters(t *testing.T) {
+	s, cleanup := setupStore(t)
+	defer cleanup()
+	ctx := context.Background()
+	model := "param-model"
+	seedHNSW(t, s, model, 6, 0.1) // persists via Build (m=16, ef=200)
+
+	other := NewHNSWIndex(400, 32, 50) // different construction parameters
+	restored, err := other.Restore(ctx, s, model)
+	if err != nil {
+		t.Fatalf("Restore: %v", err)
+	}
+	if restored {
+		t.Fatal("expected parameter-mismatched graph to be rejected")
+	}
+}