From c8a9b5df502934f64c0e1ad85c55fb3e8592d20d Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 27 Aug 2026 11:40:43 +0300
Subject: [PATCH] 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.
---
 haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 17 +++++++++++++++
 tests/sandbox/test_sandbox_multi_db.py      | 23 +++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
index e7ab6a2a..58c9bb61 100644
--- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
+++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
@@ -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."""
diff --git a/tests/sandbox/test_sandbox_multi_db.py b/tests/sandbox/test_sandbox_multi_db.py
index 2ae69a7c..623abec8 100644
--- a/tests/sandbox/test_sandbox_multi_db.py
+++ b/tests/sandbox/test_sandbox_multi_db.py
@@ -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):