Serialize the sandbox's shared connection, not its owners

The lock was applied to an owner as well, so every owner-backed file read queued
behind the capability's tool calls to guard state it does not touch. An owner is
a session of its own and is yielded straight through, which is what the
docstring already claimed.
This commit is contained in:
Yiorgis Gozadinos 2026-08-28 08:41:51 +03:00
parent bd8c1a6d15
commit b5aa0e7122
No known key found for this signature in database
2 changed files with 54 additions and 4 deletions

View file

@ -237,13 +237,15 @@ class Sandbox:
no repositories of its own, so those reads have to name their owner. An
owner is a session of its own, so it is yielded unserialized.
"""
connection = owner if owner is not None else self._rag
if connection is not None:
if owner is not None:
yield owner
return
if self._rag is not None:
if self._lock is not None:
async with self._lock:
yield connection
yield self._rag
else:
yield connection
yield self._rag
return
if self._scope.covers_multiple:
yield await self._open_connection()

View file

@ -20,6 +20,54 @@ async def _mounted(rag, sources=None):
return sandbox, docs, owners
class TestSerializingTheConnection:
"""The lock guards the shared connection, which the capability's own tool
calls also hold. An owner is a session of its own."""
@staticmethod
def _sandbox(rag, lock):
return Sandbox._covering(
rag._resolve_scope(), rag._config, AnalysisContext(), rag, lock
)
@pytest.mark.asyncio
async def test_the_shared_connection_is_serialized(self, tmp_path):
import asyncio
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
lock = asyncio.Lock()
async with HaikuRAG(config=config) as rag:
sandbox = self._sandbox(rag, lock)
async with sandbox._connection():
assert lock.locked()
assert not lock.locked()
@pytest.mark.asyncio
async def test_an_owner_is_not(self, tmp_path):
"""Serializing owner reads would queue every database's file read behind
the capability's searches, to guard state none of them touch."""
import asyncio
class Trap(asyncio.Lock):
"""Refuses rather than waits: holding a real lock would wedge the
suite on a regression instead of failing it."""
async def acquire(self):
raise AssertionError("serialized a read on an owner's own session")
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
async with HaikuRAG(config=config) as rag:
(alpha,) = await rag.clients_for(["alpha"])
sandbox = self._sandbox(rag, Trap())
async with sandbox._connection(alpha) as connection:
assert connection is alpha
class TestStandaloneAcrossDatabases:
"""Without a lent client the sandbox opens its own. The owners it hands out
are stored for later file reads, so that connection has to outlive the call