From b48ac2fdd01b2e9b0ce00b63bd37afefb01f0b1d Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 20 May 2026 13:02:25 +0300
Subject: [PATCH] sandbox + cite: guard sort precondition; require rag in cite
 fallback

---
 haiku_rag_slim/haiku/rag/sandbox/sandbox.py |  4 ++++
 haiku_rag_slim/haiku/rag/skills/_tools.py   |  4 ++--
 tests/skills/test_rag.py                    | 17 +++++++++++------
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
index 0bb9dc62..d716fb5f 100644
--- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
+++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py
@@ -62,6 +62,10 @@ def _build_toc(
     flat sibling list (see docling-project/docling#2121 for an upstream case
     where every PDF section_header is emitted at level=1).
     """
+    # Defensive: every consumer is supposed to pass items in position order,
+    # but the end_exclusive lookahead below silently miscomputes section
+    # boundaries if it's not — better to sort once than trust the caller.
+    items = sorted(items, key=lambda i: i.position)
     headers: list[DocumentItem] = [
         i for i in items if i.label == "section_header" and i.heading_level > 0
     ]
diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py
index f5907083..20bc69d4 100644
--- a/haiku_rag_slim/haiku/rag/skills/_tools.py
+++ b/haiku_rag_slim/haiku/rag/skills/_tools.py
@@ -324,8 +324,8 @@ def create_skill_tools(
                 if cid.strip("[]") not in resolved_ids
             ]
 
-            rag = ctx.deps.rag if ctx.deps else None
-            if missing and rag is not None:
+            if missing:
+                rag = _require_rag(ctx)
                 synthetic: list[SearchResult] = []
                 doc_cache: dict[str, Any] = {}
                 for cid in missing:
diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py
index 02e5929c..26dbb413 100644
--- a/tests/skills/test_rag.py
+++ b/tests/skills/test_rag.py
@@ -346,9 +346,11 @@ class TestCiteTool:
         assert "verbatim" in message
         assert "372c9ddf-not-a-real-id" in message
 
-    async def test_cite_raises_modelretry_when_no_searches_recorded(self, rag_db):
-        """If cite is called before any search has populated state.searches,
-        the retry message tells the model to call search first."""
+    async def test_cite_raises_modelretry_when_chunk_id_does_not_exist(
+        self, rag_db, rag_client
+    ):
+        """With no prior search() and a chunk_id that doesn't exist in the DB,
+        cite raises ModelRetry — the DB-fallback path tried and found nothing."""
         from pydantic_ai import ModelRetry
 
         from haiku.rag.skills.rag import RAGState, create_skill
@@ -356,11 +358,14 @@ class TestCiteTool:
         skill = create_skill(db_path=rag_db)
         cite = _get_tool(skill, "cite")
         state = RAGState()
-        ctx = _make_ctx(state)
+        ctx = _make_ctx(state, rag=rag_client)
+        assert not state.searches
 
         with pytest.raises(ModelRetry) as exc_info:
-            await cite(ctx, chunk_ids=["any-id"])
-        assert "search" in str(exc_info.value).lower()
+            await cite(ctx, chunk_ids=["nonexistent-chunk-id"])
+        message = str(exc_info.value)
+        assert "verbatim" in message
+        assert "nonexistent-chunk-id" in message
 
     async def test_cite_returns_message_when_chunk_ids_empty(self, rag_db):
         """An empty chunk_ids list is a no-op, not a retry trigger."""