Merge pull request #425 from ggozad/fix/sandbox-borrowed-bug

Fix analysis-sandbox "Already borrowed" crash under concurrent tool calls

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Fixed
+
+- Skill tools (`search`/`cite`/`list_documents`/`get_document`) and the analysis sandbox serialize access to the shared LanceDB connection through one lock, so a turn's concurrently executed tool calls no longer trigger `RuntimeError: Already borrowed`.
+
 ## [0.55.0] - 2026-06-05
 
 ### Added

haiku_rag_slim/haiku/rag/sandbox/sandbox.py

@@ -1,9 +1,8 @@
 import asyncio
 import json
 import os
-import threading
-import weakref
-from collections.abc import Callable, Coroutine
+from collections.abc import AsyncIterator, Callable, Coroutine
+from contextlib import asynccontextmanager
 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,20 @@ 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
-    single read-only connection held for the sandbox's lifetime. Monty's file
-    callbacks are synchronous and cannot ``await``, so they bridge to that loop
-    via ``run_coroutine_threadsafe``. Call ``close()`` to release the loop and
-    connection; if it is not called they are released when the instance is
-    garbage collected.
+    All database access runs on the event loop that drives ``execute()``. Monty's
+    file callbacks are synchronous and run on the interpreter's worker thread, so
+    they bridge back to that loop via ``run_coroutine_threadsafe``; the loop is
+    free during ``feed_run_async`` (the VM runs on the worker thread), so the
+    bridge does not deadlock. When a ``rag`` connection is supplied it is used
+    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"
+    _lock: "asyncio.Lock | None"
     _search_results: "list[SearchResult]"
     _doc_items: dict[str, list["DocumentItem"]]
     _doc_chunk_index: dict[str, dict[str, list[str]]]
@@ -149,19 +142,20 @@ 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,
+        lock: "asyncio.Lock | None" = None,
     ):
         self._db_path = db_path
         self._config = config
         self._context = context
+        self._rag = rag
+        self._lock = lock
         self._search_results = []
         self._doc_items = {}
         self._doc_chunk_index = {}
@@ -170,64 +164,41 @@ 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:
-        """Start the background loop and open the read-only connection once."""
-        if self._loop is not None:
+    @asynccontextmanager
+    async def _connection(self) -> "AsyncIterator[HaikuRAG]":
+        """Yield the shared connection (serialized by the lock), or an ephemeral
+        read-only one. The lock guards the whole block so a read's awaits cannot
+        interleave with another task's operation on the same connection."""
+        if self._rag is not None:
+            if self._lock is not None:
+                async with self._lock:
+                    yield self._rag
+            else:
+                yield self._rag
             return
         from haiku.rag.client import HaikuRAG
 
-        loop = asyncio.new_event_loop()
-        thread = threading.Thread(
-            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)
+        async with HaikuRAG(self._db_path, config=self._config, read_only=True) as rag:
+            yield rag
 
     def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
-        """Run a coroutine on the background loop from a synchronous caller."""
-        self._ensure_vfs_loop()
-        assert self._loop is not None
+        """Run a coroutine on the execute() loop from a synchronous callback.
+
+        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.
-
-        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)
+        """Retained for API compatibility; the sandbox owns no resources."""
+        return
 
     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 +207,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 HaikuRAG(db_path, config=config, read_only=True) as rag:
+            async with self._connection() 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 +234,7 @@ class Sandbox:
             return out
 
         async def list_documents() -> list[dict[str, Any]]:
-            from haiku.rag.client import HaikuRAG
-
-            async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+            async with self._connection() as rag:
                 docs = await rag.list_documents(filter=context.filter)
             return [
                 {
@@ -293,16 +260,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 +279,8 @@ class Sandbox:
                 return cached
 
             async def _fetch() -> list[DocumentItem]:
-                assert sandbox._vfs_rag is not None
-                return await sandbox._vfs_rag.document_item_repository.get_all_items(
-                    did
-                )
+                async with sandbox._connection() as rag:
+                    return await rag.document_item_repository.get_all_items(did)
 
             items = sandbox._run_on_loop(_fetch())
             sandbox._doc_items[did] = items
@@ -332,11 +293,13 @@ class Sandbox:
                 return cached
 
             async def _fetch() -> dict[str, list[str]]:
-                assert sandbox._vfs_rag is not None
-                index = await sandbox._vfs_rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
-                    [did]
-                )
-                return index.get(did, {})
+                async with sandbox._connection() as rag:
+                    index = (
+                        await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
+                            [did]
+                        )
+                    )
+                    return index.get(did, {})
 
             chunk_index = sandbox._run_on_loop(_fetch())
             sandbox._doc_chunk_index[did] = chunk_index
@@ -414,11 +377,9 @@ class Sandbox:
             ) -> Callable[["PurePosixPath"], str]:
                 def read_content(_path: "PurePosixPath") -> str:
                     async def _fetch() -> str:
-                        assert sandbox._vfs_rag is not None
-                        content = (
-                            await sandbox._vfs_rag.document_repository.get_content(did)
-                        )
-                        return content or ""
+                        async with sandbox._connection() as rag:
+                            content = await rag.document_repository.get_content(did)
+                            return content or ""
 
                     return sandbox._run_on_loop(_fetch())
 
@@ -469,6 +430,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

@@ -1,6 +1,7 @@
+import asyncio
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -15,6 +16,11 @@ if TYPE_CHECKING:
 @dataclass
 class RAGRunDeps(SkillRunDeps):
     rag: "HaikuRAG | None" = None
+    # pydantic-ai runs a turn's tool calls concurrently; LanceDB's per-connection
+    # state cannot take two in-flight operations at once, so every use of ``rag``
+    # (skill tools and the analysis sandbox) serializes through this lock. Always
+    # present so serialization is never accidentally skipped.
+    rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
     search_count: int = 0
 
 
@@ -71,6 +77,8 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
                 db_path=db_path,
                 config=config,
                 context=AnalysisContext(filter=doc_filter),
+                rag=rag,
+                lock=deps.rag_lock,
             )
             deps.sandbox = sandbox
             _reset_invocation_state(deps.state)

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -1,3 +1,5 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any
 
@@ -84,6 +86,23 @@ def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
     return ctx.deps.rag
 
 
+@asynccontextmanager
+async def _serialized(ctx: RunContext[RAGRunDeps]) -> AsyncIterator[None]:
+    """Serialize access to the shared connection through the run's lock.
+
+    pydantic-ai runs a turn's tool calls concurrently and LanceDB's
+    per-connection state cannot take two in-flight operations at once. The lock
+    is always present (``RAGRunDeps`` creates one by default); the no-op branch
+    is a guard for a missing deps/lock.
+    """
+    lock = ctx.deps.rag_lock if ctx.deps is not None else None
+    if lock is None:
+        yield
+    else:
+        async with lock:
+            yield
+
+
 def _register_citations(state: Any, citations: "list[Citation]") -> None:
     """Add citations to the index and record cited chunk IDs for this invocation."""
     next_index = len(state.citation_index) + 1
@@ -195,12 +214,13 @@ def create_skill_tools(
                 )
 
             state = _get_state(ctx, state_type)
-            formatted, results = await skill_search(
-                _require_rag(ctx),
-                query,
-                limit=limit,
-                document_filter=state.document_filter if state else None,
-            )
+            async with _serialized(ctx):
+                formatted, results = await skill_search(
+                    _require_rag(ctx),
+                    query,
+                    limit=limit,
+                    document_filter=state.document_filter if state else None,
+                )
             if state:
                 state.searches[query] = results
 
@@ -221,10 +241,11 @@ def create_skill_tools(
         ) -> list[dict[str, Any]]:
             """List all documents in the knowledge base."""
             state = _get_state(ctx, state_type)
-            return await skill_list_documents(
-                _require_rag(ctx),
-                filter=state.document_filter if state else None,
-            )
+            async with _serialized(ctx):
+                return await skill_list_documents(
+                    _require_rag(ctx),
+                    filter=state.document_filter if state else None,
+                )
 
         tools["list_documents"] = list_documents
 
@@ -238,7 +259,8 @@ def create_skill_tools(
             Args:
                 query: Document ID, title, or URI to look up.
             """
-            return await skill_get_document(_require_rag(ctx), query)
+            async with _serialized(ctx):
+                return await skill_get_document(_require_rag(ctx), query)
 
         tools["get_document"] = get_document
 
@@ -325,24 +347,25 @@ def create_skill_tools(
             ]
 
             if missing:
-                rag = _require_rag(ctx)
-                synthetic: list[SearchResult] = []
-                doc_cache: dict[str, Any] = {}
-                for cid in missing:
-                    chunk = await rag.get_chunk_by_id(cid)
-                    if chunk is None or not chunk.document_id:
-                        continue
-                    did = chunk.document_id
-                    if did in doc_cache:
-                        doc = doc_cache[did]
-                    else:
-                        doc = await rag.get_document_by_id(did)
-                        doc_cache[did] = doc
-                    chunk.document_uri = doc.uri if doc else None
-                    chunk.document_title = doc.title if doc else None
-                    synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
-                if synthetic:
-                    citations.extend(resolve_citations(missing, synthetic))
+                async with _serialized(ctx):
+                    rag = _require_rag(ctx)
+                    synthetic: list[SearchResult] = []
+                    doc_cache: dict[str, Any] = {}
+                    for cid in missing:
+                        chunk = await rag.get_chunk_by_id(cid)
+                        if chunk is None or not chunk.document_id:
+                            continue
+                        did = chunk.document_id
+                        if did in doc_cache:
+                            doc = doc_cache[did]
+                        else:
+                            doc = await rag.get_document_by_id(did)
+                            doc_cache[did] = doc
+                        chunk.document_uri = doc.uri if doc else None
+                        chunk.document_title = doc.title if doc else None
+                        synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
+                    if synthetic:
+                        citations.extend(resolve_citations(missing, synthetic))
 
             if citations:
                 _register_citations(state, citations)

tests/sandbox/test_sandbox.py

@@ -1,3 +1,5 @@
+import asyncio
+import threading
 from pathlib import Path
 
 import pytest
@@ -455,8 +457,83 @@ class TestSandboxHeldConnection:
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_vfs_loop_and_connection_are_reused(self, temp_db_path):
-        """The background loop + read-only connection are created once and reused."""
+    async def test_vfs_reads_use_injected_connection(self, temp_db_path):
+        """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 serialize through the shared lock, so concurrent tasks never have
+        two operations in flight on the one connection at once.
+        """
+        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:
+            lock = asyncio.Lock()
+            sb = Sandbox(
+                db_path=temp_db_path,
+                config=config,
+                context=AnalysisContext(),
+                rag=rag,
+                lock=lock,
+            )
+            read_code = (
+                "from pathlib import Path\n"
+                f"print(Path('/documents/{doc.id}/content.txt').read_text())"
+            )
+
+            async def direct_read() -> str | None:
+                async with lock:
+                    return await rag.document_repository.get_content(doc.id)
+
+            exec_result, content = await asyncio.gather(
+                sb.execute(read_code),
+                direct_read(),
+            )
+            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 +546,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
-            loop, rag = sb._loop, sb._vfs_rag
-            assert loop is not None and rag is not None
-            assert (await sb.execute(read)).success
-            assert sb._loop is loop
-            assert sb._vfs_rag is rag
+            first = await sb.execute(read)
+            second = await sb.execute(read)
+            assert first.success and second.success
+            assert "Foxes" in first.stdout
+            assert first.stdout == second.stdout
+            assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
         finally:
             sb.close()
 
@@ -516,14 +592,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):
-        """close() tears down the background loop thread and is idempotent."""
+    async def test_close_is_idempotent_after_vfs_read(self, temp_db_path):
+        """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 +615,6 @@ class TestSandboxHeldConnection:
             f"print(Path('/documents/{doc.id}/content.txt').read_text())"
         )
         assert result.success
-        thread = sb._loop_thread
-        assert thread is not None and thread.is_alive()
+        assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
         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 = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/items.jsonl"))
+    raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
     return [json.loads(line) for line in raw.strip().splitlines()] if raw else []
 
 

tests/skills/test_rag.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 
 from haiku.rag.config.models import AppConfig
@@ -169,6 +171,38 @@ class TestSearchTool:
         assert isinstance(result, str)
         assert len(result) > 0
 
+    async def test_concurrent_searches_serialize_on_lock(
+        self, rag_db, rag_client, monkeypatch
+    ):
+        """Tool calls in a turn run concurrently; the shared lock keeps only one
+        connection operation in flight (pydantic-ai borrow guard)."""
+        from haiku.rag.skills import _tools
+        from haiku.rag.skills.rag import create_skill
+
+        original = _tools.skill_search
+        inflight = {"now": 0, "max": 0}
+
+        async def tracking_search(*args, **kwargs):
+            inflight["now"] += 1
+            inflight["max"] = max(inflight["max"], inflight["now"])
+            try:
+                await asyncio.sleep(0.05)  # widen the window an overlap would use
+                return await original(*args, **kwargs)
+            finally:
+                inflight["now"] -= 1
+
+        monkeypatch.setattr(_tools, "skill_search", tracking_search)
+
+        skill = create_skill(db_path=rag_db)
+        search = _get_tool(skill, "search")
+        ctx = _make_ctx(rag=rag_client)  # one ctx -> one shared lock, as in a turn
+        results = await asyncio.gather(
+            search(ctx, query="artificial intelligence"),
+            search(ctx, query="machine learning"),
+        )
+        assert all(isinstance(r, str) and r for r in results)
+        assert inflight["max"] == 1
+
     async def test_search_updates_state(self, rag_db, rag_client):
         from haiku.rag.skills.rag import RAGState, create_skill