From 2b2de5a61fc491ea42e0f87090600e40823e6814 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 26 Aug 2026 13:03:47 +0300
Subject: [PATCH] Reuse the database a borrowed client already holds

`clients_for` hands back a client over a database the covering one owns, so
`async with` on it opened a second session and assigned it, and teardown
declined to close what this client did not open. Entry returns the client as
it stands.
---
 haiku_rag_slim/haiku/rag/client/__init__.py | 12 ++++++++++++
 tests/test_multi_db.py                      | 21 +++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py
index e5e2730b..90a114af 100644
--- a/haiku_rag_slim/haiku/rag/client/__init__.py
+++ b/haiku_rag_slim/haiku/rag/client/__init__.py
@@ -295,11 +295,19 @@ class HaikuRAG:
     async def __aenter__(self):
         """Async context manager entry — initializes store and repositories.
 
+        A client borrowing a database is already open and returns itself.
+        Opening a second session would leak it: teardown declines to close what
+        this client did not open.
+
         A client covering several databases opens none of them here: which are
         searched is a per-query choice, so they open on first use. `store` and the
         repositories stay unset in that case, since they have no unambiguous
         meaning across a set.
         """
+        if not self._owns_session:
+            assert self._session is not None
+            return self
+
         scope = self._resolve_scope()
         if scope.covers_multiple:
             if self._create:
@@ -338,6 +346,10 @@ class HaikuRAG:
         Opening is per query rather than at entry: a set of 25 configured
         databases is typically queried a few at a time, and a database nobody
         asked for must not be able to fail a query, or be opened for nothing.
+
+        The clients returned borrow their databases from this one and are valid
+        only while it is open. Closing one, or entering it as a context manager,
+        leaves the database alone; this client closes them all on teardown.
         """
         assert isinstance(self._session, FederatedSession)
         names = _without_repeats(names)
diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py
index 51f68457..92365b41 100644
--- a/tests/test_multi_db.py
+++ b/tests/test_multi_db.py
@@ -598,6 +598,27 @@ class TestBorrowedDatabases:
         assert {r.source for r in results} == {"alpha", "beta"}
         assert not store.db.is_open(), "the set left a database open"
 
+    @pytest.mark.asyncio
+    async def test_entering_a_borrowed_client_reuses_its_database(self, tmp_path):
+        """`async with` on a borrowed client is a plausible thing to write.
+        Opening a second session would leak it, since teardown declines to close
+        what this client did not open."""
+        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"])
+            borrowed = alpha.store
+
+            async with alpha as entered:
+                assert entered is alpha
+                assert alpha.store is borrowed, "entry opened a second database"
+
+            assert borrowed.db.is_open(), "exit closed a database it borrowed"
+            assert alpha.store is borrowed
+
+        assert not borrowed.db.is_open(), "the set left a database open"
+
     @pytest.mark.asyncio
     async def test_a_borrowed_client_releases_what_it_built(self, tmp_path):
         """Its reranker is its own; the database it wraps is not."""