Run analysis sandbox VFS reads on the calling loop via the skill connection

haiku_rag_slim/haiku/rag/sandbox/sandbox.py

@@ -1,9 +1,8 @@
 import asyncio
 import json
 import os
-import threading
+from collections.abc import AsyncIterator, Callable, Coroutine
-import weakref
+from contextlib import asynccontextmanager
-from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
@@ -31,15 +30,6 @@ class SandboxResult:
     success: bool
 
 
-def _stop_loop(loop: asyncio.AbstractEventLoop, thread: threading.Thread) -> None:
-    """Stop a background loop and join its thread. Safe if already stopped."""
-    if loop.is_closed():
-        return
-    loop.call_soon_threadsafe(loop.stop)
-    thread.join()
-    loop.close()
-
-
 def _build_toc(
     items: list["DocumentItem"],
     chunk_index: dict[str, list[str]],
@@ -130,17 +120,19 @@ class Sandbox:
         result = await sandbox.execute("x = await search('query')")
         result = await sandbox.execute("print(x[0]['content'])")  # x persists
 
-    The virtual filesystem reads run on a dedicated background event loop with a
+    All database access runs on the event loop that drives ``execute()``. Monty's
-    single read-only connection held for the sandbox's lifetime. Monty's file
+    file callbacks are synchronous and run on the interpreter's worker thread, so
-    callbacks are synchronous and cannot ``await``, so they bridge to that loop
+    they bridge back to that loop via ``run_coroutine_threadsafe``; the loop is
-    via ``run_coroutine_threadsafe``. Call ``close()`` to release the loop and
+    free during ``feed_run_async`` (the VM runs on the worker thread), so the
-    connection; if it is not called they are released when the instance is
+    bridge does not deadlock. When a ``rag`` connection is supplied it is used
-    garbage collected.
+    for every read, so an analysis run drives a single connection on a single
+    loop; otherwise each read opens an ephemeral read-only connection.
     """
 
     _db_path: Path
     _config: AppConfig
     _context: AnalysisContext
+    _rag: "HaikuRAG | None"
     _search_results: "list[SearchResult]"
     _doc_items: dict[str, list["DocumentItem"]]
     _doc_chunk_index: dict[str, dict[str, list[str]]]
@@ -149,19 +141,18 @@ class Sandbox:
     _repl: MontyRepl | None
     _vfs: OSAccess | None
     _loop: asyncio.AbstractEventLoop | None
-    _loop_thread: threading.Thread | None
-    _vfs_rag: "HaikuRAG | None"
-    _finalizer: weakref.finalize | None
 
     def __init__(
         self,
         db_path: Path,
         config: AppConfig,
         context: AnalysisContext,
+        rag: "HaikuRAG | None" = None,
     ):
         self._db_path = db_path
         self._config = config
         self._context = context
+        self._rag = rag
         self._search_results = []
         self._doc_items = {}
         self._doc_chunk_index = {}
@@ -170,64 +161,35 @@ class Sandbox:
         self._repl = None
         self._vfs = None
         self._loop = None
-        self._loop_thread = None
-        self._vfs_rag = None
-        self._finalizer = None
 
-    def _ensure_vfs_loop(self) -> None:
+    @asynccontextmanager
-        """Start the background loop and open the read-only connection once."""
+    async def _connection(self) -> "AsyncIterator[HaikuRAG]":
-        if self._loop is not None:
+        """Yield the supplied connection, or an ephemeral read-only one."""
+        if self._rag is not None:
+            yield self._rag
             return
         from haiku.rag.client import HaikuRAG
 
-        loop = asyncio.new_event_loop()
+        async with HaikuRAG(self._db_path, config=self._config, read_only=True) as rag:
-        thread = threading.Thread(
+            yield rag
-            target=loop.run_forever, daemon=True, name="sandbox-vfs"
-        )
-        thread.start()
-        self._loop = loop
-        self._loop_thread = thread
-
-        async def _open() -> "HaikuRAG":
-            rag = HaikuRAG(self._db_path, config=self._config, read_only=True)
-            await rag.__aenter__()
-            return rag
-
-        self._vfs_rag = self._run_on_loop(_open())
-        # GC backstop: capture only the loop + thread, never ``self`` (which
-        # would pin the instance and prevent collection).
-        self._finalizer = weakref.finalize(self, _stop_loop, loop, thread)
 
     def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
-        """Run a coroutine on the background loop from a synchronous caller."""
+        """Run a coroutine on the execute() loop from a synchronous callback.
-        self._ensure_vfs_loop()
+
-        assert self._loop is not None
+        Called from Monty's worker thread while ``feed_run_async`` leaves the
+        loop free, so scheduling onto it and blocking for the result is safe.
+        """
+        assert self._loop is not None, (
+            "VFS reads happen during execute(); the loop must be captured first."
+        )
         return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
 
     def close(self) -> None:
-        """Release the background loop and read-only connection.
+        """Retained for API compatibility; the sandbox owns no resources."""
-
+        return
-        Idempotent and safe to call when no VFS read ever started the loop.
-        """
-        loop, thread, rag = self._loop, self._loop_thread, self._vfs_rag
-        if loop is None or thread is None:
-            return
-        self._loop = None
-        self._loop_thread = None
-        self._vfs_rag = None
-        if self._finalizer is not None:
-            self._finalizer.detach()
-            self._finalizer = None
-        if rag is not None and not loop.is_closed():
-            asyncio.run_coroutine_threadsafe(
-                rag.__aexit__(None, None, None), loop
-            ).result()
-        _stop_loop(loop, thread)
 
     def _build_external_functions(self) -> dict[str, Any]:
         """Build async external functions for the Monty interpreter."""
-        db_path = self._db_path
-        config = self._config
         context = self._context
 
         async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
@@ -236,9 +198,7 @@ class Sandbox:
             # agent's Python can't do anything with them. The driving model
             # gets figures through the top-level `search` tool when the
             # question is visual; in-code search is for structural work.
-            from haiku.rag.client import HaikuRAG
+            async with self._connection() as rag:
-
-            async with HaikuRAG(db_path, config=config, read_only=True) as rag:
                 results = await rag.search(query, limit=limit, filter=context.filter)
                 expanded = await rag.expand_context(results)
             self._search_results.extend(expanded)
@@ -265,9 +225,7 @@ class Sandbox:
             return out
 
         async def list_documents() -> list[dict[str, Any]]:
-            from haiku.rag.client import HaikuRAG
+            async with self._connection() as rag:
-
-            async with HaikuRAG(db_path, config=config, read_only=True) as rag:
                 docs = await rag.list_documents(filter=context.filter)
             return [
                 {
@@ -293,16 +251,12 @@ class Sandbox:
         - items.jsonl: CallbackFile (lazy, bulk-cached)
         - toc.json: CallbackFile (lazy, bulk-cached)
         """
-        from haiku.rag.client import HaikuRAG
-
-        db_path = self._db_path
-        config = self._config
         files: list[MemoryFile | CallbackFile] = []
 
         def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None:
             raise PermissionError(f"Document files are read-only: {_path}")
 
-        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+        async with self._connection() as rag:
             docs = await rag.list_documents(filter=self._context.filter)
 
         doc_titles = {doc.id: doc.title for doc in docs if doc.id}
@@ -316,10 +270,8 @@ class Sandbox:
                 return cached
 
             async def _fetch() -> list[DocumentItem]:
-                assert sandbox._vfs_rag is not None
+                async with sandbox._connection() as rag:
-                return await sandbox._vfs_rag.document_item_repository.get_all_items(
+                    return await rag.document_item_repository.get_all_items(did)
-                    did
-                )
 
             items = sandbox._run_on_loop(_fetch())
             sandbox._doc_items[did] = items
@@ -332,11 +284,13 @@ class Sandbox:
                 return cached
 
             async def _fetch() -> dict[str, list[str]]:
-                assert sandbox._vfs_rag is not None
+                async with sandbox._connection() as rag:
-                index = await sandbox._vfs_rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
+                    index = (
-                    [did]
+                        await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
-                )
+                            [did]
-                return index.get(did, {})
+                        )
+                    )
+                    return index.get(did, {})
 
             chunk_index = sandbox._run_on_loop(_fetch())
             sandbox._doc_chunk_index[did] = chunk_index
@@ -414,11 +368,9 @@ class Sandbox:
             ) -> Callable[["PurePosixPath"], str]:
                 def read_content(_path: "PurePosixPath") -> str:
                     async def _fetch() -> str:
-                        assert sandbox._vfs_rag is not None
+                        async with sandbox._connection() as rag:
-                        content = (
+                            content = await rag.document_repository.get_content(did)
-                            await sandbox._vfs_rag.document_repository.get_content(did)
+                            return content or ""
-                        )
-                        return content or ""
 
                     return sandbox._run_on_loop(_fetch())
 
@@ -469,6 +421,8 @@ class Sandbox:
 
         Variables persist across calls within the same Sandbox instance.
         """
+        # Monty's synchronous file callbacks bridge DB reads back to this loop.
+        self._loop = asyncio.get_running_loop()
         repl, vfs = await self._ensure_initialized()
         external_fns = self._build_external_functions()
 

haiku_rag_slim/haiku/rag/skills/_deps.py

@@ -71,6 +71,7 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
                 db_path=db_path,
                 config=config,
                 context=AnalysisContext(filter=doc_filter),
+                rag=rag,
             )
             deps.sandbox = sandbox
             _reset_invocation_state(deps.state)

tests/sandbox/test_sandbox.py

@@ -1,3 +1,5 @@
+import asyncio
+import threading
 from pathlib import Path
 
 import pytest
@@ -455,8 +457,76 @@ class TestSandboxHeldConnection:
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_vfs_loop_and_connection_are_reused(self, temp_db_path):
+    async def test_vfs_reads_use_injected_connection(self, temp_db_path):
-        """The background loop + read-only connection are created once and reused."""
+        """An injected connection services VFS reads on the calling loop.
+
+        No dedicated background loop/thread is spawned: all DB access for the
+        sandbox runs on the loop driving execute(), through the one connection.
+        """
+        config = AppConfig()
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            doc = await client.create_document(
+                content="Foxes and dogs roam the quiet hills.",
+                uri="test://doc",
+                title="Doc",
+            )
+        assert doc.id
+
+        async with HaikuRAG(temp_db_path, config=config, read_only=True) as rag:
+            sb = Sandbox(
+                db_path=temp_db_path,
+                config=config,
+                context=AnalysisContext(),
+                rag=rag,
+            )
+            result = await sb.execute(
+                "from pathlib import Path\n"
+                f"print(Path('/documents/{doc.id}/content.txt').read_text())"
+            )
+            assert result.success, result.stderr
+            assert "Foxes" in result.stdout
+            assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
+
+    @pytest.mark.asyncio
+    @pytest.mark.vcr()
+    async def test_vfs_read_concurrent_with_connection_read(self, temp_db_path):
+        """A VFS read and a direct read on the shared connection run together.
+
+        Both go through one connection on one loop, so the concurrency that
+        used to cross two event loops no longer borrows the same LanceDB state.
+        """
+        config = AppConfig()
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            doc = await client.create_document(
+                content="Foxes and dogs roam the quiet hills.",
+                uri="test://doc",
+                title="Doc",
+            )
+        assert doc.id
+
+        async with HaikuRAG(temp_db_path, config=config, read_only=True) as rag:
+            sb = Sandbox(
+                db_path=temp_db_path,
+                config=config,
+                context=AnalysisContext(),
+                rag=rag,
+            )
+            read_code = (
+                "from pathlib import Path\n"
+                f"print(Path('/documents/{doc.id}/content.txt').read_text())"
+            )
+            exec_result, content = await asyncio.gather(
+                sb.execute(read_code),
+                rag.document_repository.get_content(doc.id),
+            )
+            assert exec_result.success, exec_result.stderr
+            assert "Foxes" in exec_result.stdout
+            assert content and "Foxes" in content
+
+    @pytest.mark.asyncio
+    @pytest.mark.vcr()
+    async def test_vfs_reads_repeatable_across_executes(self, temp_db_path):
+        """Repeated VFS reads across execute() calls return consistent content."""
         config = AppConfig()
         async with HaikuRAG(temp_db_path, create=True) as client:
             doc = await client.create_document(
@@ -469,17 +539,16 @@ class TestSandboxHeldConnection:
         context = AnalysisContext()
         sb = Sandbox(db_path=temp_db_path, config=config, context=context)
         try:
-            assert sb._loop is None and sb._vfs_rag is None
             read = (
                 "from pathlib import Path\n"
                 f"print(Path('/documents/{doc.id}/content.txt').read_text())"
             )
-            assert (await sb.execute(read)).success
+            first = await sb.execute(read)
-            loop, rag = sb._loop, sb._vfs_rag
+            second = await sb.execute(read)
-            assert loop is not None and rag is not None
+            assert first.success and second.success
-            assert (await sb.execute(read)).success
+            assert "Foxes" in first.stdout
-            assert sb._loop is loop
+            assert first.stdout == second.stdout
-            assert sb._vfs_rag is rag
+            assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
         finally:
             sb.close()
 
@@ -516,14 +585,13 @@ class TestSandboxHeldConnection:
         async with HaikuRAG(temp_db_path, create=True):
             config = AppConfig()
             sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
-            assert sb._loop is None
             sb.close()
             sb.close()
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_close_stops_loop_thread(self, temp_db_path):
+    async def test_close_is_idempotent_after_vfs_read(self, temp_db_path):
-        """close() tears down the background loop thread and is idempotent."""
+        """close() is a safe no-op after a VFS read; no background thread lingers."""
         config = AppConfig()
         async with HaikuRAG(temp_db_path, create=True) as client:
             doc = await client.create_document(
@@ -540,8 +608,6 @@ class TestSandboxHeldConnection:
             f"print(Path('/documents/{doc.id}/content.txt').read_text())"
         )
         assert result.success
-        thread = sb._loop_thread
+        assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
-        assert thread is not None and thread.is_alive()
         sb.close()
-        assert not thread.is_alive()
         sb.close()

tests/sandbox/test_sandbox_toc.py

@@ -7,6 +7,7 @@ list of siblings; HTML/markdown corpora with real heading hierarchy get a
 nested tree. Items with no section_header at all produce `tree: []`.
 """
 
+import asyncio
 import json
 from pathlib import PurePosixPath
 
@@ -56,15 +57,21 @@ def _header(
     )
 
 
-async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict:
+async def _read_vfs_text(sandbox: Sandbox, path: str) -> str:
+    """Read a VFS file the way execute() does: the synchronous reader runs on a
+    worker thread and bridges DB access back to the (free) calling loop."""
     vfs = await sandbox._build_vfs()
-    raw = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/toc.json"))
+    sandbox._loop = asyncio.get_running_loop()
+    return await asyncio.to_thread(vfs.path_read_text, PurePosixPath(path))
+
+
+async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict:
+    raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/toc.json")
     return json.loads(raw)
 
 
 async def _read_items_jsonl(sandbox: Sandbox, doc_id: str) -> list[dict]:
-    vfs = await sandbox._build_vfs()
+    raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
-    raw = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/items.jsonl"))
     return [json.loads(line) for line in raw.strip().splitlines()] if raw else []