Ask the client what it covers

covers_multiple, source_names, source and reader_for replace the private
state seven modules were reading to work out how many databases they had.
The configured selection is kept intact, so entering a client twice derives
the same database rather than the last derivation.

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -127,7 +127,7 @@ def covers_several_databases(
     covers. Instructions follow coverage, not configuration.
     """
     if rag is not None:
-        return bool(rag._federated)
+        return rag.covers_multiple
     return db_path is None and len(config.lancedb.databases) > 1
 
 
@@ -543,7 +543,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
                     owner, chunk = found
                     if not chunk.document_id:
                         continue
-                    key = (owner._source, chunk.document_id)
+                    key = (owner.source, chunk.document_id)
                     if key not in documents:
                         documents[key] = await owner.get_document_by_id(
                             chunk.document_id
@@ -553,7 +553,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
                     chunk.document_title = document.title if document else None
                     chunk.document_meta = document.metadata if document else {}
                     result = SearchResult.from_chunk(chunk, score=1.0)
-                    result.source = owner._source
+                    result.source = owner.source
                     synthetic.append(result)
                 citations.extend(resolve_citations(missing, synthetic))
 

haiku_rag_slim/haiku/rag/chat/app.py

@@ -357,12 +357,10 @@ class ChatApp(App):
 
         citation = selected_widgets[0].citation
         # Chunks, pages and bounding boxes all come from the database holding the
-        # cited chunk. A client covering a set has no repositories of its own.
-        client = self.client
-        if client._federated:
-            if citation.source is None:
-                return
-            (client,) = await client.clients_for([citation.source])
+        # cited chunk, which a client covering a set has to be asked for.
+        client = await self.client.reader_for(citation.source)
+        if client is None:
+            return
         chunk_ids = citation.chunk_ids or [citation.chunk_id]
         chunks = []
         for cid in chunk_ids:

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -41,7 +41,7 @@ from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
 from haiku.rag.store.repositories.document_item import DocumentItemRepository
 from haiku.rag.store.repositories.settings import SettingsRepository
-from haiku.rag.utils import escape_sql_string, locate_database
+from haiku.rag.utils import escape_sql_string
 
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
@@ -128,7 +128,11 @@ class HaikuRAG:
                 None means all of them. Ignored when a single ``uri`` or an
                 explicit ``db_path`` is given.
         """
-        self._config = config if config is not None else get_config()
+        self._configured = config if config is not None else get_config()
+        # What the caller configured, kept intact: entering derives a
+        # single-database configuration from it, and asking again has to see the
+        # same set rather than the answer from last time.
+        self._config = self._configured
         self._db_path_given = db_path is not None
         if db_path is None:
             db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
@@ -139,20 +143,47 @@ class HaikuRAG:
         self._read_only = read_only
         self._requested_sources = sources
         self._clients: dict[str, HaikuRAG] = {}
-        self._source: str | None = None
         self._session: SingleDatabaseSession | FederatedSession | None = None
         self._owns_session = True
 
     @property
-    def _federated(self) -> dict[str, str]:
-        """The databases covered, name to location, empty while covering one.
+    def covers_multiple(self) -> bool:
+        """Whether this client reads from more than one database."""
+        return isinstance(self._session, FederatedSession)
 
-        Derived from the session rather than stored beside it: two answers to
-        "how many databases?" is what left every call site deciding for itself.
+    @property
+    def source_names(self) -> tuple[str, ...]:
+        """The configured databases this client covers, in configured order.
+
+        A single database contributes its own name, or nothing where the
+        configuration named none.
         """
         if isinstance(self._session, FederatedSession):
-            return self._session.locations
-        return {}
+            return self._session.names
+        return () if self.source is None else (self.source,)
+
+    @property
+    def source(self) -> str | None:
+        """The configured database this client reads, or None while covering a
+        set or reading a database the configuration did not name."""
+        if isinstance(self._session, SingleDatabaseSession):
+            return self._session.source
+        return None
+
+    async def reader_for(self, source: str | None) -> "HaikuRAG | None":
+        """The client that can read `source` — itself, where it reads one
+        database.
+
+        None where the database cannot be placed: a client covering a set needs
+        the name, and evidence recorded before databases could be named carries
+        none.
+        """
+        if not self.covers_multiple:
+            return self
+        if source is None:
+            return None
+        (owner,) = await self.clients_for([source])
+        return owner
 
     def _single_session(self, operation: str) -> SingleDatabaseSession:
         """The one database this operation works on.
@@ -163,7 +194,7 @@ class HaikuRAG:
         """
         if isinstance(self._session, SingleDatabaseSession):
             return self._session
-        covered = ", ".join(sorted(self._federated))
+        covered = ", ".join(sorted(self.source_names))
         raise AmbiguousDatabaseError(
             f"{operation} works on one database, and this client covers "
             f"{covered}; select one with clients_for([name])"
@@ -216,7 +247,7 @@ class HaikuRAG:
         covering a set has an unambiguous embedder without opening any of them.
         Built on first use and owned by this client, which closes it.
         """
-        if self._federated:
+        if self.covers_multiple:
             return get_embedder(config=self._config)
         return self.store.embedder
 
@@ -235,7 +266,7 @@ class HaikuRAG:
         Empty when the caller named a database itself: an explicit `db_path` says
         which one to open, so it is not overridden by a configured set.
         """
-        declared = self._config.lancedb.databases
+        declared = self._configured.lancedb.databases
         if not declared or self._db_path_given:
             return {}
         if self._requested_sources is not None and not self._requested_sources:
@@ -283,12 +314,12 @@ class HaikuRAG:
                 read_only=self._read_only,
             )
             return self
+        source: str | None = None
         if selected:
-            [(self._source, location)] = selected.items()
-            uri, db_path = locate_database(location)
-            self._config = self._config.model_copy(deep=True)
-            self._config.lancedb.databases = {}
-            self._config.lancedb.uri = uri
+            [(source, location)] = selected.items()
+            self._config, db_path = DatabaseRef.configured(source, location).connection(
+                self._configured
+            )
             if db_path is not None:
                 self._db_path = db_path
 
@@ -298,7 +329,7 @@ class HaikuRAG:
             skip_validation=self._skip_validation,
             create=self._create,
             read_only=self._read_only,
-            source=self._source,
+            source=source,
         ).open()
         return self
 
@@ -337,7 +368,6 @@ class HaikuRAG:
             session.db_path, config=session.config, read_only=session.read_only
         )
         client._session = session
-        client._source = session.source
         client._owns_session = False
         return client
 
@@ -354,7 +384,7 @@ class HaikuRAG:
         `SettingsRepository` reports on open.
         """
         recorded = [
-            (client._source, client.store.stored_embedding)
+            (client.source, client.store.stored_embedding)
             for client in clients
             if client.store.stored_embedding is not None
         ]
@@ -569,7 +599,7 @@ class HaikuRAG:
         Returns:
             The Document instance if found, None otherwise.
         """
-        if self._federated:
+        if self.covers_multiple:
             return await self._from_any_covered(
                 lambda owner: owner.get_document_by_id(document_id)
             )
@@ -586,7 +616,7 @@ class HaikuRAG:
         Returns:
             The Chunk instance if found, None otherwise.
         """
-        if self._federated:
+        if self.covers_multiple:
             return await self._from_any_covered(
                 lambda owner: owner.get_chunk_by_id(chunk_id)
             )
@@ -605,7 +635,7 @@ class HaikuRAG:
         Returns:
             The picture bytes if found, None otherwise.
         """
-        if not self._federated:
+        if not self.covers_multiple:
             return await self.document_item_repository.get_picture_bytes(
                 document_id, self_ref
             )
@@ -627,7 +657,7 @@ class HaikuRAG:
         Returns:
             The Document instance if found, None otherwise.
         """
-        if self._federated:
+        if self.covers_multiple:
             return await self._from_any_covered(
                 lambda owner: owner.get_document_by_uri(uri)
             )
@@ -685,7 +715,7 @@ class HaikuRAG:
         Returns:
             List of Document instances matching the criteria.
         """
-        if self._federated:
+        if self.covers_multiple:
             # Each database is asked for enough rows to satisfy the window, and
             # the window is applied to the merged listing: a limit means that
             # many documents in total, not that many per database.
@@ -718,7 +748,7 @@ class HaikuRAG:
         Returns:
             Number of documents matching the criteria.
         """
-        if self._federated:
+        if self.covers_multiple:
             counts = await asyncio.gather(
                 *(
                     owner.count_documents(filter=filter)
@@ -745,19 +775,19 @@ class HaikuRAG:
         `None` for all of them. Every read honouring `sources` decides through
         this, so the rule cannot differ between one operation and another.
         """
-        if self._federated:
+        if self.covers_multiple:
             return await self.clients_for(
-                list(self._federated) if sources is None else sources
+                list(self.source_names) if sources is None else sources
             )
         if sources is None:
             return [self]
         sources = _without_repeats(sources)
         if not sources:
             return []
-        if sources != [self._source]:
+        if sources != [self.source]:
             raise KeyError(
                 f"unknown database(s) {', '.join(sources) or '(none)'}; this "
-                f"client covers {self._source or 'a single unnamed database'}"
+                f"client covers {self.source or 'a single unnamed database'}"
             )
         return [self]
 
@@ -772,7 +802,7 @@ class HaikuRAG:
     ) -> list[SearchResult]:
         from haiku.rag.client.search import search, search_sources
 
-        if self._federated:
+        if self.covers_multiple:
             return await search_sources(
                 self, query, limit, search_type, filter, include_images, sources
             )
@@ -782,7 +812,7 @@ class HaikuRAG:
         # A database named in config keeps its name even when it is the only one
         # this client covers. Only a legacy single `uri` leaves source unset.
         for result in results:
-            result.source = self._source
+            result.source = self.source
         return results
 
     async def expand_context(

haiku_rag_slim/haiku/rag/client/agents.py

@@ -64,7 +64,7 @@ async def ask(
     from haiku.rag.utils import get_model
 
     capability = create_capability(
-        db_path=None if client._federated else client.store.db_path,
+        db_path=None if client.covers_multiple else client.store.db_path,
         config=client._config,
         rag=client,
         defer_loading=False,
@@ -124,7 +124,7 @@ async def analyze(
     from haiku.rag.utils import get_model
 
     capability = create_capability(
-        db_path=None if client._federated else client.store.db_path,
+        db_path=None if client.covers_multiple else client.store.db_path,
         config=client._config,
         rag=client,
         defer_loading=False,

haiku_rag_slim/haiku/rag/client/search.py

@@ -82,7 +82,7 @@ async def search_sources(
     if limit is None:
         limit = client._config.search.limit
 
-    names = list(client._federated) if sources is None else list(sources)
+    names = list(client.source_names) if sources is None else list(sources)
     if not names:
         return []
     selected = await client.clients_for(names)
@@ -107,7 +107,7 @@ async def search_sources(
     results: list[SearchResult] = []
     for owner, chunk, score in ranked:
         result = SearchResult.from_chunk(chunk, score)
-        result.source = owner._source
+        result.source = owner.source
         results.append(result)
     results = _dedup_picture_chunks(results)
 
@@ -116,10 +116,11 @@ async def search_sources(
         for result in results:
             if result.source:
                 by_owner.setdefault(result.source, []).append(result)
+        owners = await client.clients_for(list(by_owner))
         await asyncio.gather(
             *(
-                _populate_image_data(client._clients[name], owned)
-                for name, owned in by_owner.items()
+                _populate_image_data(owner, by_owner[name])
+                for name, owner in zip(by_owner, owners, strict=True)
             )
         )
 
@@ -430,7 +431,7 @@ async def expand_context(
     """
     # A federating client has no repositories of its own, so each result expands
     # through the database it came from.
-    if client._federated:
+    if client.covers_multiple:
         by_source: dict[str, list[SearchResult]] = {}
         unsourced: list[SearchResult] = []
         for result in search_results:
@@ -441,9 +442,9 @@ async def expand_context(
         owners = await client.clients_for(list(by_source))
         expanded_groups = await asyncio.gather(
             *(
-                expand_context(owner, by_source[owner._source])
+                expand_context(owner, by_source[owner.source])
                 for owner in owners
-                if owner._source
+                if owner.source
             )
         )
         merged = unsourced + [r for group in expanded_groups for r in group]

haiku_rag_slim/haiku/rag/client/session.py

@@ -280,9 +280,9 @@ class FederatedSession:
         self._lock = asyncio.Lock()
 
     @property
-    def locations(self) -> dict[str, str]:
-        """The databases covered, name to location as configured."""
-        return {name: ref.uri or str(ref.db_path) for name, ref in self._refs.items()}
+    def names(self) -> tuple[str, ...]:
+        """The databases covered, in configured order."""
+        return tuple(self._refs)
 
     async def sessions_for(self, names: list[str]) -> list[SingleDatabaseSession]:
         """The sessions for these databases, opening any not yet open.

haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py

@@ -22,7 +22,7 @@ def reported_database(client: "HaikuRAG", db_path: "Path | None") -> "Path | Non
     one named database opens it rather than covering a set. Only what the client
     ended up covering says which of the two this is.
     """
-    if client._federated:
+    if client.covers_multiple:
         return None
     return db_path if db_path is not None else client.store.db_path
 
@@ -194,7 +194,7 @@ class InfoModal(ModalScreen):
             # rather than the whole panel. Names only, no paths — a location
             # belongs in the configuration.
             blocks = await asyncio.gather(
-                *(self._report(name) for name in sorted(self.client._federated))
+                *(self._report(name) for name in sorted(self.client.source_names))
             )
             for block in blocks:
                 lines.extend(block)

haiku_rag_slim/haiku/rag/sandbox/sandbox.py

@@ -214,7 +214,7 @@ class Sandbox:
         a database the question excluded cannot be mounted.
         """
         async with self._connection() as rag:
-            if not rag._federated:
+            if not rag.covers_multiple:
                 if not await rag.clients_covering(self._context.sources):
                     return [], {}
                 docs = await rag.list_documents(filter=self._context.filter)
@@ -251,10 +251,10 @@ class Sandbox:
                 if doc.id in holders:
                     raise ValueError(
                         f"document {doc.id} is in databases {held_by[doc.id]!r} and "
-                        f"{owner._source!r}; analysis mounts one document per id"
+                        f"{owner.source!r}; analysis mounts one document per id"
                     )
                 holders[doc.id] = owner
-                held_by[doc.id] = owner._source
+                held_by[doc.id] = owner.source
         return holders
 
     def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
@@ -354,7 +354,7 @@ class Sandbox:
                     "title": d.title,
                     "uri": d.uri,
                     "created_at": str(d.created_at),
-                    "source": owners[d.id]._source if d.id in owners else None,
+                    "source": owners[d.id].source if d.id in owners else None,
                 }
                 for d in docs
             ]

haiku_rag_slim/haiku/rag/utils.py

@@ -394,7 +394,7 @@ async def format_citations_rich(
         idx = c.index if c.index is not None else (i + 1)
 
         header_parts: list[str] = [f"[{idx}] {_citation_label(c)}"]
-        if c.source and client is not None and client._federated:
+        if c.source and client is not None and client.covers_multiple:
             header_parts.append(c.source)
         pages = _citation_pages(c)
         if pages:

tests/capabilities/test_capabilities.py

@@ -144,13 +144,15 @@ def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
 def _single_database_client() -> AsyncMock:
     """A stand-in for a client covering one unnamed database.
 
-    A bare AsyncMock answers every attribute with a truthy Mock, so `_federated`
-    would read as a set of databases, `_source` would reach a validated field, and
-    `clients_covering` would return a Mock where the code iterates clients.
+    A bare AsyncMock answers every attribute with a truthy Mock, so
+    `covers_multiple` would read as a set of databases, `source` would reach a
+    validated field, and `clients_covering` would return a Mock where the code
+    iterates clients.
     """
     client = AsyncMock()
-    client._federated = {}
-    client._source = None
+    client.covers_multiple = False
+    client.source_names = ()
+    client.source = None
     client.clients_covering.return_value = [client]
     return client
 
@@ -1537,7 +1539,8 @@ class TestSeveralDatabasesInstructions:
     @staticmethod
     def _client(federated):
         client = AsyncMock()
-        client._federated = federated
+        client.covers_multiple = len(federated) > 1
+        client.source_names = tuple(federated)
         return client
 
     def test_one_database_is_instructed_as_before(self):

tests/chat/test_chat_app.py

@@ -118,10 +118,11 @@ def _make_mock_client():
     mock_client = AsyncMock()
     mock_client.__aenter__ = AsyncMock(return_value=mock_client)
     mock_client.__aexit__ = AsyncMock(return_value=None)
-    # Covers one database: a bare AsyncMock answers `_federated` with a truthy
+    # Covers one database: a bare AsyncMock answers `covers_multiple` with a truthy
     # Mock, which would send every read down the covering-a-set branch.
-    mock_client._federated = {}
-    mock_client._source = None
+    mock_client.covers_multiple = False
+    mock_client.source_names = ()
+    mock_client.source = None
     return mock_client
 
 
@@ -510,14 +511,15 @@ async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path)
     from haiku.rag.store.models.citation import Citation
 
     owner = _make_mock_client()
-    owner._source = "beta"
+    owner.source = "beta"
     owner.get_chunk_by_id.return_value = Chunk(
         id="c1", document_id="d1", content="cited body"
     )
 
     covering = _make_mock_client()
-    covering._federated = {"alpha": "/a.lancedb", "beta": "/b.lancedb"}
-    covering.clients_for = AsyncMock(return_value=[owner])
+    covering.covers_multiple = True
+    covering.source_names = ("alpha", "beta")
+    covering.reader_for = AsyncMock(return_value=owner)
 
     app, _ = _make_app(tmp_path / "unused.lancedb", covering)
     with patch("haiku.rag.chat.app.HaikuRAG", return_value=covering):
@@ -540,7 +542,7 @@ async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path)
             with patch.object(app, "push_screen", new=AsyncMock()) as push:
                 await app.action_show_visual()
 
-    covering.clients_for.assert_awaited_once_with(["beta"])
+    covering.reader_for.assert_awaited_once_with("beta")
     owner.get_chunk_by_id.assert_awaited_once_with("c1")
     assert push.await_args is not None
     assert push.await_args.args[0].client is owner
@@ -558,7 +560,8 @@ class TestDocumentSelectionIdentity:
         from haiku.rag.store.models.document import Document
 
         client = AsyncMock()
-        client._federated = {"arxiv": "a", "wiki": "b"}
+        client.covers_multiple = True
+        client.source_names = ("arxiv", "wiki")
         client.list_documents.return_value = [
             Document(id="id-one", content="", title="Capital region", source="arxiv"),
             Document(id="id-two", content="", title="Capital region", source="wiki"),

tests/sandbox/test_sandbox_multi_db.py

@@ -33,7 +33,7 @@ class TestDocumentsAcrossDatabases:
             "test://alpha/alpha document about cats",
             "test://beta/beta document about cats",
         }
-        assert {owner._source for owner in owners.values()} == {"alpha", "beta"}
+        assert {owner.source for owner in owners.values()} == {"alpha", "beta"}
 
     @pytest.mark.asyncio
     async def test_selected_databases_bound_the_corpus(self, tmp_path):
@@ -46,7 +46,7 @@ class TestDocumentsAcrossDatabases:
             _, docs, owners = await _mounted(rag, sources=["alpha"])
 
         assert [d.uri for d in docs] == ["test://alpha/alpha document about cats"]
-        assert {owner._source for owner in owners.values()} == {"alpha"}
+        assert {owner.source for owner in owners.values()} == {"alpha"}
 
     @pytest.mark.asyncio
     async def test_one_database_needs_no_owners(self, tmp_path, temp_db_path):
@@ -76,8 +76,8 @@ class TestDocumentsAcrossDatabases:
                 async with sandbox._connection(owners[doc.id]) as owner:
                     content = await owner.document_repository.get_content(doc.id)
                 assert content is not None
-                assert owners[doc.id]._source is not None
-                assert owners[doc.id]._source in content
+                assert owners[doc.id].source is not None
+                assert owners[doc.id].source in content
 
 
 class TestExecutingAcrossDatabases:

tests/test_inspector.py

@@ -339,7 +339,8 @@ class TestReportedDatabase:
     @staticmethod
     def _client(federated, store_path=None):
         client = MagicMock()
-        client._federated = federated
+        client.covers_multiple = len(federated) > 1
+        client.source_names = tuple(federated)
         client.store.db_path = store_path
         return client
 

tests/test_multi_db.py

@@ -126,8 +126,8 @@ class TestNamingADatabaseDirectly:
         await _seed(config, "alpha", ["alpha document about cats"])
 
         async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
-            assert rag._federated == {}
-            assert rag._source is None
+            assert not rag.covers_multiple
+            assert rag.source is None
             assert rag.store.db_path == temp_db_path
 
     @pytest.mark.asyncio
@@ -137,8 +137,8 @@ class TestNamingADatabaseDirectly:
         await _seed(config, "alpha", ["alpha document about cats"])
 
         async with HaikuRAG(config=config) as rag:
-            assert rag._federated == {}
-            assert rag._source == "alpha"
+            assert not rag.covers_multiple
+            assert rag.source == "alpha"
             results = await rag.search("cats", search_type="fts", limit=10)
 
         assert [r.source for r in results] == ["alpha"]
@@ -168,7 +168,7 @@ class TestOpeningDatabases:
             rag._session._open = gated
             clients = await asyncio.wait_for(rag.clients_for(names), timeout=15)
 
-        assert {client._source for client in clients} == set(names)
+        assert {client.source for client in clients} == set(names)
 
     @pytest.mark.asyncio
     async def test_a_failed_open_does_not_leak_the_ones_that_worked(self, tmp_path):
@@ -195,7 +195,7 @@ class TestOpeningDatabases:
         async with HaikuRAG(config=config) as rag:
             clients = await rag.clients_for(["alpha", "alpha", "beta"])
 
-        assert [client._source for client in clients] == ["alpha", "beta"]
+        assert [client.source for client in clients] == ["alpha", "beta"]
 
     @pytest.mark.asyncio
     async def test_a_database_named_twice_returns_each_result_once(self, tmp_path):
@@ -220,7 +220,7 @@ class TestOpeningDatabases:
         async with HaikuRAG(config=config) as rag:
             covering = await rag.clients_covering(["alpha", "alpha"])
 
-        assert [client._source for client in covering] == ["alpha"]
+        assert [client.source for client in covering] == ["alpha"]
 
 
 class TestListingAcrossDatabases:
@@ -379,6 +379,84 @@ class TestClosingASet:
         assert sorted(drained) == ["alpha", "beta"]
 
 
+class TestPlacingADatabase:
+    """What a client says about the databases it covers, so nothing outside has
+    to read its private state to find out."""
+
+    @pytest.mark.asyncio
+    async def test_a_set_names_every_database_it_covers(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha one"])
+        await _seed(config, "beta", ["beta one"])
+
+        async with HaikuRAG(config=config, read_only=True) as rag:
+            assert rag.covers_multiple
+            assert rag.source_names == ("alpha", "beta")
+            assert rag.source is None
+
+    @pytest.mark.asyncio
+    async def test_one_named_database_names_itself(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha one"])
+
+        async with HaikuRAG(config=config, read_only=True, sources=["alpha"]) as rag:
+            assert not rag.covers_multiple
+            assert rag.source_names == ("alpha",)
+            assert rag.source == "alpha"
+
+    @pytest.mark.asyncio
+    async def test_a_named_database_keeps_its_name_on_re_entry(self, tmp_path):
+        """Entering derives a single-database configuration from what was
+        configured. Deriving it from the last derivation loses the name."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        rag = HaikuRAG(config=config, read_only=True, sources=["alpha"])
+        async with rag:
+            assert rag.source == "alpha"
+        async with rag:
+            assert rag.source == "alpha"
+            assert rag.source_names == ("alpha",)
+            results = await rag.search("cats", search_type="fts")
+
+        assert {r.source for r in results} == {"alpha"}
+
+    @pytest.mark.asyncio
+    async def test_an_unnamed_database_names_nothing(self, temp_db_path):
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            assert rag.source_names == ()
+            assert rag.source is None
+
+    @pytest.mark.asyncio
+    async def test_the_reader_for_a_database_is_the_client_holding_it(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha one"])
+        await _seed(config, "beta", ["beta one"])
+
+        async with HaikuRAG(config=config, read_only=True) as rag:
+            reader = await rag.reader_for("beta")
+
+            assert reader is not None
+            assert reader.source == "beta"
+            # Asked twice, the same wrapper comes back.
+            assert await rag.reader_for("beta") is reader
+
+    @pytest.mark.asyncio
+    async def test_a_client_reading_one_database_is_its_own_reader(self, temp_db_path):
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            assert await rag.reader_for(None) is rag
+            assert await rag.reader_for("anything") is rag
+
+    @pytest.mark.asyncio
+    async def test_a_set_cannot_place_evidence_that_names_no_database(self, tmp_path):
+        """Evidence recorded before databases could be named carries no source."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha one"])
+
+        async with HaikuRAG(config=config, read_only=True) as rag:
+            assert await rag.reader_for(None) is None
+
+
 class TestBorrowedDatabases:
     """A client for one of a set wraps a database the set opened."""
 

tests/test_multi_db_ask.py

@@ -293,7 +293,7 @@ class TestStandaloneCapabilities:
             await run._close()
 
         assert len(docs) == 2
-        assert {owner._source for owner in owners.values()} == {"alpha", "beta"}
+        assert {owner.source for owner in owners.values()} == {"alpha", "beta"}
 
     @pytest.mark.asyncio
     async def test_a_single_configured_database_is_still_opened(self, tmp_path):
@@ -304,7 +304,7 @@ class TestStandaloneCapabilities:
         capability = create_capability(config=config, defer_loading=False)
         rag = await capability._ensure_rag()
         try:
-            assert rag._source == "alpha"
+            assert rag.source == "alpha"
         finally:
             await capability._close()
 

tests/test_utils.py

@@ -785,7 +785,8 @@ async def test_format_citations_rich_names_the_database_when_federating():
         source="medic",
     )
     client = AsyncMock()
-    client._federated = {"medic": "/data/medic.lancedb", "st": "/data/st.lancedb"}
+    client.covers_multiple = True
+    client.source_names = ("medic", "st")
 
     output = _render_rich(await format_citations_rich([citation], client))
 
@@ -808,7 +809,8 @@ async def test_format_citations_rich_omits_the_database_for_one_database():
         source="medic",
     )
     client = AsyncMock()
-    client._federated = {}
+    client.covers_multiple = False
+    client.source_names = ()
 
     output = _render_rich(await format_citations_rich([citation], client))