fix: keep checkpoint lifecycle work off the event loop

reme/application.py

@@ -15,6 +15,7 @@
 from .plugin import resolve_plugin_runtime
 from .schema import ComponentConfig, Response, StreamChunk
 from .utils import execute_stream_task, print_logo, get_logger
+from .utils.async_utils import complete_in_thread
 
 T = TypeVar("T", bound=BaseComponent)
 _NodeKey = tuple[str, str]
@@ -229,8 +230,9 @@ async def _close_started_components(self) -> None:
                 self.logger.exception(f"Failed to close {component_type_name(c.component_type)}:{c.name}: {e}")
         self._started_components.clear()
         if self.context.thread_pool is not None:
-            self.context.thread_pool.shutdown(wait=True)
+            thread_pool = self.context.thread_pool
             self.context.thread_pool = None
+            await complete_in_thread(thread_pool.shutdown, True)
 
     async def update_component(self, component_enum: ComponentType, name: str, /, **kwargs) -> BaseComponent:
         """Update an existing component by type/name; never creates missing components."""

reme/components/file_catalog/local_file_catalog.py

@@ -5,6 +5,7 @@
 from .base_file_catalog import BaseFileCatalog
 from ..component_registry import R
 from ...schema import FileNode
+from ...utils.async_utils import complete_in_thread
 from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst
 
 
@@ -24,12 +25,18 @@ async def load(self) -> None:
         async with self._io_lock:
             if not self._catalog_file.exists():
                 return
-            await self._read_jsonl()
+            loaded = await complete_in_thread(
+                self._read_jsonl_sync,
+                self._catalog_file,
+                self.encoding,
+                self._nodes,
+            )
+            self._nodes = loaded
             self.logger.debug(f"Loaded {len(self._nodes)} nodes from {self._catalog_file}")
 
     async def dump(self) -> None:
         async with self._io_lock:
-            await self._write_jsonl()
+            await complete_in_thread(self._write_jsonl_sync)
             self.logger.info(f"Saved {len(self._nodes)} nodes to {self._catalog_file}")
 
     async def upsert(self, nodes: list[FileNode]) -> None:
@@ -49,11 +56,16 @@ async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]:
                 return list(self._nodes.values())
             return [self._nodes[p] for p in paths if p in self._nodes]
 
-    async def _read_jsonl(self) -> None:
-        for line in read_jsonl_zst(self._catalog_file, self.encoding):
+    @staticmethod
+    def _read_jsonl_sync(path, encoding: str, existing: dict[str, FileNode]) -> dict[str, FileNode]:
+        """Read, decompress, and parse a catalog checkpoint off-loop."""
+        nodes = dict(existing)
+        for line in read_jsonl_zst(path, encoding):
             if stripped := line.strip():
                 node = FileNode.model_validate_json(stripped)
-                self._nodes[node.path] = node
+                nodes[node.path] = node
+        return nodes
 
-    async def _write_jsonl(self) -> None:
+    def _write_jsonl_sync(self) -> None:
+        """Serialize, compress, and atomically publish the locked catalog state."""
         write_jsonl_zst(self._catalog_file, (n.model_dump_json() for n in self._nodes.values()), self.encoding)

reme/components/file_graph/local_file_graph.py

@@ -1,11 +1,13 @@
 """Pure-Python file-graph backend (no external deps)."""
 
+import asyncio
 from pathlib import Path
 
 from .base_file_graph import BaseFileGraph
 from ..component_registry import R
 from ...enumeration import LinkScopeEnum
 from ...schema import FileLink, FileNode
+from ...utils.async_utils import complete_in_thread
 from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst
 
 
@@ -18,33 +20,64 @@ def __init__(self, **kwargs):
         self._nodes: dict[str, FileNode] = {}
         self._inverse: dict[str, set[str]] = {}  # real target → sources
         self._pending: dict[str, set[str]] = {}  # virtual target → sources
+        self._io_lock = asyncio.Lock()
         self._graph_file: Path = self.component_metadata_path / f"{self.name}.jsonl.zst"
 
     # -- Lifecycle ---------------------------------------------------------
 
     async def _start(self) -> None:
         self.component_metadata_path.mkdir(parents=True, exist_ok=True)
         await super()._start()  # base calls load()
-        await self.rebuild_links()
 
     async def load(self) -> None:
-        if not self._graph_file.exists():
-            return
-        try:
-            for line in read_jsonl_zst(self._graph_file):
-                if line.strip():
-                    node = FileNode.model_validate_json(line)
-                    self._nodes[node.path] = node
-            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 with self._io_lock:
+            if not self._graph_file.exists():
+                return
+            try:
+                nodes, inverse, pending = await complete_in_thread(
+                    self._load_sync,
+                    self._graph_file,
+                    self._nodes,
+                )
+                self._nodes = nodes
+                self._inverse = inverse
+                self._pending = pending
+                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:
-        try:
-            write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values()))
-            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}")
+        async with self._io_lock:
+            try:
+                await complete_in_thread(self._dump_sync)
+                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}")
+
+    @classmethod
+    def _load_sync(
+        cls,
+        path: Path,
+        existing: dict[str, FileNode],
+    ) -> tuple[dict[str, FileNode], dict[str, set[str]], dict[str, set[str]]]:
+        """Restore nodes and rebuild adjacency entirely outside the event loop."""
+        nodes = dict(existing)
+        for line in read_jsonl_zst(path):
+            if stripped := line.strip():
+                node = FileNode.model_validate_json(stripped)
+                nodes[node.path] = node
+
+        inverse: dict[str, set[str]] = {}
+        pending: dict[str, set[str]] = {}
+        for source, node in nodes.items():
+            for target in cls._targets(node):
+                bucket = inverse if target in nodes else pending
+                bucket.setdefault(target, set()).add(source)
+        return nodes, inverse, pending
+
+    def _dump_sync(self) -> None:
+        """Serialize and atomically publish the graph while its state is locked."""
+        write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values()))
 
     # -- Internals ---------------------------------------------------------
 
@@ -78,68 +111,89 @@ def _scope_match(self, target: str, scope: LinkScopeEnum | str) -> bool:
     # -- Node CRUD ---------------------------------------------------------
 
     async def upsert_nodes(self, nodes: list[FileNode]) -> None:
-        for node in nodes:
-            path = node.path
-            old = self._nodes.get(path)
-            if old is not None:
-                for target in self._targets(old):
-                    self._remove_edge(path, target)
-            self._nodes[path] = node
-            for target in self._targets(node):
-                self._add_edge(path, target)
-            promoted = self._pending.pop(path, None)
-            if promoted:
-                self._inverse.setdefault(path, set()).update(promoted)
+        async with self._io_lock:
+            for node in nodes:
+                path = node.path
+                old = self._nodes.get(path)
+                if old is not None:
+                    for target in self._targets(old):
+                        self._remove_edge(path, target)
+                self._nodes[path] = node
+                for target in self._targets(node):
+                    self._add_edge(path, target)
+                promoted = self._pending.pop(path, None)
+                if promoted:
+                    self._inverse.setdefault(path, set()).update(promoted)
 
     async def delete_nodes(self, paths: list[str]) -> None:
-        for path in paths:
-            node = self._nodes.pop(path, None)
-            if node is None:
-                continue
-            for target in self._targets(node):
-                self._remove_edge(path, target)
-            demoted = self._inverse.pop(path, None)
-            if demoted:
-                self._pending.setdefault(path, set()).update(demoted)
+        async with self._io_lock:
+            for path in paths:
+                node = self._nodes.pop(path, None)
+                if node is None:
+                    continue
+                for target in self._targets(node):
+                    self._remove_edge(path, target)
+                demoted = self._inverse.pop(path, None)
+                if demoted:
+                    self._pending.setdefault(path, set()).update(demoted)
 
     async def get_nodes(self, paths: list[str] | None = None) -> list[FileNode]:
-        if paths is None:
-            return list(self._nodes.values())
-        return [self._nodes[p] for p in paths if p in self._nodes]
+        async with self._io_lock:
+            if paths is None:
+                return list(self._nodes.values())
+            return [self._nodes[p] for p in paths if p in self._nodes]
 
     async def rebuild_links(self) -> None:
-        self._inverse.clear()
-        self._pending.clear()
-        for src, node in self._nodes.items():
-            for target in self._targets(node):
-                self._add_edge(src, target)
+        async with self._io_lock:
+            nodes, inverse, pending = await complete_in_thread(self._rebuild_links_sync, self._nodes)
+            self._nodes = nodes
+            self._inverse = inverse
+            self._pending = pending
+
+    @classmethod
+    def _rebuild_links_sync(
+        cls,
+        existing: dict[str, FileNode],
+    ) -> tuple[dict[str, FileNode], dict[str, set[str]], dict[str, set[str]]]:
+        """Rebuild adjacency from an in-memory node generation off-loop."""
+        nodes = dict(existing)
+        inverse: dict[str, set[str]] = {}
+        pending: dict[str, set[str]] = {}
+        for source, node in nodes.items():
+            for target in cls._targets(node):
+                bucket = inverse if target in nodes else pending
+                bucket.setdefault(target, set()).add(source)
+        return nodes, inverse, pending
 
     async def clear(self):
-        self._nodes.clear()
-        self._inverse.clear()
-        self._pending.clear()
-        self._graph_file.unlink(missing_ok=True)
+        async with self._io_lock:
+            self._nodes.clear()
+            self._inverse.clear()
+            self._pending.clear()
+            self._graph_file.unlink(missing_ok=True)
 
     # -- Link access -------------------------------------------------------
 
     async def get_outlinks(self, path: str, scope: LinkScopeEnum | str = LinkScopeEnum.REAL) -> list[FileLink]:
-        scope = self._normalize_scope(scope)
-        node = self._nodes.get(path)
-        if node is None:
-            return []
-        return [lnk for lnk in node.links if lnk.target_path and self._scope_match(lnk.target_path, scope)]
+        async with self._io_lock:
+            scope = self._normalize_scope(scope)
+            node = self._nodes.get(path)
+            if node is None:
+                return []
+            return [lnk for lnk in node.links if lnk.target_path and self._scope_match(lnk.target_path, scope)]
 
     async def get_inlinks(self, path: str, scope: LinkScopeEnum | str = LinkScopeEnum.REAL) -> list[FileLink]:
-        scope = self._normalize_scope(scope)
-        sources: set[str] = set()
-        if scope in (LinkScopeEnum.REAL, LinkScopeEnum.ALL):
-            sources |= self._inverse.get(path, set())
-        if scope in (LinkScopeEnum.VIRTUAL, LinkScopeEnum.ALL):
-            sources |= self._pending.get(path, set())
-        return [
-            link
-            for src in sorted(sources)
-            if src in self._nodes
-            for link in self._nodes[src].links
-            if link.target_path == path
-        ]
+        async with self._io_lock:
+            scope = self._normalize_scope(scope)
+            sources: set[str] = set()
+            if scope in (LinkScopeEnum.REAL, LinkScopeEnum.ALL):
+                sources |= self._inverse.get(path, set())
+            if scope in (LinkScopeEnum.VIRTUAL, LinkScopeEnum.ALL):
+                sources |= self._pending.get(path, set())
+            return [
+                link
+                for src in sorted(sources)
+                if src in self._nodes
+                for link in self._nodes[src].links
+                if link.target_path == path
+            ]