Serialize shared-connection access across skill tools and the sandbox

haiku_rag_slim/haiku/rag/sandbox/sandbox.py

@@ -133,6 +133,7 @@ class Sandbox:
     _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]]]
@@ -148,11 +149,13 @@ class Sandbox:
         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 = {}
@@ -164,9 +167,15 @@ class Sandbox:
 
     @asynccontextmanager
     async def _connection(self) -> "AsyncIterator[HaikuRAG]":
-        """Yield the supplied connection, or an ephemeral read-only one."""
+        """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:
-            yield self._rag
+            if self._lock is not None:
+                async with self._lock:
+                    yield self._rag
+            else:
+                yield self._rag
             return
         from haiku.rag.client import HaikuRAG
 

haiku_rag_slim/haiku/rag/skills/_deps.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -15,6 +16,10 @@ 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.
+    rag_lock: "asyncio.Lock | None" = None
     search_count: int = 0
 
 
@@ -50,6 +55,7 @@ def make_rag_lifespan(db_path: Path, config: AppConfig):
 
         async with HaikuRAG(db_path, config=config, read_only=True) as rag:
             deps.rag = rag
+            deps.rag_lock = asyncio.Lock()
             deps.search_count = 0
             _reset_invocation_state(deps.state)
             yield
@@ -66,12 +72,14 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
         doc_filter = getattr(deps.state, "document_filter", None)
         async with HaikuRAG(db_path, config=config, read_only=True) as rag:
             deps.rag = rag
+            deps.rag_lock = asyncio.Lock()
             deps.search_count = 0
             sandbox = Sandbox(
                 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,4 @@
+import asyncio
 from pathlib import Path
 from typing import Any
 
@@ -84,6 +85,13 @@ def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
     return ctx.deps.rag
 
 
+def _require_lock(ctx: RunContext[RAGRunDeps]) -> asyncio.Lock:
+    assert ctx.deps is not None and ctx.deps.rag_lock is not None, (
+        "RAGRunDeps.rag_lock is not set — skill lifespan must run before tools."
+    )
+    return ctx.deps.rag_lock
+
+
 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 +203,13 @@ def create_skill_tools(
                 )
 
             state = _get_state(ctx, state_type)
-            formatted, results = await skill_search(
+            async with _require_lock(ctx):
-                _require_rag(ctx),
+                formatted, results = await skill_search(
-                query,
+                    _require_rag(ctx),
-                limit=limit,
+                    query,
-                document_filter=state.document_filter if state else None,
+                    limit=limit,
-            )
+                    document_filter=state.document_filter if state else None,
+                )
             if state:
                 state.searches[query] = results
 
@@ -221,10 +230,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(
+            async with _require_lock(ctx):
-                _require_rag(ctx),
+                return await skill_list_documents(
-                filter=state.document_filter if state else None,
+                    _require_rag(ctx),
-            )
+                    filter=state.document_filter if state else None,
+                )
 
         tools["list_documents"] = list_documents
 
@@ -238,7 +248,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 _require_lock(ctx):
+                return await skill_get_document(_require_rag(ctx), query)
 
         tools["get_document"] = get_document
 
@@ -325,24 +336,25 @@ def create_skill_tools(
             ]
 
             if missing:
-                rag = _require_rag(ctx)
+                async with _require_lock(ctx):
-                synthetic: list[SearchResult] = []
+                    rag = _require_rag(ctx)
-                doc_cache: dict[str, Any] = {}
+                    synthetic: list[SearchResult] = []
-                for cid in missing:
+                    doc_cache: dict[str, Any] = {}
-                    chunk = await rag.get_chunk_by_id(cid)
+                    for cid in missing:
-                    if chunk is None or not chunk.document_id:
+                        chunk = await rag.get_chunk_by_id(cid)
-                        continue
+                        if chunk is None or not chunk.document_id:
-                    did = chunk.document_id
+                            continue
-                    if did in doc_cache:
+                        did = chunk.document_id
-                        doc = doc_cache[did]
+                        if did in doc_cache:
-                    else:
+                            doc = doc_cache[did]
-                        doc = await rag.get_document_by_id(did)
+                        else:
-                        doc_cache[did] = doc
+                            doc = await rag.get_document_by_id(did)
-                    chunk.document_uri = doc.uri if doc else None
+                            doc_cache[did] = doc
-                    chunk.document_title = doc.title if doc else None
+                        chunk.document_uri = doc.uri if doc else None
-                    synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
+                        chunk.document_title = doc.title if doc else None
-                if synthetic:
+                        synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
-                    citations.extend(resolve_citations(missing, synthetic))
+                    if synthetic:
+                        citations.extend(resolve_citations(missing, synthetic))
 
             if citations:
                 _register_citations(state, citations)

tests/sandbox/test_sandbox.py

@@ -492,8 +492,8 @@ class TestSandboxHeldConnection:
     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
+        Both serialize through the shared lock, so concurrent tasks never have
-        used to cross two event loops no longer borrows the same LanceDB state.
+        two operations in flight on the one connection at once.
         """
         config = AppConfig()
         async with HaikuRAG(temp_db_path, create=True) as client:
@@ -505,19 +505,26 @@ class TestSandboxHeldConnection:
         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),
-                rag.document_repository.get_content(doc.id),
+                direct_read(),
             )
             assert exec_result.success, exec_result.stderr
             assert "Foxes" in exec_result.stdout

tests/skills/conftest.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 from unittest.mock import MagicMock
 
@@ -17,10 +18,11 @@ def _make_ctx(state=None, rag=None, sandbox=None):
     from haiku.rag.skills.analysis import AnalysisState
 
     ctx = MagicMock(spec=RunContext)
+    lock = asyncio.Lock()
     if isinstance(state, AnalysisState) or sandbox is not None:
-        ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox)
+        ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox, rag_lock=lock)
     else:
-        ctx.deps = RAGRunDeps(state=state, rag=rag)
+        ctx.deps = RAGRunDeps(state=state, rag=rag, rag_lock=lock)
     return ctx
 
 

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