fix(persistence): align dumps with component ownership

reme/components/file_graph/local_file_graph.py

@@ -19,7 +19,6 @@ class LocalFileGraph(BaseFileGraph):
         self._inverse: dict[str, set[str]] = {}  # real target → sources
         self._pending: dict[str, set[str]] = {}  # virtual target → sources
         self._graph_file: Path = self.component_metadata_path / f"{self.name}.jsonl.zst"
-        self._needs_persist = True
 
     # -- Lifecycle ---------------------------------------------------------
 
@@ -36,17 +35,13 @@ class LocalFileGraph(BaseFileGraph):
                 if line.strip():
                     node = FileNode.model_validate_json(line)
                     self._nodes[node.path] = node
-            self._needs_persist = False
             self.logger.debug(f"Loaded {len(self._nodes)} nodes from {self._graph_file}")
         except Exception as e:
             self.logger.exception(f"Failed to load {self._graph_file}: {e}")
 
     async def dump(self) -> None:
-        if not self._needs_persist:
-            return
         try:
             write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values()))
-            self._needs_persist = False
             self.logger.info(f"Saved {len(self._nodes)} nodes to {self._graph_file}")
         except Exception as e:
             self.logger.exception(f"Failed to write {self._graph_file}: {e}")
@@ -83,8 +78,6 @@ class LocalFileGraph(BaseFileGraph):
     # -- Node CRUD ---------------------------------------------------------
 
     async def upsert_nodes(self, nodes: list[FileNode]) -> None:
-        if nodes:
-            self._needs_persist = True
         for node in nodes:
             path = node.path
             old = self._nodes.get(path)
@@ -103,7 +96,6 @@ class LocalFileGraph(BaseFileGraph):
             node = self._nodes.pop(path, None)
             if node is None:
                 continue
-            self._needs_persist = True
             for target in self._targets(node):
                 self._remove_edge(path, target)
             demoted = self._inverse.pop(path, None)
@@ -127,7 +119,6 @@ class LocalFileGraph(BaseFileGraph):
         self._inverse.clear()
         self._pending.clear()
         self._graph_file.unlink(missing_ok=True)
-        self._needs_persist = False
 
     # -- Link access -------------------------------------------------------
 

reme/components/file_graph/nx_file_graph.py

@@ -25,7 +25,6 @@ class NxFileGraph(BaseFileGraph):
         self._graph = nx.MultiDiGraph()
         self.component_metadata_path.mkdir(parents=True, exist_ok=True)
         self._graph_file: Path = self.component_metadata_path / f"{self.name}.pkl"
-        self._needs_persist = True
 
     # -- Lifecycle ---------------------------------------------------------
 
@@ -35,20 +34,16 @@ class NxFileGraph(BaseFileGraph):
         try:
             with open(self._graph_file, "rb") as f:
                 self._graph = pickle.load(f)
-            self._needs_persist = False
             self.logger.info(f"Loaded {self._real_count()} nodes from {self._graph_file}")
         except Exception as e:
             self.logger.exception(f"Failed to load {self._graph_file}: {e}")
 
     async def dump(self) -> None:
-        if not self._needs_persist:
-            return
         try:
             tmp = self._graph_file.with_suffix(".tmp")
             with open(tmp, "wb") as f:
                 pickle.dump(self._graph, f, protocol=pickle.HIGHEST_PROTOCOL)
             tmp.replace(self._graph_file)
-            self._needs_persist = False
             self.logger.info(f"Saved {self._real_count()} nodes to {self._graph_file}")
         except Exception as e:
             self.logger.exception(f"Failed to write {self._graph_file}: {e}")
@@ -74,8 +69,6 @@ class NxFileGraph(BaseFileGraph):
     # -- Node CRUD ---------------------------------------------------------
 
     async def upsert_nodes(self, nodes: list[FileNode]) -> None:
-        if nodes:
-            self._needs_persist = True
         for node in nodes:
             path = node.path
             if self._graph.has_node(path):
@@ -87,7 +80,6 @@ class NxFileGraph(BaseFileGraph):
         for path in paths:
             if not self._graph.has_node(path):
                 continue
-            self._needs_persist = True
             self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True)))
             self._graph.nodes[path].pop("node", None)  # demote to virtual
             if self._graph.in_degree(path) == 0:
@@ -109,7 +101,6 @@ class NxFileGraph(BaseFileGraph):
     async def clear(self):
         self._graph.clear()
         self._graph_file.unlink(missing_ok=True)
-        self._needs_persist = False
 
     # -- Link access -------------------------------------------------------
 

reme/components/file_store/faiss_local_file_store.py

@@ -496,10 +496,10 @@ class FaissLocalFileStore(LocalFileStore):
             self.faiss_idmap_path.unlink(missing_ok=True)
             return False
 
-    async def dump(self) -> None:
-        """Persist chunks JSONL via the parent, then write the FAISS sidecar atomically."""
+    async def _dump_owned_state(self) -> None:
+        """Persist chunks and the FAISS sidecar, excluding dependency snapshots."""
         async with self._faiss_dump_lock:
-            await super().dump()
+            await super()._dump_owned_state()
             if self._faiss_index is None or self.embedding_store is None:
                 return
             try:

reme/components/file_store/local_file_store.py

@@ -100,7 +100,11 @@ class LocalFileStore(BaseFileStore):
     async def _close(self) -> None:
         self._closing = True
         await self._cancel_embedding_backfill()
-        await self.dump()
+        # Dependencies are closed separately by Application (reverse
+        # topological order) or BaseComponent (owned standalone dependencies).
+        # Persist only this store's local state here so each component writes
+        # exactly once during shutdown.
+        await self._dump_owned_state()
         self.file_chunks.clear()
         await super()._close()
 
@@ -529,9 +533,8 @@ class LocalFileStore(BaseFileStore):
         elapsed = time.monotonic() - started_at
         self.logger.info(f"{self.name}: keyword index rebuild complete: total={total}, elapsed={elapsed:.2f}s")
 
-    async def dump(self) -> None:
-        """Atomically rewrite the JSONL, then cascade dump into keyword_index and file_graph."""
-        assert self.file_graph is not None
+    async def _dump_owned_state(self) -> None:
+        """Persist state owned by this store, excluding dependency snapshots."""
         try:
             write_jsonl_zst(
                 self.chunks_path,
@@ -541,6 +544,11 @@ class LocalFileStore(BaseFileStore):
             self.logger.info(f"Saved {len(self.file_chunks)} chunks to {self.chunks_path}")
         except Exception as e:
             self.logger.exception(f"Failed to write {self.chunks_path}: {e}")
+
+    async def dump(self) -> None:
+        """Persist a complete store/index/graph consistency checkpoint."""
+        assert self.file_graph is not None
+        await self._dump_owned_state()
         if self.keyword_index:
             await self.keyword_index.dump()
         await self.file_graph.dump()

reme/components/file_store/zvec_local_file_store.py

@@ -319,9 +319,9 @@ class ZvecLocalFileStore(LocalFileStore):
             f"elapsed={time.monotonic() - started_at:.3f}s",
         )
 
-    async def dump(self) -> None:
-        """Persist chunks JSONL via the parent, then flush zvec and write the sidecar."""
-        await super().dump()
+    async def _dump_owned_state(self) -> None:
+        """Persist chunks and zvec state, excluding dependency snapshots."""
+        await super()._dump_owned_state()
         if self._collection is None or self.embedding_store is None:
             return
         try:

reme/components/keyword_index/bm25_index.py

@@ -52,10 +52,6 @@ class BM25Index(BaseKeywordIndex):
 
         # IDF cache; invalidated whenever live-doc count or postings change.
         self._idf_cache: dict[int, float] = {}
-        # A fresh component still needs one persistence pass so an empty stale
-        # snapshot can be removed. Successful dump/load operations clear this
-        # flag; mutations set it again.
-        self._needs_persist = True
 
     # -- Properties -----------------------------------------------------------
 
@@ -249,8 +245,6 @@ class BM25Index(BaseKeywordIndex):
         if not docs_dict:
             return
 
-        self._needs_persist = True
-
         new_doc_ids: list[str] = []
         new_doc_lens: list[int] = []
         new_doc_token_ids: list[np.ndarray] = []
@@ -278,8 +272,6 @@ class BM25Index(BaseKeywordIndex):
 
     async def delete_docs(self, doc_ids: list[str]) -> None:
         """Lazy-delete a batch of doc_ids; physical reclaim happens in optimize_index."""
-        if doc_ids:
-            self._needs_persist = True
         for doc_id in doc_ids:
             self._remove_doc(doc_id)
         self._idf_cache = {}
@@ -362,11 +354,8 @@ class BM25Index(BaseKeywordIndex):
 
     async def dump(self) -> None:
         """Persist the index via temp file + atomic rename to avoid torn writes."""
-        if not self._needs_persist:
-            return
         if self.n_docs == 0 and not self.vocab:
             self.index_file.unlink(missing_ok=True)
-            self._needs_persist = False
             return
         try:
             self.index_file.parent.mkdir(parents=True, exist_ok=True)
@@ -374,7 +363,6 @@ class BM25Index(BaseKeywordIndex):
             with open(tmp, "wb") as f:
                 pickle.dump(self._snapshot(), f)
             tmp.replace(self.index_file)
-            self._needs_persist = False
             self.logger.info(f"Saved {self.n_docs} docs to {self.index_file}")
         except Exception as e:
             self.logger.exception(f"Failed to write {self.index_file}: {e}")
@@ -388,7 +376,6 @@ class BM25Index(BaseKeywordIndex):
             with open(self.index_file, "rb") as f:
                 data = pickle.load(f)
             self._restore(data)
-            self._needs_persist = False
             self.logger.info(f"Loaded {self.n_docs} docs from {self.index_file}")
         except Exception as e:
             self.logger.exception(f"Failed to load index: {e}")
@@ -407,7 +394,6 @@ class BM25Index(BaseKeywordIndex):
         self._posting_tfs = {}
         self._idf_cache = {}
         self.index_file.unlink(missing_ok=True)
-        self._needs_persist = False
 
     # -- Compaction -----------------------------------------------------------
 
@@ -493,4 +479,3 @@ class BM25Index(BaseKeywordIndex):
         self._posting_doc_idxs = new_posting_idxs
         self._posting_tfs = new_posting_tfs
         self._idf_cache = {}
-        self._needs_persist = True

tests/unit/test_file_graph.py

@@ -329,6 +329,68 @@ def test_rebuild_links_idempotent(backend_cls):
     asyncio.run(run())
 
 
+def test_nx_rebuild_links_repair_survives_restart():
+    """A repaired NetworkX snapshot must not revert to stale serialized edges."""
+
+    async def run():
+        with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
+            graph = NxFileGraph()
+            await graph.start()
+            await graph.upsert_nodes(
+                [
+                    make_node("a.md", [("c.md", None)]),
+                    make_node("b.md"),
+                    make_node("c.md"),
+                ],
+            )
+            await graph.dump()
+
+            # Model an old/torn snapshot: the node payload says a -> b while
+            # the separately serialized NetworkX edge still says a -> c.
+            graph._graph.nodes["a.md"]["node"] = make_node("a.md", [("b.md", None)])
+            await graph.dump()
+            await graph.close()
+
+            repaired = NxFileGraph()
+            await repaired.start()
+            assert {link.target_path for link in await repaired.get_outlinks("a.md")} == {"c.md"}
+            await repaired.rebuild_links()
+            assert {link.target_path for link in await repaired.get_outlinks("a.md")} == {"b.md"}
+            await repaired.close()
+
+            reopened = NxFileGraph()
+            await reopened.start()
+            assert {link.target_path for link in await reopened.get_outlinks("a.md")} == {"b.md"}
+            await reopened.close()
+
+    asyncio.run(run())
+
+
+def test_local_load_merged_state_survives_restart():
+    """Loading a snapshot into a modified local graph must not suppress its later dump."""
+
+    async def run():
+        with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
+            seed = LocalFileGraph()
+            await seed.start()
+            await seed.upsert_nodes([make_node("a.md")])
+            await seed.close()
+
+            merged = LocalFileGraph()
+            await merged.start()
+            await merged.upsert_nodes([make_node("b.md")])
+            await merged.load()
+            assert {node.path for node in await merged.get_nodes()} == {"a.md", "b.md"}
+            await merged.close()
+
+            reopened = LocalFileGraph()
+            await reopened.start()
+            assert {node.path for node in await reopened.get_nodes()} == {"a.md", "b.md"}
+            await reopened.close()
+
+    asyncio.run(run())
+
+
 @pytest.mark.parametrize("backend_cls", BACKENDS)
 def test_persistence_roundtrip(backend_cls):
     """close() dumps; a fresh instance loads the same nodes from disk."""

tests/unit/test_file_store_consistency.py

@@ -285,8 +285,8 @@ def test_keyword_only_upsert_removes_old_chunks_and_docs():
     run(go())
 
 
-def test_close_does_not_rewrite_keyword_index_and_graph_after_store_dump(monkeypatch):
-    """A store cascade leaves child components clean for their later close()."""
+def test_close_persists_each_component_once(monkeypatch):
+    """Store and owned dependency shutdown each persist their own state once."""
     bm25_writes = 0
     graph_writes = 0
     real_pickle_dump = bm25_index_module.pickle.dump
@@ -1433,7 +1433,7 @@ def test_faiss_rejects_stale_sidecar_after_partial_dump():
             # t2: simulate a crash between the two writes in dump(): only the
             # parent's JSONL write lands; the sidecar stays at the alpha
             # generation. (No close() -- the process is presumed dead.)
-            await LocalFileStore.dump(store_a)
+            await LocalFileStore._dump_owned_state(store_a)
 
             # t3: restart. The stale sidecar must be rejected by the digest.
             store_b = _new_faiss_store("t_faiss_stale_sidecar")

tests/unit/test_keyword_index.py

@@ -666,6 +666,28 @@ def test_dump_load_roundtrip_preserves_state():
     run(go())
 
 
+def test_runtime_parameter_update_survives_restart():
+    """Persist k1/b updates made after loading an existing index."""
+
+    async def go():
+        with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
+            seed = await create_bm25()
+            await seed.add_docs({"d1": "hello world"})
+            await seed.close()
+
+            changed = await create_bm25()
+            assert (changed.k1, changed.b) == (1.5, 0.75)
+            changed.k1 = 2.0
+            changed.b = 0.4
+            await changed.close()
+
+            reopened = await create_bm25()
+            assert (reopened.k1, reopened.b) == (2.0, 0.4)
+            await reopened.close()
+
+    run(go())
+
+
 def test_load_missing_file_keeps_empty_state():
     """Calling load() with no file on disk is a no-op."""