Always create the shared-connection lock so serialization can't be skipped

This commit is contained in:
Yiorgis Gozadinos 2026-06-06 10:51:07 +03:00
parent 458f365e29
commit 4450c1c908
No known key found for this signature in database
3 changed files with 9 additions and 12 deletions

View file

@ -1,7 +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
@ -18,8 +18,9 @@ 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
# (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
@ -55,7 +56,6 @@ 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
@ -72,7 +72,6 @@ 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,

View file

@ -91,9 +91,9 @@ 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. When no
lock is present (tools invoked directly, without a skill lifespan) there is
no concurrency to guard, so this is a no-op.
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:

View file

@ -1,4 +1,3 @@
import asyncio
import random
from unittest.mock import MagicMock
@ -18,11 +17,10 @@ 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, rag_lock=lock)
ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox)
else:
ctx.deps = RAGRunDeps(state=state, rag=rag, rag_lock=lock)
ctx.deps = RAGRunDeps(state=state, rag=rag)
return ctx