ensuring embeddings are removed and coutns are updated on delets

Author: orneryd

pkg/nornicdb/db.go

@@ -951,15 +951,20 @@ func Open(dataDir string, config *Config) (*DB, error) {
 
 	// Wire up storage event callbacks to keep search indexes synchronized
 	// Storage is the single source of truth - it notifies when changes happen
-	// This works whether storage is BadgerEngine directly or wrapped in AsyncEngine
+	// The storage chain can be: AsyncEngine -> WALEngine -> BadgerEngine
 	var underlyingEngine storage.Engine = db.storage
 
-	// If storage is wrapped in AsyncEngine, get the underlying BadgerEngine
-	if asyncEngine, ok := db.storage.(*storage.AsyncEngine); ok {
-		underlyingEngine = asyncEngine.GetUnderlying()
+	// Unwrap AsyncEngine if present
+	if asyncEngine, ok := underlyingEngine.(*storage.AsyncEngine); ok {
+		underlyingEngine = asyncEngine.GetEngine()
+	}
+
+	// Unwrap WALEngine if present
+	if walEngine, ok := underlyingEngine.(*storage.WALEngine); ok {
+		underlyingEngine = walEngine.GetEngine()
 	}
 
-	// Set callbacks on the actual storage engine (where operations actually happen)
+	// Set callbacks on the actual storage engine (BadgerEngine which implements StorageEventNotifier)
 	if notifier, ok := underlyingEngine.(storage.StorageEventNotifier); ok {
 		// When a node is created, automatically index it for search
 		notifier.OnNodeCreated(func(node *storage.Node) {

pkg/nornicdb/db_test.go

@@ -2539,6 +2539,94 @@ func TestDeleteNode_RemovesFromSearchIndex(t *testing.T) {
 		assert.Error(t, err)
 		assert.ErrorIs(t, err, ErrClosed)
 	})
+
+	t.Run("delete_decrements_embedding_count_via_callback", func(t *testing.T) {
+		// This test verifies the full callback chain:
+		// Storage.DeleteNode -> notifyNodeDeleted callback -> SearchService.RemoveNode
+		// This ensures embeddings are cleaned up automatically when nodes are deleted.
+
+		db, err := Open(t.TempDir(), nil)
+		require.NoError(t, err)
+		defer db.Close()
+
+		// Create nodes with embeddings directly via storage layer
+		// We use the storage layer directly to have precise control over the Embedding field
+		nodes := []*storage.Node{
+			{
+				ID:         "embed-callback-test-1",
+				Labels:     []string{"TestNode"},
+				Properties: map[string]any{"name": "Test Node 1"},
+				Embedding:  make([]float32, 1024), // Match default dimension
+			},
+			{
+				ID:         "embed-callback-test-2",
+				Labels:     []string{"TestNode"},
+				Properties: map[string]any{"name": "Test Node 2"},
+				Embedding:  make([]float32, 1024),
+			},
+			{
+				ID:         "embed-callback-test-3",
+				Labels:     []string{"TestNode"},
+				Properties: map[string]any{"name": "Test Node 3"},
+				Embedding:  make([]float32, 1024),
+			},
+		}
+
+		// Set distinct embedding values
+		nodes[0].Embedding[0] = 1.0
+		nodes[1].Embedding[1] = 1.0
+		nodes[2].Embedding[2] = 1.0
+
+		// Get initial embedding count
+		initialCount := db.EmbeddingCount()
+
+		// Create nodes via storage (this triggers OnNodeCreated callback)
+		for _, node := range nodes {
+			err := db.storage.CreateNode(node)
+			require.NoError(t, err)
+		}
+
+		// Wait for async callbacks to fire
+		time.Sleep(100 * time.Millisecond)
+
+		// Verify embedding count increased by 3
+		afterCreate := db.EmbeddingCount()
+		assert.Equal(t, initialCount+3, afterCreate,
+			"Embedding count should increase by 3 after creating nodes (was %d, now %d)",
+			initialCount, afterCreate)
+
+		// Delete node1 via storage (this triggers OnNodeDeleted callback)
+		err = db.storage.DeleteNode("embed-callback-test-1")
+		require.NoError(t, err)
+
+		// Wait for async callback
+		time.Sleep(100 * time.Millisecond)
+
+		afterDelete1 := db.EmbeddingCount()
+		assert.Equal(t, initialCount+2, afterDelete1,
+			"Embedding count should decrease by 1 after deleting node1 (expected %d, got %d)",
+			initialCount+2, afterDelete1)
+
+		// Delete node2
+		err = db.storage.DeleteNode("embed-callback-test-2")
+		require.NoError(t, err)
+		time.Sleep(100 * time.Millisecond)
+
+		afterDelete2 := db.EmbeddingCount()
+		assert.Equal(t, initialCount+1, afterDelete2,
+			"Embedding count should decrease by 2 after deleting node2 (expected %d, got %d)",
+			initialCount+1, afterDelete2)
+
+		// Delete node3
+		err = db.storage.DeleteNode("embed-callback-test-3")
+		require.NoError(t, err)
+		time.Sleep(100 * time.Millisecond)
+
+		afterDelete3 := db.EmbeddingCount()
+		assert.Equal(t, initialCount, afterDelete3,
+			"Embedding count should return to initial after deleting all test nodes (expected %d, got %d)",
+			initialCount, afterDelete3)
+	})
 }
 
 func TestDeleteUserData_RemovesFromSearchIndex(t *testing.T) {

pkg/search/search_test.go

@@ -685,6 +685,126 @@ func TestSearchService_RemoveNode(t *testing.T) {
 	}
 }
 
+// TestSearchService_RemoveNode_DecrementsEmbeddingCount verifies that removing a node
+// decrements the embedding count in stats. This is critical for ensuring that when nodes
+// are deleted via Cypher, the embedding count is updated correctly without requiring
+// a manual "regenerate" operation.
+func TestSearchService_RemoveNode_DecrementsEmbeddingCount(t *testing.T) {
+	engine := storage.NewMemoryEngine()
+	defer engine.Close()
+
+	svc := NewServiceWithDimensions(engine, 4) // 4-dimensional embeddings for test
+
+	// Initial state: no embeddings
+	assert.Equal(t, 0, svc.EmbeddingCount(), "Should start with 0 embeddings")
+
+	// Create and index 3 nodes with embeddings
+	// Note: Node.Embedding is the struct field used by IndexNode, not Properties["embedding"]
+	nodes := []*storage.Node{
+		{
+			ID:         "node1",
+			Labels:     []string{"Person"},
+			Properties: map[string]any{"name": "Alice"},
+			Embedding:  []float32{1, 0, 0, 0},
+		},
+		{
+			ID:         "node2",
+			Labels:     []string{"Person"},
+			Properties: map[string]any{"name": "Bob"},
+			Embedding:  []float32{0, 1, 0, 0},
+		},
+		{
+			ID:         "node3",
+			Labels:     []string{"Person"},
+			Properties: map[string]any{"name": "Charlie"},
+			Embedding:  []float32{0, 0, 1, 0},
+		},
+	}
+
+	for _, node := range nodes {
+		require.NoError(t, engine.CreateNode(node))
+		require.NoError(t, svc.IndexNode(node))
+	}
+
+	// Verify all 3 nodes are indexed
+	assert.Equal(t, 3, svc.EmbeddingCount(), "Should have 3 embeddings after indexing")
+
+	// Remove node1
+	require.NoError(t, svc.RemoveNode("node1"))
+	assert.Equal(t, 2, svc.EmbeddingCount(), "Should have 2 embeddings after removing node1")
+
+	// Remove node2
+	require.NoError(t, svc.RemoveNode("node2"))
+	assert.Equal(t, 1, svc.EmbeddingCount(), "Should have 1 embedding after removing node2")
+
+	// Remove node3
+	require.NoError(t, svc.RemoveNode("node3"))
+	assert.Equal(t, 0, svc.EmbeddingCount(), "Should have 0 embeddings after removing all nodes")
+
+	// Removing non-existent node should not affect count
+	require.NoError(t, svc.RemoveNode("non-existent"))
+	assert.Equal(t, 0, svc.EmbeddingCount(), "Count should remain 0 after removing non-existent node")
+}
+
+// TestSearchService_RemoveNode_OnlyRemovesTargetNode ensures RemoveNode is precise
+// and doesn't affect other nodes' embeddings.
+func TestSearchService_RemoveNode_OnlyRemovesTargetNode(t *testing.T) {
+	engine := storage.NewMemoryEngine()
+	defer engine.Close()
+
+	svc := NewServiceWithDimensions(engine, 4) // 4-dimensional embeddings for test
+
+	// Create distinct embeddings for easy search verification
+	// Note: Node.Embedding is the struct field used by IndexNode
+	node1 := &storage.Node{
+		ID:         "target-to-remove",
+		Labels:     []string{"Document"},
+		Properties: map[string]any{"content": "unique alpha content"},
+		Embedding:  []float32{1, 0, 0, 0},
+	}
+	node2 := &storage.Node{
+		ID:         "should-remain-1",
+		Labels:     []string{"Document"},
+		Properties: map[string]any{"content": "unique beta content"},
+		Embedding:  []float32{0, 1, 0, 0},
+	}
+	node3 := &storage.Node{
+		ID:         "should-remain-2",
+		Labels:     []string{"Document"},
+		Properties: map[string]any{"content": "unique gamma content"},
+		Embedding:  []float32{0, 0, 1, 0},
+	}
+
+	for _, node := range []*storage.Node{node1, node2, node3} {
+		require.NoError(t, engine.CreateNode(node))
+		require.NoError(t, svc.IndexNode(node))
+	}
+
+	assert.Equal(t, 3, svc.EmbeddingCount())
+
+	// Remove only the target node
+	require.NoError(t, svc.RemoveNode("target-to-remove"))
+
+	// Verify count decreased
+	assert.Equal(t, 2, svc.EmbeddingCount())
+
+	// Verify remaining nodes are still searchable
+	opts := DefaultSearchOptions()
+
+	// Search for remaining nodes by their unique content
+	response, err := svc.Search(context.Background(), "beta", nil, opts)
+	require.NoError(t, err)
+	found := false
+	for _, r := range response.Results {
+		if r.ID == "should-remain-1" {
+			found = true
+		}
+		// The removed node should NOT appear
+		assert.NotEqual(t, "target-to-remove", r.ID, "Removed node should not appear in results")
+	}
+	assert.True(t, found, "Remaining node 'should-remain-1' should be searchable")
+}
+
 // TestSearchService_HybridSearch tests the hybrid RRF search.
 func TestSearchService_HybridSearch(t *testing.T) {
 	engine := storage.NewMemoryEngine()