Merge pull request #489 from ggozad/fix/excessive-citations

Scope citations to the matched content and make visual grounding faithful

CHANGELOG.md

@@ -1,6 +1,13 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `SearchResult.chunk_ids` and `Citation.chunk_ids` carry the chunk ids merged into an expanded result.
+- `Citation.doc_item_refs` carries the items the model saw; `visualize_chunk` accepts a `refs` argument and resolves bounding boxes from them, so visualizations match the cited content instead of re-expanding.
+- Chunk visualizations draw matched content in a stronger highlight than surrounding context.
+- `haiku-rag visualize --no-expand` highlights only the chunk itself, without its expanded context (`visualize_chunk` gains an `expand` argument).
+
 ### Changed
 
 - Duplicate images within a document produce a single picture chunk.
@@ -9,6 +16,9 @@
 ### Fixed
 
 - vLLM embedding and vLLM/Jina reranking reuse one HTTP client across requests instead of opening one per request.
+- Picture and table search hits expand within their section, never across section boundaries.
+- A merged search result anchors its `chunk_id` on the highest-scoring constituent chunk instead of the one earliest in the document.
+- A clipped citation's `page_numbers` and `doc_item_refs` reflect only the content that survived the budget, not the full expanded range.
 
 ## [0.64.0] - 2026-07-08
 

app/backend/main.py

@@ -203,22 +203,43 @@ async def db_info(_: Request) -> JSONResponse:
 
 
 async def visualize_chunk(request: Request) -> JSONResponse:
-    """Return visual grounding images for a chunk as base64."""
+    """Return visual grounding images for one or more chunks as base64.
+
+    The path param accepts comma-separated chunk ids (a merged citation's
+    constituent chunks). The optional ``refs`` query param is a JSON-encoded
+    list of the citation's ``doc_item_refs`` — the exact items the model saw —
+    so the highlight matches the cited content instead of re-expanding.
+    """
     import base64
+    import json
     from io import BytesIO
 
     chunk_id = request.path_params["chunk_id"]
 
+    refs: list[str] | None = None
+    refs_param = request.query_params.get("refs")
+    if refs_param:
+        try:
+            parsed = json.loads(refs_param)
+        except ValueError:
+            parsed = None
+        if isinstance(parsed, list):
+            refs = [str(x) for x in parsed]
+
     if not db_path.exists():
         return JSONResponse({"error": "Database not found"}, status_code=404)
 
     client = await get_client()
 
-    chunk = await client.chunk_repository.get_by_id(chunk_id)
-    if not chunk:
+    chunks = []
+    for cid in chunk_id.split(","):
+        chunk = await client.chunk_repository.get_by_id(cid)
+        if chunk:
+            chunks.append(chunk)
+    if not chunks:
         return JSONResponse({"error": "Chunk not found"}, status_code=404)
 
-    images = await client.visualize_chunk(chunk)
+    images = await client.visualize_chunk(chunks, refs)
     if not images:
         return JSONResponse({"images": [], "message": "No visual grounding available"})
 
@@ -233,7 +254,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
         {
             "images": base64_images,
             "chunk_id": chunk_id,
-            "document_uri": chunk.document_uri,
+            "document_uri": chunks[0].document_uri,
         }
     )
 

app/frontend/app/api/visualize/[chunk_id]/route.ts

@@ -1,14 +1,18 @@
 import { NextResponse } from "next/server";
 
 export async function GET(
-	_request: Request,
+	request: Request,
 	{ params }: { params: Promise<{ chunk_id: string }> },
 ) {
 	const { chunk_id } = await params;
 	const backendUrl = process.env.BACKEND_URL || "http://backend:8000";
+	const refs = new URL(request.url).searchParams.get("refs");
+	const query = refs ? `?refs=${encodeURIComponent(refs)}` : "";
 
 	try {
-		const response = await fetch(`${backendUrl}/api/visualize/${chunk_id}`);
+		const response = await fetch(
+			`${backendUrl}/api/visualize/${chunk_id}${query}`,
+		);
 		const data = await response.json();
 
 		if (!response.ok) {

app/frontend/components/CitationBlock.tsx

@@ -20,7 +20,7 @@ function CitationItem({
 	onViewInDocument,
 }: {
 	citation: Citation;
-	onViewInDocument: (chunkId: string) => void;
+	onViewInDocument: (chunkId: string, refs?: string[]) => void;
 }) {
 	const [expanded, setExpanded] = useState(false);
 
@@ -55,7 +55,14 @@ function CitationItem({
 					<button
 						type="button"
 						className="citation-view-btn"
-						onClick={() => onViewInDocument(citation.chunk_id)}
+						onClick={() =>
+							onViewInDocument(
+								citation.chunk_ids?.length
+									? citation.chunk_ids.join(",")
+									: citation.chunk_id,
+								citation.doc_item_refs,
+							)
+						}
 					>
 						View in Document
 					</button>
@@ -82,45 +89,54 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
 		return () => abortRef.current?.abort();
 	}, []);
 
-	const fetchVisualGrounding = useCallback(async (chunkId: string) => {
-		abortRef.current?.abort();
-		const controller = new AbortController();
-		abortRef.current = controller;
+	const fetchVisualGrounding = useCallback(
+		async (chunkId: string, refs?: string[]) => {
+			abortRef.current?.abort();
+			const controller = new AbortController();
+			abortRef.current = controller;
 
-		setVisualGrounding({
-			isOpen: true,
-			chunkId,
-			images: [],
-			loading: true,
-			error: null,
-		});
-
-		try {
-			const response = await fetch(`/api/visualize/${chunkId}`, {
-				signal: controller.signal,
+			setVisualGrounding({
+				isOpen: true,
+				chunkId,
+				images: [],
+				loading: true,
+				error: null,
 			});
-			const data = await response.json();
 
-			if (controller.signal.aborted) return;
-			if (!response.ok) {
-				throw new Error(data.error || "Failed to fetch visual grounding");
+			const query = refs?.length
+				? `?refs=${encodeURIComponent(JSON.stringify(refs))}`
+				: "";
+			try {
+				const response = await fetch(
+					`/api/visualize/${encodeURIComponent(chunkId)}${query}`,
+					{
+						signal: controller.signal,
+					},
+				);
+				const data = await response.json();
+
+				if (controller.signal.aborted) return;
+				if (!response.ok) {
+					throw new Error(data.error || "Failed to fetch visual grounding");
+				}
+
+				setVisualGrounding((prev) => ({
+					...prev,
+					images: data.images || [],
+					loading: false,
+					error: data.images?.length === 0 ? data.message : null,
+				}));
+			} catch (err) {
+				if (controller.signal.aborted) return;
+				setVisualGrounding((prev) => ({
+					...prev,
+					loading: false,
+					error: err instanceof Error ? err.message : "Unknown error",
+				}));
 			}
-
-			setVisualGrounding((prev) => ({
-				...prev,
-				images: data.images || [],
-				loading: false,
-				error: data.images?.length === 0 ? data.message : null,
-			}));
-		} catch (err) {
-			if (controller.signal.aborted) return;
-			setVisualGrounding((prev) => ({
-				...prev,
-				loading: false,
-				error: err instanceof Error ? err.message : "Unknown error",
-			}));
-		}
-	}, []);
+		},
+		[],
+	);
 
 	const closeVisualGrounding = useCallback(() => {
 		abortRef.current?.abort();

app/frontend/lib/sessionStorage.ts

@@ -2,11 +2,13 @@ export interface Citation {
 	index: number;
 	document_id: string;
 	chunk_id: string;
+	chunk_ids?: string[];
 	document_uri: string;
 	document_title: string | null;
 	page_numbers: number[];
 	headings: string[] | null;
 	content: string;
+	doc_item_refs?: string[];
 }
 
 // Matches RAGState from the backend skill

docs/cli.md

@@ -246,7 +246,9 @@ Display visual grounding for a chunk - shows page images with highlighted boundi
 haiku-rag visualize <chunk_id>
 ```
 
-This renders the source document pages with the chunk's location highlighted. Useful for verifying chunk boundaries and understanding document structure.
+This renders the source document pages with the chunk's location highlighted. The chunk itself draws in a strong highlight, while surrounding context swept in by expansion draws fainter. Useful for verifying chunk boundaries and understanding document structure.
+
+Pass `--no-expand` to highlight only the chunk itself, without its expanded context.
 
 !!! note
     Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored.

docs/configuration/qa.md

@@ -13,7 +13,7 @@ search:
 - **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 10
 - **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
 
-Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
+Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
 
 !!! note "Reranking behavior"
     When a reranker is configured, search automatically retrieves 10x the requested limit, then reranks to return the final count. This improves result quality without requiring you to adjust `limit`.

docs/python.md

@@ -298,7 +298,7 @@ for result in expanded_results:
     print(f"Expanded content: {result.content}")
 ```
 
-Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
+Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
 
 Configuration:
 

haiku_rag_slim/haiku/rag/app.py

@@ -457,7 +457,7 @@ class HaikuRAGApp:  # pragma: no cover
             for result in results:
                 self._rich_print_search_result(result)
 
-    async def visualize_chunk(self, chunk_id: str):
+    async def visualize_chunk(self, chunk_id: str, expand: bool = True):
         """Display visual grounding images for a chunk."""
         from textual_image.renderable import Image as RichImage
 
@@ -472,7 +472,7 @@ class HaikuRAGApp:  # pragma: no cover
                 self.console.print(f"[red]Chunk with id {chunk_id} not found.[/red]")
                 return
 
-            images = await self.client.visualize_chunk(chunk)
+            images = await self.client.visualize_chunk(chunk, expand=expand)
             if not images:
                 self.console.print(
                     "[yellow]No visual grounding available for this chunk.[/yellow]"

haiku_rag_slim/haiku/rag/chat/app.py

@@ -405,13 +405,24 @@ class ChatApp(App):
             return
 
         citation = selected_widgets[0].citation
-        chunk = await self.client.get_chunk_by_id(citation.chunk_id)
-        if not chunk:
+        chunk_ids = citation.chunk_ids or [citation.chunk_id]
+        chunks = []
+        for cid in chunk_ids:
+            chunk = await self.client.get_chunk_by_id(cid)
+            if chunk:
+                chunks.append(chunk)
+        if not chunks:
             return
 
         from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
 
-        await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client))
+        await self.push_screen(
+            VisualGroundingModal(
+                chunk=chunks,
+                client=self.client,
+                refs=citation.doc_item_refs or None,
+            )
+        )
 
     async def action_show_info(self) -> None:
         """Show database info modal."""

haiku_rag_slim/haiku/rag/cli.py

@@ -346,9 +346,14 @@ def visualize(  # pragma: no cover
         "--db",
         help="Path to the LanceDB database file",
     ),
+    no_expand: bool = typer.Option(
+        False,
+        "--no-expand",
+        help="Highlight only the chunk itself, without its expanded context",
+    ),
 ):
     app = create_app(db)
-    asyncio.run(app.visualize_chunk(chunk_id=chunk_id))
+    asyncio.run(app.visualize_chunk(chunk_id=chunk_id, expand=not no_expand))
 
 
 @_cli.command("ask", help="Ask a question using the QA agent")

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -4,7 +4,7 @@ import json
 import logging
 import mimetypes
 import tempfile
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Sequence
 from datetime import datetime
 from enum import Enum
 from functools import cached_property
@@ -504,10 +504,15 @@ class HaikuRAG:
 
         return await analyze(self, question, filter)
 
-    async def visualize_chunk(self, chunk: Chunk) -> list:
+    async def visualize_chunk(
+        self,
+        chunk: Chunk | Sequence[Chunk],
+        refs: list[str] | None = None,
+        expand: bool = True,
+    ) -> list:
         from haiku.rag.client.search import visualize_chunk
 
-        return await visualize_chunk(self, chunk)
+        return await visualize_chunk(self, chunk, refs, expand)
 
     async def rebuild_database(
         self, mode: RebuildMode = RebuildMode.FULL

haiku_rag_slim/haiku/rag/client/search.py

@@ -1,4 +1,5 @@
 import base64
+from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
@@ -223,12 +224,24 @@ async def expand_context(
     return expanded_results
 
 
-async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
-    """Render page images with bounding box highlights for a chunk.
+async def visualize_chunk(
+    client: "HaikuRAG",
+    chunk: "Chunk | Sequence[Chunk]",
+    refs: list[str] | None = None,
+    expand: bool = True,
+) -> list:
+    """Render page images with bounding box highlights for one or more chunks.
 
-    Expands the chunk's context to find the full section, then resolves
-    bounding boxes from all items in the expanded range. This ensures
-    visualization covers all pages the expanded content spans.
+    When ``refs`` is given (the ``doc_item_refs`` of the citation, i.e. the
+    exact items the model saw), bounding boxes are resolved from them directly
+    so the visualization matches the cited context precisely. Otherwise, with
+    ``expand=True`` (default) the chunks' context is re-expanded to recover the
+    surrounding section; with ``expand=False`` only the chunks' own items are
+    drawn, so the visualization shows just the retrieved chunk with no context.
+
+    The chunks' own items draw in a strong highlight; the remaining items draw
+    fainter, so the matched content stands out from its surrounding context.
+    Chunks from a different document than the first are ignored.
 
     Returns a list of PIL Image objects, one per page with bounding boxes.
     Empty list if no bounding boxes or page images available.
@@ -239,10 +252,15 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
 
     from haiku.rag.store.models.chunk import ChunkMetadata
 
-    if not chunk.document_id:
+    chunks = [chunk] if isinstance(chunk, Chunk) else list(chunk)
+    if not chunks:
         return []
+    document_id = chunks[0].document_id
+    if not document_id:
+        return []
+    chunks = [c for c in chunks if c.document_id == document_id]
 
-    doc = await client.document_repository.get_docling_data(chunk.document_id)
+    doc = await client.document_repository.get_docling_data(document_id)
     if not doc:
         return []
 
@@ -250,35 +268,60 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
     if not docling_doc:
         return []
 
-    # Expand context to get all doc_item_refs in the section
-    chunk_meta = chunk.get_chunk_metadata()
-    if chunk_meta.doc_item_refs:
-        search_result = SearchResult(
-            content=chunk.content,
-            score=1.0,
-            chunk_id=chunk.id,
-            document_id=chunk.document_id,
-            doc_item_refs=chunk_meta.doc_item_refs,
-            page_numbers=chunk_meta.page_numbers,
-        )
-        expanded = await expand_context(client, [search_result])
-        refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
-        meta = ChunkMetadata(doc_item_refs=refs)
+    matched_refs = {r for c in chunks for r in c.get_chunk_metadata().doc_item_refs}
+
+    if refs is not None:
+        all_refs = list(refs)
+    elif not expand:
+        # Chunk-only: draw just the retrieved chunks' own items, no context.
+        all_refs = list(matched_refs)
     else:
-        meta = chunk_meta
-    bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
-    if not bounding_boxes:
+        # No stored context: re-expand the chunks to recover their section.
+        search_results = [
+            SearchResult(
+                content=c.content,
+                score=1.0,
+                chunk_id=c.id,
+                document_id=c.document_id,
+                doc_item_refs=meta.doc_item_refs,
+                page_numbers=meta.page_numbers,
+            )
+            for c in chunks
+            if (meta := c.get_chunk_metadata()).doc_item_refs
+        ]
+        if search_results:
+            expanded = await expand_context(client, search_results)
+            all_refs = []
+            for result in expanded:
+                all_refs.extend(r for r in result.doc_item_refs if r not in all_refs)
+            if not all_refs:
+                all_refs = [r for sr in search_results for r in sr.doc_item_refs]
+        else:
+            all_refs = list(chunks[0].get_chunk_metadata().doc_item_refs)
+
+    matched_draw = [r for r in all_refs if r in matched_refs]
+    swept_refs = [r for r in all_refs if r not in matched_refs]
+
+    matched_boxes = ChunkMetadata(doc_item_refs=matched_draw).resolve_bounding_boxes(
+        docling_doc
+    )
+    swept_boxes = ChunkMetadata(doc_item_refs=swept_refs).resolve_bounding_boxes(
+        docling_doc
+    )
+    if not matched_boxes and not swept_boxes:
         return []
 
-    # Group bounding boxes by page
+    # Group bounding boxes by page; swept boxes first so matched draw on top
     boxes_by_page: dict[int, list] = {}
-    for bbox in bounding_boxes:
+    for bbox, is_matched in [(b, False) for b in swept_boxes] + [
+        (b, True) for b in matched_boxes
+    ]:
         if bbox.page_no not in boxes_by_page:
             boxes_by_page[bbox.page_no] = []
-        boxes_by_page[bbox.page_no].append(bbox)
+        boxes_by_page[bbox.page_no].append((bbox, is_matched))
 
     # Load only the needed page images
-    pages_doc = await client.document_repository.get_pages_data(chunk.document_id)
+    pages_doc = await client.document_repository.get_pages_data(document_id)
     if not pages_doc:
         return []
     page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
@@ -303,7 +346,7 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
         image = deepcopy(pil_image)
         draw = ImageDraw.Draw(image, "RGBA")
 
-        for bbox in boxes_by_page[page_no]:
+        for bbox, is_matched in boxes_by_page[page_no]:
             # Document coords are bottom-left origin; PIL uses top-left
             x0 = bbox.left * scale_x
             y0 = (page_height - bbox.top) * scale_y
@@ -313,8 +356,12 @@ async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
             if y0 > y1:
                 y0, y1 = y1, y0
 
-            fill_color = (255, 255, 0, 40)  # Yellow with transparency
-            outline_color = (255, 165, 0, 100)  # Orange outline
+            if is_matched:
+                fill_color = (255, 150, 0, 55)  # Orange, matched content
+                outline_color = (240, 130, 0, 150)  # Orange outline
+            else:
+                fill_color = (255, 255, 0, 40)  # Yellow, surrounding context
+                outline_color = (255, 165, 0, 100)
 
             draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
             draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1)

haiku_rag_slim/haiku/rag/context.py

@@ -12,7 +12,8 @@ For STRUCTURED documents (containing section_header or title labels):
   5. If the section is too small (under 20% of max_context_chars), expand
      item-by-item crossing into adjacent sections until the budget is filled.
      This lets small sections (e.g., title+authors) grow into neighboring
-     content.
+     content. Picture and table matches are exempt: they return their
+     enclosing section as-is, never crossing section boundaries.
   6. Merge overlapping ranges from multiple results in the same document.
      Adjacent but non-overlapping ranges stay separate to preserve section
      independence.
@@ -36,6 +37,10 @@ from haiku.rag.store.repositories.document_item import DocumentItemRepository
 _NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
 _SECTION_BOUNDARY_LABELS = {"section_header", "title"}
 
+# Labels whose pertinent unit is the item plus its own section: expansion
+# never crosses section boundaries for these matches.
+_TIGHT_LABELS = {"picture", "table"}
+
 # Sections with fewer chars than this fraction of max_context_chars are
 # considered too small — expansion falls through to item-by-item outward
 # growth, which naturally crosses into adjacent sections.
@@ -99,16 +104,18 @@ def _evidence_anchors(content: str, max_chars: int) -> list[str]:
     return anchors
 
 
-def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
-    """Clip expanded content to ``max_chars``, keeping the matched evidence.
+def _clip_window(
+    content: str, results: list[SearchResult], max_chars: int
+) -> tuple[int, int]:
+    """Return the ``[start, end)`` char window to keep when clipping ``content``.
 
     Anchors on the first locatable result in ``results`` order (the primary
     chunk that supplies the expanded result's identity) and returns a
-    ``max_chars``-wide window centered on it. Falls back to a prefix cut only
+    ``max_chars``-wide window centered on it. Falls back to a prefix window only
     when no anchor is locatable (heavy drift).
     """
     if max_chars <= 0:
-        return ""
+        return (0, 0)
     evidence_start, evidence_len = -1, 0
     for result in results:
         for anchor in _evidence_anchors(result.content, max_chars):
@@ -119,14 +126,45 @@ def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -
         if evidence_start != -1:
             break
     if evidence_start == -1:
-        return content[:max_chars]
+        return (0, min(len(content), max_chars))
     center = evidence_start + evidence_len // 2
     start = max(0, center - max_chars // 2)
     end = min(len(content), start + max_chars)
     start = max(0, end - max_chars)
+    return (start, end)
+
+
+def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
+    """Clip expanded content to ``max_chars``, keeping the matched evidence."""
+    start, end = _clip_window(content, results, max_chars)
     return content[start:end]
 
 
+def _collect_meta(
+    spans: list[tuple[int, int, DocumentItem]],
+) -> tuple[set[int], list[str], set[str]]:
+    """Union the page numbers, refs, and labels of the given item spans."""
+    pages: set[int] = set()
+    refs: list[str] = []
+    labels: set[str] = set()
+    for _start, _end, item in spans:
+        refs.append(item.self_ref)
+        if item.label:
+            labels.add(item.label)
+        pages.update(item.page_numbers)
+    return pages, refs, labels
+
+
+def _span_in_window(
+    span: tuple[int, int, DocumentItem], win_start: int, win_end: int
+) -> bool:
+    """Whether an item's char span overlaps the ``[win_start, win_end]`` clip window."""
+    start, end, _item = span
+    if start == end:  # zero-width picture position
+        return win_start <= start <= win_end
+    return start < win_end and end > win_start
+
+
 def _expand_outward(
     items: list[DocumentItem],
     center_idx: int,
@@ -222,6 +260,11 @@ def _find_expansion_range(
             hi_bound=sec_end,
         )
 
+    # Picture/table hits stay section-bounded: their pertinent unit is the
+    # figure or table plus its section, never neighboring sections.
+    if any(items[i].label in _TIGHT_LABELS for i in matched_indices):
+        return (items[sec_start].position, items[sec_end].position)
+
     # Section too small (e.g., title+authors) — expand across boundaries
     return _expand_outward(items, center_idx, max_chars, skip_noise=True)
 
@@ -280,9 +323,11 @@ async def expand_with_items(
     final_results: list[SearchResult] = []
     for range_start, range_end, original_results in merged:
         content_parts: list[str] = []
-        refs: list[str] = []
-        pages: set[int] = set()
-        labels: set[str] = set()
+        # Char span of each contributing item within the joined content, so
+        # metadata can be narrowed to whatever survives a budget clip.
+        item_spans: list[tuple[int, int, DocumentItem]] = []
+        cursor = 0
+        separator = "\n\n"
 
         for pos in range(range_start, range_end + 1):
             item = pos_to_item.get(pos)
@@ -291,56 +336,96 @@ async def expand_with_items(
             if has_sections and item.label in _NOISE_LABELS:
                 continue
             if item.text:
+                if content_parts:
+                    cursor += len(separator)
+                start = cursor
                 content_parts.append(item.text)
-                refs.append(item.self_ref)
-                if item.label:
-                    labels.add(item.label)
-                pages.update(item.page_numbers)
+                cursor += len(item.text)
+                item_spans.append((start, cursor, item))
             elif item.label == "picture":
                 # Pictures may legitimately have empty text (no VLM
                 # description configured). Keep their self_ref so the
-                # downstream image_data lookup can still attach bytes.
-                refs.append(item.self_ref)
-                labels.add(item.label)
-                pages.update(item.page_numbers)
+                # downstream image_data lookup can still attach bytes. They
+                # occupy a zero-width position in reading order.
+                item_spans.append((cursor, cursor, item))
 
         all_headings: list[str] = []
         for r in original_results:
             if r.headings:
                 all_headings.extend(h for h in r.headings if h not in all_headings)
 
-        # Carry image_data and picture_captions through expansion so that
-        # only pictures from the originally retrieved chunks get attached.
-        # Pictures swept in by section expansion are referenced in `refs`
-        # for cross-referencing but their bytes are not re-fetched —
-        # otherwise a single search can balloon the response with adjacent
-        # figures the model did not actually retrieve.
+        # Anchor identity (chunk_id, content/refs fallbacks) on the
+        # best-scoring constituent — the chunk that earned the result its
+        # rank — rather than whichever sits earliest in the document.
+        first = max(original_results, key=lambda r: r.score)
+
+        chunk_ids: list[str] = []
+        for r in original_results:
+            if r.chunk_id and r.chunk_id not in chunk_ids:
+                chunk_ids.append(r.chunk_id)
+
+        joined = separator.join(content_parts)
+
+        # Expansion should never return less content than the original chunk.
+        # This can happen when item texts are fragmented (e.g., docling splits
+        # formatted HTML list items into many small text nodes); fall back to
+        # the chunk's own content, described by the chunk's own metadata.
+        if len(joined) < len(first.content):
+            base_content, base_spans = first.content, None
+        else:
+            base_content, base_spans = joined, item_spans
+
+        if len(base_content) > max_chars:
+            # Clip to the budget, anchored on the primary (highest-scoring)
+            # chunk, and narrow metadata to whatever survives the window.
+            win_start, win_end = _clip_window(
+                base_content, [first, *original_results], max_chars
+            )
+            expanded_content = base_content[win_start:win_end]
+            if base_spans is None:
+                pages, refs, labels = (
+                    set(first.page_numbers),
+                    list(first.doc_item_refs),
+                    set(first.labels),
+                )
+            else:
+                pages, refs, labels = _collect_meta(
+                    [s for s in base_spans if _span_in_window(s, win_start, win_end)]
+                )
+        else:
+            expanded_content = base_content
+            if base_spans is None:
+                pages, refs, labels = (
+                    set(first.page_numbers),
+                    list(first.doc_item_refs),
+                    set(first.labels),
+                )
+            else:
+                pages, refs, labels = _collect_meta(base_spans)
+
+        # Carry image_data and picture_captions from the originally retrieved
+        # chunks, but only for constituents whose refs survive the window — a
+        # chunk clipped out of the budget must not still ship its image to the
+        # model. Pictures swept in by section expansion are referenced in
+        # ``refs`` but their bytes are never re-fetched, so the multimodal
+        # payload stays bounded to what was actually retrieved and shown.
+        surviving_refs = set(refs)
         merged_image_data: dict[str, str] = {}
         merged_captions: dict[str, str] = {}
         for r in original_results:
+            if r.doc_item_refs and not surviving_refs.intersection(r.doc_item_refs):
+                continue
             if r.image_data:
                 merged_image_data.update(r.image_data)
             if r.picture_captions:
                 merged_captions.update(r.picture_captions)
 
-        first = original_results[0]
-
-        # Expansion should never return less content than the original chunk.
-        # This can happen when item texts are fragmented (e.g., docling splits
-        # formatted HTML list items into many small text nodes).
-        expanded_content = "\n\n".join(content_parts)
-        if len(expanded_content) < len(first.content):
-            expanded_content = first.content
-        if len(expanded_content) > max_chars:
-            expanded_content = _clip_to_budget(
-                expanded_content, original_results, max_chars
-            )
-
         final_results.append(
             SearchResult(
                 content=expanded_content,
                 score=max(r.score for r in original_results),
                 chunk_id=first.chunk_id,
+                chunk_ids=chunk_ids,
                 document_id=first.document_id,
                 document_uri=first.document_uri,
                 document_title=first.document_title,

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

@@ -57,14 +57,16 @@ class VisualGroundingModal(Screen):
 
     def __init__(
         self,
-        chunk: "Chunk",
+        chunk: "Chunk | list[Chunk]",
         client: "HaikuRAG",
         document_uri: str | None = None,
+        refs: list[str] | None = None,
     ):
         super().__init__()
-        self.chunk = chunk
+        self.chunks = chunk if isinstance(chunk, list) else [chunk]
         self.client = client
-        self.document_uri = document_uri or chunk.document_uri
+        self.refs = refs
+        self.document_uri = document_uri or self.chunks[0].document_uri
         self.images: list[PILImage] = []
         self.current_page_idx = 0
         self._image_widget: Widget = Static("Loading...", id="image-display")
@@ -81,7 +83,7 @@ class VisualGroundingModal(Screen):
 
     async def on_mount(self) -> None:
         """Load images and display the first page."""
-        self.images = await self.client.visualize_chunk(self.chunk)
+        self.images = await self.client.visualize_chunk(self.chunks, self.refs)
         await self._render_current_page()
 
     async def _render_current_page(self) -> None:

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -131,14 +131,19 @@ def create_skill_extras(
     - 'visualize_chunk': returns visualizations for chunks in the database
     """
 
-    async def visualize_chunk(chunk_id: str) -> list:
+    async def visualize_chunk(chunk_id: str | list[str]) -> list:
         from haiku.rag.client import HaikuRAG
 
+        chunk_ids = [chunk_id] if isinstance(chunk_id, str) else chunk_id
         async with HaikuRAG(db_path, config=config, read_only=True) as rag:
-            chunk = await rag.get_chunk_by_id(chunk_id)
-            if chunk is None:
+            chunks = []
+            for cid in chunk_ids:
+                chunk = await rag.get_chunk_by_id(cid)
+                if chunk is not None:
+                    chunks.append(chunk)
+            if not chunks:
                 return []
-            return await rag.visualize_chunk(chunk)
+            return await rag.visualize_chunk(chunks)
 
     async def list_documents(
         limit: int | None = None,

haiku_rag_slim/haiku/rag/store/models/chunk.py

@@ -124,11 +124,17 @@ class SearchResult(BaseModel):
     when the caller asked to omit them via ``include_images=False`` on
     ``client.search``. Same shape is used everywhere — MCP, in-process search,
     agent toolsets — so non-vision callers see ``None`` and pay nothing.
+
+    ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
+    into this result; empty means just ``chunk_id``. It lets citation
+    consumers (visual grounding) reproduce a merged expansion and is never
+    part of ``format_for_agent`` output.
     """
 
     content: str
     score: float
     chunk_id: str | None = None
+    chunk_ids: list[str] = []
     document_id: str | None = None
     document_uri: str | None = None
     document_title: str | None = None

haiku_rag_slim/haiku/rag/store/models/citation.py

@@ -18,16 +18,26 @@ class Citation(BaseModel):
     cited chunk. Empty for text-only citations. UIs can fetch the picture
     bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
     and render them alongside the text content.
+
+    ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
+    into the cited result (always includes ``chunk_id``).
+
+    ``doc_item_refs`` are the ``self_ref`` values of every item in the cited
+    content — the exact items the model saw. Visual grounding resolves bounding
+    boxes from them so the rendered pages match the citation precisely.
+    ``picture_refs`` is the picture-labeled subset.
     """
 
     index: int | None = None
     document_id: str
     chunk_id: str
+    chunk_ids: list[str] = Field(default_factory=list)
     document_uri: str
     document_title: str | None = None
     page_numbers: list[int] = Field(default_factory=list)
     headings: list[str] | None = None
     content: str
+    doc_item_refs: list[str] = Field(default_factory=list)
     picture_refs: list[str] = Field(default_factory=list)
 
 
@@ -51,11 +61,13 @@ def resolve_citations(
             Citation(
                 document_id=r.document_id or "",
                 chunk_id=chunk_id,
+                chunk_ids=r.chunk_ids or [chunk_id],
                 document_uri=r.document_uri or "",
                 document_title=r.document_title,
                 page_numbers=r.page_numbers,
                 headings=r.headings,
                 content=r.content,
+                doc_item_refs=list(r.doc_item_refs),
                 picture_refs=picture_refs,
             )
         )

tests/store/test_citation.py

@@ -0,0 +1,41 @@
+from haiku.rag.store.models.chunk import SearchResult
+from haiku.rag.store.models.citation import resolve_citations
+
+
+def _result(chunk_id: str, chunk_ids: list[str] | None = None) -> SearchResult:
+    return SearchResult(
+        content="content",
+        score=0.9,
+        chunk_id=chunk_id,
+        chunk_ids=chunk_ids or [],
+        document_id="doc-1",
+        document_uri="test://doc",
+    )
+
+
+def test_resolve_citations_copies_merged_chunk_ids():
+    result = _result("c1", chunk_ids=["c1", "c2"])
+    citations = resolve_citations(["c1"], [result])
+    assert len(citations) == 1
+    assert citations[0].chunk_id == "c1"
+    assert citations[0].chunk_ids == ["c1", "c2"]
+
+
+def test_resolve_citations_falls_back_to_chunk_id():
+    result = _result("c1")
+    citations = resolve_citations(["c1"], [result])
+    assert len(citations) == 1
+    assert citations[0].chunk_ids == ["c1"]
+
+
+def test_resolve_citations_strips_brackets():
+    result = _result("c1")
+    citations = resolve_citations(["[c1]"], [result])
+    assert len(citations) == 1
+    assert citations[0].chunk_id == "c1"
+
+
+def test_resolve_citations_skips_unknown_ids():
+    result = _result("c1")
+    citations = resolve_citations(["c1", "missing"], [result])
+    assert len(citations) == 1

tests/test_client.py

@@ -1590,6 +1590,269 @@ async def test_client_visualize_chunk_multi_page(temp_db_path):
             assert img.tobytes() != blank.tobytes()
 
 
+async def test_client_visualize_chunk_merged_chunks_union_pages(temp_db_path):
+    """Visualizing all chunks of a merged result covers the union of their
+    expansions, which a single constituent chunk alone does not reach."""
+    from docling_core.types.doc.base import BoundingBox, Size
+    from docling_core.types.doc.document import (
+        DoclingDocument,
+        ImageRef,
+        ProvenanceItem,
+    )
+    from docling_core.types.doc.labels import DocItemLabel
+    from PIL import Image as PilImageModule
+
+    docling_doc = DoclingDocument(name="merged-viz-test")
+    page_size = Size(width=612.0, height=792.0)
+    for page_no in (1, 2):
+        docling_doc.add_page(
+            page_no=page_no,
+            size=page_size,
+            image=ImageRef.from_pil(
+                PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
+            ),
+        )
+
+    # Two sections, one per page, each large enough to be returned whole and
+    # under budget (so no cross-boundary expansion and no clip). A single
+    # chunk visualizes its own section's page; both chunks cover both pages.
+    layout = [
+        (DocItemLabel.SECTION_HEADER, "Section One", 1),
+        (DocItemLabel.PARAGRAPH, "Page one body. " + "x" * 3000, 1),
+        (DocItemLabel.SECTION_HEADER, "Section Two", 2),
+        (DocItemLabel.PARAGRAPH, "Page two body. " + "y" * 3000, 2),
+    ]
+    for i, (label, text, page_no) in enumerate(layout):
+        docling_doc.add_text(
+            label=label,
+            text=text,
+            prov=ProvenanceItem(
+                page_no=page_no,
+                bbox=BoundingBox(l=50, t=700 - (i % 2) * 100, r=550, b=650),
+                charspan=(0, 20),
+            ),
+        )
+
+    chunks = [
+        Chunk(
+            content="Page one body. " + "x" * 3000,
+            metadata={
+                "doc_item_refs": ["#/texts/1"],
+                "page_numbers": [1],
+                "labels": ["paragraph"],
+            },
+            order=0,
+            embedding=[0.1] * 2560,
+        ),
+        Chunk(
+            content="Page two body. " + "y" * 3000,
+            metadata={
+                "doc_item_refs": ["#/texts/3"],
+                "page_numbers": [2],
+                "labels": ["paragraph"],
+            },
+            order=1,
+            embedding=[0.1] * 2560,
+        ),
+    ]
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(docling_doc, chunks, uri="test://merged")
+
+        stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
+        stored_chunks.sort(key=lambda c: c.order)
+        assert len(stored_chunks) == 2
+        c1, c2 = stored_chunks
+
+        solo_images = await client.visualize_chunk(c1)
+        assert len(solo_images) == 1
+
+        merged_images = await client.visualize_chunk([c1, c2])
+        assert len(merged_images) == 2
+
+
+async def test_client_visualize_chunk_two_tone_highlights(temp_db_path):
+    """Matched content draws stronger than context swept in by expansion."""
+    from docling_core.types.doc.base import BoundingBox, Size
+    from docling_core.types.doc.document import (
+        DoclingDocument,
+        ImageRef,
+        ProvenanceItem,
+    )
+    from docling_core.types.doc.labels import DocItemLabel
+    from PIL import Image as PilImageModule
+
+    docling_doc = DoclingDocument(name="two-tone-test")
+    page_size = Size(width=612.0, height=792.0)
+    docling_doc.add_page(
+        page_no=1,
+        size=page_size,
+        image=ImageRef.from_pil(
+            PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
+        ),
+    )
+
+    # Three small paragraphs; the chunk matches only the middle one, so
+    # expansion sweeps in its neighbors.
+    for i in range(3):
+        docling_doc.add_text(
+            label=DocItemLabel.PARAGRAPH,
+            text=f"Paragraph {i}.",
+            prov=ProvenanceItem(
+                page_no=1,
+                bbox=BoundingBox(l=50, t=700 - i * 100, r=550, b=650 - i * 100),
+                charspan=(0, 12),
+            ),
+        )
+
+    chunks = [
+        Chunk(
+            content="Paragraph 1.",
+            metadata={
+                "doc_item_refs": ["#/texts/1"],
+                "page_numbers": [1],
+                "labels": ["paragraph"],
+            },
+            order=0,
+            embedding=[0.1] * 2560,
+        )
+    ]
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(docling_doc, chunks, uri="test://two-tone")
+        stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
+        assert len(stored_chunks) == 1
+
+        images = await client.visualize_chunk(stored_chunks[0])
+        assert len(images) == 1
+        image = images[0]
+
+        # Page dpi 72 == document coords, bottom-left origin flipped to
+        # top-left: item i's box spans y = 92 + i * 100 .. 142 + i * 100.
+        matched = image.getpixel((300, 217))  # inside #/texts/1
+        swept = image.getpixel((300, 117))  # inside #/texts/0
+        background = image.getpixel((300, 30))  # outside all boxes
+
+        assert matched != background
+        assert swept != background
+        assert matched != swept
+
+
+async def test_client_visualize_chunk_uses_given_refs(temp_db_path):
+    """Explicit refs (the citation's doc_item_refs) restrict the visualization
+    to exactly those items, instead of re-expanding the chunk's context."""
+    from docling_core.types.doc.base import BoundingBox, Size
+    from docling_core.types.doc.document import (
+        DoclingDocument,
+        ImageRef,
+        ProvenanceItem,
+    )
+    from docling_core.types.doc.labels import DocItemLabel
+    from PIL import Image as PilImageModule
+
+    docling_doc = DoclingDocument(name="refs-test")
+    page_size = Size(width=612.0, height=792.0)
+    for page_no in (1, 2):
+        docling_doc.add_page(
+            page_no=page_no,
+            size=page_size,
+            image=ImageRef.from_pil(
+                PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
+            ),
+        )
+    docling_doc.add_text(
+        label=DocItemLabel.PARAGRAPH,
+        text="Content on page one.",
+        prov=ProvenanceItem(
+            page_no=1, bbox=BoundingBox(l=50, t=700, r=550, b=650), charspan=(0, 20)
+        ),
+    )
+    docling_doc.add_text(
+        label=DocItemLabel.PARAGRAPH,
+        text="Content on page two.",
+        prov=ProvenanceItem(
+            page_no=2, bbox=BoundingBox(l=50, t=700, r=550, b=650), charspan=(0, 20)
+        ),
+    )
+
+    chunks = [
+        Chunk(
+            content="Content on page one.\nContent on page two.",
+            metadata={
+                "doc_item_refs": ["#/texts/0", "#/texts/1"],
+                "page_numbers": [1, 2],
+                "labels": ["paragraph", "paragraph"],
+            },
+            order=0,
+            embedding=[0.1] * 2560,
+        )
+    ]
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(docling_doc, chunks, uri="test://refs")
+        chunk = (await client.chunk_repository.get_by_document_id(doc.id))[0]
+
+        # No refs: re-expands the chunk's own refs → both pages.
+        assert len(await client.visualize_chunk(chunk)) == 2
+        # Given only the page-one ref → only page one is rendered.
+        assert len(await client.visualize_chunk(chunk, refs=["#/texts/0"])) == 1
+
+
+async def test_client_visualize_chunk_no_expand_shows_only_chunk(temp_db_path):
+    """expand=False draws only the chunk's own items, not the expanded section."""
+    from docling_core.types.doc.base import BoundingBox, Size
+    from docling_core.types.doc.document import (
+        DoclingDocument,
+        ImageRef,
+        ProvenanceItem,
+    )
+    from docling_core.types.doc.labels import DocItemLabel
+    from PIL import Image as PilImageModule
+
+    docling_doc = DoclingDocument(name="no-expand-test")
+    page_size = Size(width=612.0, height=792.0)
+    for page_no in (1, 2):
+        docling_doc.add_page(
+            page_no=page_no,
+            size=page_size,
+            image=ImageRef.from_pil(
+                PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
+            ),
+        )
+    for page_no in (1, 2):
+        docling_doc.add_text(
+            label=DocItemLabel.PARAGRAPH,
+            text=f"Short paragraph on page {page_no}.",
+            prov=ProvenanceItem(
+                page_no=page_no,
+                bbox=BoundingBox(l=50, t=700, r=550, b=650),
+                charspan=(0, 20),
+            ),
+        )
+
+    chunks = [
+        Chunk(
+            content="Short paragraph on page 1.",
+            metadata={
+                "doc_item_refs": ["#/texts/0"],
+                "page_numbers": [1],
+                "labels": ["paragraph"],
+            },
+            order=0,
+            embedding=[0.1] * 2560,
+        )
+    ]
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.import_document(docling_doc, chunks, uri="test://no-expand")
+        chunk = (await client.chunk_repository.get_by_document_id(doc.id))[0]
+
+        # Default expands the chunk's context outward → reaches page two.
+        assert len(await client.visualize_chunk(chunk)) == 2
+        # expand=False draws only the chunk's own page-one item.
+        assert len(await client.visualize_chunk(chunk, expand=False)) == 1
+
+
 # =============================================================================
 # convert() method tests
 # =============================================================================

tests/test_context.py

@@ -244,6 +244,46 @@ class TestFindExpansionRange:
         # Should expand outward into the first section
         assert hi >= 2
 
+    def test_picture_in_small_section_stays_section_bounded(self):
+        items = [
+            _item(0, label="section_header", text="Chapter 1"),
+            _item(1, text="Chapter 1 prose. " * 100),
+            _item(2, label="section_header", text="Figure heading"),
+            _item(3, label="picture", text="Diagram description."),
+            _item(4, label="caption", text="Figure 2-3. Balance arm."),
+            _item(5, label="section_header", text="Chapter 3"),
+            _item(6, text="Chapter 3 prose. " * 100),
+        ]
+        # Figure section (items 2-4) is far under 20% of 5000 chars.
+        lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000)
+        # Never crosses either header
+        assert (lo, hi) == (2, 4)
+
+    def test_table_in_small_section_stays_section_bounded(self):
+        items = [
+            _item(0, label="section_header", text="Chapter 1"),
+            _item(1, text="Chapter 1 prose. " * 100),
+            _item(2, label="section_header", text="Table heading"),
+            _item(3, label="table", text="Header | Value"),
+            _item(4, label="section_header", text="Chapter 3"),
+            _item(5, text="Chapter 3 prose. " * 100),
+        ]
+        lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000)
+        assert (lo, hi) == (2, 3)
+
+    def test_text_in_small_section_still_expands_outward(self):
+        items = [
+            _item(0, label="section_header", text="Chapter 1"),
+            _item(1, text="Chapter 1 prose. " * 100),
+            _item(2, label="section_header", text="Short note"),
+            _item(3, text="A brief remark."),
+            _item(4, label="section_header", text="Chapter 3"),
+            _item(5, text="Chapter 3 prose. " * 100),
+        ]
+        lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000)
+        # Text hits keep growing across section boundaries
+        assert lo < 2 or hi > 3
+
 
 class TestEvidenceAnchors:
     def test_empty_content(self):
@@ -369,6 +409,85 @@ class TestExpandWithItems:
             # real content — so we get expanded content, not the fallback.
             assert len(expanded[0].content) > 0
 
+    async def test_picture_expansion_stays_within_section_pages(self, temp_db_path):
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=0,
+                    self_ref="#/texts/0",
+                    label="section_header",
+                    text="Chapter 1",
+                    page_numbers=[10],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="text",
+                    text="Chapter 1 prose. " * 200,
+                    page_numbers=[10],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=2,
+                    self_ref="#/texts/2",
+                    label="section_header",
+                    text="Figure heading",
+                    page_numbers=[13],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=3,
+                    self_ref="#/pictures/0",
+                    label="picture",
+                    text="Diagram description.",
+                    page_numbers=[13],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=4,
+                    self_ref="#/texts/3",
+                    label="caption",
+                    text="Figure 2-3. Balance arm.",
+                    page_numbers=[13],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=5,
+                    self_ref="#/texts/4",
+                    label="section_header",
+                    text="Chapter 3",
+                    page_numbers=[15],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=6,
+                    self_ref="#/texts/5",
+                    label="text",
+                    text="Chapter 3 prose. " * 200,
+                    page_numbers=[15],
+                ),
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            result = SearchResult(
+                content="Diagram description.",
+                score=0.9,
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/0"],
+                page_numbers=[13],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [result], 5000
+            )
+            assert len(expanded) == 1
+            assert expanded[0].page_numbers == [13]
+            assert "Chapter 1 prose" not in expanded[0].content
+            assert "Chapter 3 prose" not in expanded[0].content
+
     async def test_fragmented_items_preserve_chunk(self, temp_db_path):
         """When items are fragmented (e.g., list_item children), the original
         chunk content is preserved if expansion produces less text."""
@@ -578,6 +697,187 @@ class TestExpandWithItems:
             assert len(expanded[0].content) <= 5000
             assert marker in expanded[0].content
 
+    async def test_solo_result_carries_own_chunk_id(self, temp_db_path):
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=i,
+                    self_ref=f"#/texts/{i}",
+                    label="text",
+                    text=f"Paragraph {i}. " * 10,
+                )
+                for i in range(5)
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            result = SearchResult(
+                content="Paragraph 2.",
+                score=0.9,
+                chunk_id="c1",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/2"],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [result], 5000
+            )
+            assert len(expanded) == 1
+            assert expanded[0].chunk_ids == ["c1"]
+
+    async def test_merged_results_carry_all_chunk_ids(self, temp_db_path):
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=i,
+                    self_ref=f"#/texts/{i}",
+                    label="text",
+                    text=f"Paragraph {i}. " * 10,
+                )
+                for i in range(5)
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            r1 = SearchResult(
+                content="Paragraph 1.",
+                score=0.9,
+                chunk_id="c1",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/1"],
+            )
+            r2 = SearchResult(
+                content="Paragraph 3.",
+                score=0.85,
+                chunk_id="c2",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/3"],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r1, r2], 5000
+            )
+            # Ranges around positions 1 and 3 overlap → one merged result.
+            assert len(expanded) == 1
+            assert expanded[0].chunk_id == "c1"
+            assert expanded[0].chunk_ids == ["c1", "c2"]
+            # Sibling chunk ids are plumbing for visualization, never shown
+            # to the model.
+            assert "c2" not in expanded[0].format_for_agent()
+
+    async def test_merged_anchor_is_highest_scoring_constituent(self, temp_db_path):
+        """A merged result's chunk_id anchors on the best-scoring constituent,
+        not whichever chunk sits earliest in the document."""
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=i,
+                    self_ref=f"#/texts/{i}",
+                    label="text",
+                    text=f"Paragraph {i}. " * 10,
+                )
+                for i in range(5)
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            # earlier in the document, lower score
+            r_early = SearchResult(
+                content="Paragraph 1.",
+                score=0.40,
+                chunk_id="c-early",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/1"],
+            )
+            # later in the document, higher score — the real hit
+            r_best = SearchResult(
+                content="Paragraph 3.",
+                score=0.95,
+                chunk_id="c-best",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/3"],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r_early, r_best], 5000
+            )
+            assert len(expanded) == 1
+            assert expanded[0].chunk_id == "c-best"
+            assert expanded[0].score == 0.95
+            # provenance still lists both
+            assert set(expanded[0].chunk_ids) == {"c-early", "c-best"}
+
+    async def test_clipped_merged_result_keeps_anchor_evidence_and_pages(
+        self, temp_db_path
+    ):
+        """When a merged result is clipped to budget, the surviving window is
+        centered on the anchor (highest-scoring) chunk, and page_numbers reflect
+        only the content that survived — not the full merged range."""
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=0,
+                    self_ref="#/texts/0",
+                    label="text",
+                    text="LOWMARK " + "a" * 400,
+                    page_numbers=[1],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="text",
+                    text="b" * 400,
+                    page_numbers=[2],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=2,
+                    self_ref="#/texts/2",
+                    label="text",
+                    text="c" * 400 + " HIGHMARK",
+                    page_numbers=[3],
+                ),
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            r_low = SearchResult(
+                content="LOWMARK " + "a" * 400,
+                score=0.4,
+                chunk_id="c-low",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/0"],
+                page_numbers=[1],
+            )
+            r_high = SearchResult(
+                content="c" * 400 + " HIGHMARK",
+                score=0.9,
+                chunk_id="c-high",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/2"],
+                page_numbers=[3],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r_low, r_high], 500
+            )
+            assert len(expanded) == 1
+            e = expanded[0]
+            # anchor is the high-scoring chunk, and its evidence survives clipping
+            assert e.chunk_id == "c-high"
+            assert "HIGHMARK" in e.content
+            assert "LOWMARK" not in e.content
+            # page_numbers reflect only the surviving window, not the full range
+            assert 3 in e.page_numbers
+            assert 1 not in e.page_numbers
+            # refs likewise exclude the clipped-out item
+            assert "#/texts/0" not in e.doc_item_refs
+
     async def test_fuzzy_match_preserves_central_marker(self, temp_db_path):
         """The chunk's text need not be verbatim in the joined item text: a clean
         central marker is still located via the central-slice anchor."""
@@ -715,3 +1015,61 @@ class TestExpandWithItemsPictureBytes:
                 "#/pictures/1": "A",
                 "#/pictures/3": "B",
             }
+
+    async def test_clipped_out_picture_bytes_dropped(self, temp_db_path):
+        """A lower-scoring picture chunk clipped out of the budget window no
+        longer contributes its image bytes — the model must not receive an
+        image the citation and visualization omit."""
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=0,
+                    self_ref="#/pictures/0",
+                    label="picture",
+                    text="LOWPIC " + "a" * 400,
+                    page_numbers=[1],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="text",
+                    text="b" * 400,
+                    page_numbers=[2],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=2,
+                    self_ref="#/pictures/1",
+                    label="picture",
+                    text="c" * 400 + " HIGHPIC",
+                    page_numbers=[3],
+                ),
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            r_low = SearchResult(
+                content="LOWPIC " + "a" * 400,
+                score=0.4,
+                chunk_id="c-low",
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/0"],
+                image_data={"#/pictures/0": "LOWBYTES"},
+            )
+            r_high = SearchResult(
+                content="c" * 400 + " HIGHPIC",
+                score=0.9,
+                chunk_id="c-high",
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/1"],
+                image_data={"#/pictures/1": "HIGHBYTES"},
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r_low, r_high], 500
+            )
+            assert len(expanded) == 1
+            assert "#/pictures/0" not in expanded[0].doc_item_refs
+            assert expanded[0].image_data == {"#/pictures/1": "HIGHBYTES"}

tests/test_skill_extras.py

@@ -0,0 +1,42 @@
+"""Tests for non-tool utilities from ``haiku.rag.skills._tools.create_skill_extras``."""
+
+from docling_core.types.doc.document import DoclingDocument
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import AppConfig
+from haiku.rag.skills._tools import create_skill_extras
+from haiku.rag.store.models.chunk import Chunk
+
+
+async def _seed_chunk(db_path) -> str:
+    docling_doc = DoclingDocument(name="extras-test")
+    chunk = Chunk(
+        content="Some content.",
+        metadata={"doc_item_refs": ["#/texts/0"], "page_numbers": [1]},
+        order=0,
+        embedding=[0.1] * 2560,
+    )
+    async with HaikuRAG(db_path, create=True) as client:
+        doc = await client.import_document(docling_doc, [chunk], uri="test://extras")
+        stored = await client.chunk_repository.get_by_document_id(doc.id)
+        return stored[0].id
+
+
+async def test_extras_visualize_chunk_accepts_str_and_list(temp_db_path):
+    chunk_id = await _seed_chunk(temp_db_path)
+    extras = create_skill_extras(temp_db_path, AppConfig())
+    visualize_chunk = extras["visualize_chunk"]
+
+    # A document imported without page images yields no visualizations, but the
+    # str and list inputs must both resolve the chunk and reach visualize_chunk.
+    assert await visualize_chunk(chunk_id) == []
+    assert await visualize_chunk([chunk_id]) == []
+
+
+async def test_extras_visualize_chunk_unknown_id_returns_empty(temp_db_path):
+    await _seed_chunk(temp_db_path)
+    extras = create_skill_extras(temp_db_path, AppConfig())
+    visualize_chunk = extras["visualize_chunk"]
+
+    assert await visualize_chunk("does-not-exist") == []
+    assert await visualize_chunk(["does-not-exist"]) == []