From 729a3910d85e38f8b2fd57d3d9c3ea793aeba0c9 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Mon, 26 Jan 2026 10:58:33 +0200
Subject: [PATCH] Fix visual grounding citation selection in TUI

---
 haiku_rag_slim/haiku/rag/chat/app.py | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py
index 64ce2c67..7091ead8 100644
--- a/haiku_rag_slim/haiku/rag/chat/app.py
+++ b/haiku_rag_slim/haiku/rag/chat/app.py
@@ -103,7 +103,6 @@ class ChatApp(App):
         self._is_processing = False
         self._tool_call_widgets: dict[str, Any] = {}
         self._last_citations: list[Citation] = []
-        self._selected_citation_idx: int | None = None
         self._current_worker: Worker[None] | None = None
         self._message_history: list[ModelMessage] = []
 
@@ -200,7 +199,6 @@ class ChatApp(App):
         # Clear for new query
         self._tool_call_widgets.clear()
         self._last_citations.clear()
-        self._selected_citation_idx = None
 
         # Run agent in a worker to keep UI responsive
         self._is_processing = True
@@ -271,7 +269,6 @@ class ChatApp(App):
         chat_history = self.query_one(ChatHistory)
         await chat_history.clear_messages()
         self._last_citations.clear()
-        self._selected_citation_idx = None
         self._message_history.clear()
         # Reset session state for fresh conversation
         self.session_state = ChatSessionState(
@@ -289,7 +286,6 @@ class ChatApp(App):
         chat_history = self.query_one(ChatHistory)
         for widget in chat_history.query(CitationWidget):
             widget.remove_class("selected")
-        self._selected_citation_idx = None
 
     def on_descendant_focus(self, _event: object) -> None:
         """Clear citation selection when input is focused."""
@@ -298,15 +294,16 @@ class ChatApp(App):
 
     async def action_show_visual(self) -> None:
         """Show visual grounding for the selected citation."""
-        if not self.client or not self._last_citations:
+        if not self.client:
             return
 
-        idx = (
-            self._selected_citation_idx
-            if self._selected_citation_idx is not None
-            else 0
-        )
-        citation = self._last_citations[idx]
+        # Get citation from selected widget directly
+        chat_history = self.query_one(ChatHistory)
+        selected_widgets = list(chat_history.query(CitationWidget).filter(".selected"))
+        if not selected_widgets:
+            return
+
+        citation = selected_widgets[0].citation
         chunk = await self.client.chunk_repository.get_by_id(citation.chunk_id)
         if not chunk:
             return
@@ -342,4 +339,3 @@ class ChatApp(App):
         citation_widgets = list(chat_history.query(CitationWidget))
         if 0 <= event.citation_index < len(citation_widgets):
             citation_widgets[event.citation_index].add_class("selected")
-            self._selected_citation_idx = event.citation_index