diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 64642be1..181aa2ff 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/skills/_deps.py b/haiku_rag_slim/haiku/rag/skills/_deps.py index 1f8efb01..b8bf6e02 100644 --- a/haiku_rag_slim/haiku/rag/skills/_deps.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 20bc69d4..ecd1fe2c 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/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( - _require_rag(ctx), - query, - limit=limit, - document_filter=state.document_filter if state else None, - ) + async with _require_lock(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 +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( - _require_rag(ctx), - filter=state.document_filter if state else None, - ) + async with _require_lock(ctx): + return await skill_list_documents( + _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) - 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 _require_lock(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) diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index ed370f24..f8b33635 100644 --- a/tests/sandbox/test_sandbox.py +++ b/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 - used to cross two event loops no longer borrows the same LanceDB state. + 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: @@ -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 diff --git a/tests/skills/conftest.py b/tests/skills/conftest.py index 94776924..f72cb932 100644 --- a/tests/skills/conftest.py +++ b/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 diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py index e72eaec3..a583ba1e 100644 --- a/tests/skills/test_rag.py +++ b/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