From cfb880e8a00e33868d54f3a8cc423c489a4f6545 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 27 Aug 2026 11:40:24 +0300
Subject: [PATCH] Keep every database a cancelled fan-out opened

`sessions_for` recorded what opened from `gather`'s results, which arrive only
when it runs to completion. Cancelling it discarded them, so a database that
opened while a sibling was still pending was never recorded and `aclose` never
closed it. `return_exceptions=True` covers a failing child, not a cancelled
parent.

`_open` now registers its own session, so what opened is reachable however the
fan-out ends.
---
 haiku_rag_slim/haiku/rag/client/session.py | 23 ++++++-------
 tests/multi_db/test_lifecycle.py           | 40 +++++++++++++++++++---
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/client/session.py b/haiku_rag_slim/haiku/rag/client/session.py
index bfe9afb3..c0f7343c 100644
--- a/haiku_rag_slim/haiku/rag/client/session.py
+++ b/haiku_rag_slim/haiku/rag/client/session.py
@@ -288,25 +288,22 @@ class FederatedSession:
             missing = [name for name in names if name not in self._sessions]
             if missing:
                 opened = await asyncio.gather(
-                    *(self._open(self._refs[name]) for name in missing),
+                    *(self._open(name) for name in missing),
                     return_exceptions=True,
                 )
-                # Whatever opened is tracked before the failure is reported, so
-                # teardown closes it: `gather` does not cancel the siblings of the
-                # one that raised, and an untracked connection leaks.
-                failure: BaseException | None = None
-                for name, result in zip(missing, opened, strict=True):
+                for result in opened:
                     if isinstance(result, BaseException):
-                        failure = failure or result
-                    else:
-                        self._sessions[name] = result
-                if failure is not None:
-                    raise failure
+                        raise result
         return [self._sessions[name] for name in names]
 
-    async def _open(self, ref: DatabaseRef) -> SingleDatabaseSession:
+    async def _open(self, name: str) -> None:
+        """Open and register one database before returning to the fan-out.
+
+        Registered here because a cancelled `gather` discards its results.
+        """
+        ref = self._refs[name]
         one, db_path = ref.connection(self._config)
-        return await SingleDatabaseSession(
+        self._sessions[name] = await SingleDatabaseSession(
             db_path if db_path is not None else default_db_path(one),
             one,
             skip_validation=self._skip_validation,
diff --git a/tests/multi_db/test_lifecycle.py b/tests/multi_db/test_lifecycle.py
index c179a837..2cab37f1 100644
--- a/tests/multi_db/test_lifecycle.py
+++ b/tests/multi_db/test_lifecycle.py
@@ -35,11 +35,11 @@ class TestOpeningDatabases:
             barrier = asyncio.Barrier(len(names))
             open_one = rag._session._open
 
-            async def gated(ref):
+            async def gated(name):
                 # Every open has to be in flight before any of them finishes, so
                 # a serial loop cannot get past this and the wait times out.
                 await barrier.wait()
-                return await open_one(ref)
+                await open_one(name)
 
             rag._session._open = gated
             clients = await asyncio.wait_for(rag.clients_for(names), timeout=15)
@@ -61,6 +61,38 @@ class TestOpeningDatabases:
             assert isinstance(rag._session, FederatedSession)
             assert set(rag._session._sessions) == {"alpha"}
 
+    @pytest.mark.asyncio
+    async def test_a_cancelled_open_does_not_leak_the_ones_that_worked(self, tmp_path):
+        """Cancellation discards the fan-out's results rather than returning them,
+        so a database that opened while a sibling was still pending is reachable
+        only because the opener recorded it."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            assert isinstance(rag._session, FederatedSession)
+            open_one = rag._session._open
+            alpha_open = asyncio.Event()
+
+            async def staged(name):
+                if name == "beta":
+                    await asyncio.sleep(60)
+                await open_one(name)
+                alpha_open.set()
+
+            rag._session._open = staged
+            fanout = asyncio.create_task(rag.clients_for(["alpha", "beta"]))
+            await asyncio.wait_for(alpha_open.wait(), timeout=15)
+            fanout.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await fanout
+
+            assert set(rag._session._sessions) == {"alpha"}
+            alpha = rag._session._sessions["alpha"]
+
+        assert not alpha.store.db.is_open()
+
     @pytest.mark.asyncio
     async def test_a_database_named_twice_is_opened_once(self, tmp_path):
         """Fusion would count a repeated database as two rank lists."""
@@ -392,8 +424,8 @@ class TestDatabaseIndependentWork:
 
         opened: list[str] = []
 
-        async def refuse(self, ref):
-            opened.append(ref.name)
+        async def refuse(self, name):
+            opened.append(name)
             raise AssertionError("opened a database to chunk a document")
 
         monkeypatch.setattr(FederatedSession, "_open", refuse)