Keep the sandbox's connection open for the owners it hands out

`_documents` took its owner clients inside an ephemeral connection and used
them after that connection closed, then stored them for the file reads that
follow, so every standalone multi-database read went through a closed database.

A sandbox covering a set now retains the client it opened until `close()`. One
database has no owners and keeps its connection no longer than the read that
opened it, so a write from elsewhere is still visible to the next read.
This commit is contained in:
Yiorgis Gozadinos 2026-08-27 11:40:43 +03:00
parent 9210642f65
commit c8a9b5df50
No known key found for this signature in database
2 changed files with 40 additions and 0 deletions

View file

@ -149,6 +149,7 @@ class Sandbox:
_doc_chunk_index: dict[str, dict[str, list[str]]]
_items_jsonl_cache: dict[str, str]
_toc_json_cache: dict[str, str]
_opened: "HaikuRAG | None"
_pool: AsyncMonty | None
_session: AsyncMontySession | None
_vfs: OSAccess | None
@ -206,6 +207,7 @@ class Sandbox:
self._config = config
self._context = context
self._rag = rag
self._opened = None
self._owners = {}
self._lock = lock
self._search_results = []
@ -239,11 +241,23 @@ class Sandbox:
else:
yield connection
return
if self._scope.covers_multiple:
yield await self._open_connection()
return
from haiku.rag.client import HaikuRAG
async with HaikuRAG._covering(self._scope, self._config, read_only=True) as rag:
yield rag
async def _open_connection(self) -> "HaikuRAG":
"""Open and retain a federated client for owner-backed reads."""
if self._opened is None:
from haiku.rag.client import HaikuRAG
self._opened = HaikuRAG._covering(self._scope, self._config, read_only=True)
await self._opened.__aenter__()
return self._opened
async def _documents(self) -> "tuple[list[Any], dict[str, HaikuRAG]]":
"""Every document in scope, and the client holding each of them.
@ -338,6 +352,9 @@ class Sandbox:
if self._pool is not None:
await self._pool.__aexit__(None, None, None)
self._pool = None
if self._opened is not None:
await self._opened.__aexit__(None, None, None)
self._opened = None
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""

View file

@ -20,6 +20,29 @@ async def _mounted(rag, sources=None):
return sandbox, docs, owners
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
that produced them."""
@pytest.mark.asyncio
async def test_owners_stay_open_for_later_reads(self, tmp_path):
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
sandbox = Sandbox(db_path=None, config=config, context=AnalysisContext())
try:
_, owners = await sandbox._documents()
assert len(owners) == 2
assert all(owner.store.db.is_open() for owner in owners.values())
finally:
await sandbox.close()
assert not any(owner.store.db.is_open() for owner in owners.values())
class TestDocumentsAcrossDatabases:
@pytest.mark.asyncio
async def test_the_corpus_covers_every_configured_database(self, tmp_path):