From 49580228c2b9b34731159441b4df32c6f209e286 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 26 Aug 2026 10:33:29 +0300
Subject: [PATCH] Read the document's own database in the sandbox listing

`list_documents()` reached into the ownership map for a name the document
already carries. The map is still built, so a document id two databases
claim is still refused before anything mounts. A listing over one named
database now reports that name, as in-code `search()` already does.
---
 haiku_rag_slim/haiku/rag/sandbox/sandbox.py |  4 ++--
 tests/sandbox/test_sandbox.py               |  5 ++++-
 tests/sandbox/test_sandbox_multi_db.py      | 19 +++++++++++++++++++
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
index 10bb7f65..c424e891 100644
--- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
+++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
@@ -347,14 +347,14 @@ class Sandbox:
             return out
 
         async def list_documents() -> list[dict[str, Any]]:
-            docs, owners = await self._documents()
+            docs, _ = await self._documents()
             return [
                 {
                     "id": d.id,
                     "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": d.source,
                 }
                 for d in docs
             ]
diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py
index 90423a5a..a5883345 100644
--- a/tests/sandbox/test_sandbox.py
+++ b/tests/sandbox/test_sandbox.py
@@ -104,11 +104,14 @@ class TestSandboxListDocuments:
             result = await sb.execute(
                 "docs = await list_documents()\n"
                 "print(len(docs))\n"
-                "print(docs[0]['title'])"
+                "print(docs[0]['title'])\n"
+                "print(docs[0]['source'])"
             )
             assert result.success
             assert "1" in result.stdout
             assert "Test Document" in result.stdout
+            # Nothing names this database, so there is no name to report.
+            assert "None" in result.stdout
 
 
 class TestSandboxSearch:
diff --git a/tests/sandbox/test_sandbox_multi_db.py b/tests/sandbox/test_sandbox_multi_db.py
index ba8dfa18..d842ceaa 100644
--- a/tests/sandbox/test_sandbox_multi_db.py
+++ b/tests/sandbox/test_sandbox_multi_db.py
@@ -244,3 +244,22 @@ class TestListingOrder:
 
         assert "source" in rows[0]
         assert {r["source"] for r in rows} == {"alpha", "beta"}
+
+    @pytest.mark.asyncio
+    async def test_in_code_list_documents_names_one_database_too(self, tmp_path):
+        """A document knows which database it came from whether or not the
+        analysis spans several."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha one"])
+        await _seed(config, "beta", ["beta one"])
+
+        async with HaikuRAG(config=config, sources=["alpha"]) as rag:
+            sandbox = Sandbox(
+                db_path=None,
+                config=config,
+                context=AnalysisContext(),
+                rag=rag,
+            )
+            rows = await sandbox._build_external_functions()["list_documents"]()
+
+        assert [r["source"] for r in rows] == ["alpha"]