Fix citation selection when we call visualize

haiku_rag_slim/haiku/rag/chat/app.py

@@ -317,8 +317,8 @@ class ChatApp(App):
             widget.remove_class("selected")
 
     def on_descendant_focus(self, _event: object) -> None:
-        """Clear citation selection when input is focused."""
-        if isinstance(self.focused, Input):
+        """Clear citation selection when chat input is focused."""
+        if isinstance(self.focused, Input) and self.focused.id == "chat-input":
             self._clear_citation_selection()
 
     async def action_show_visual(self) -> None:
@@ -364,10 +364,8 @@ class ChatApp(App):
         for widget in chat_history.query(CitationWidget):
             widget.remove_class("selected")
 
-        # Add selected class to the newly selected citation
-        citation_widgets = list(chat_history.query(CitationWidget))
-        if 0 <= event.citation_index < len(citation_widgets):
-            citation_widgets[event.citation_index].add_class("selected")
+        # Add selected class to the widget that was focused
+        event.widget.add_class("selected")
 
     async def action_show_filter(self) -> None:
         """Show document filter modal."""

haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py

@@ -90,9 +90,9 @@ class CitationWidget(Collapsible):
     class Selected(Message):
         """Message sent when a citation is selected."""
 
-        def __init__(self, citation_index: int) -> None:
+        def __init__(self, widget: "CitationWidget") -> None:
             super().__init__()
-            self.citation_index = citation_index
+            self.widget = widget
 
     def __init__(self, citation: Citation, **kwargs) -> None:
         title = f"[{citation.index}] {citation.document_title or citation.document_uri}"
@@ -120,8 +120,7 @@ class CitationWidget(Collapsible):
 
     def on_focus(self) -> None:
         """When focused, mark as selected."""
-        index = self.citation.index or 1
-        self.post_message(self.Selected(index - 1))
+        self.post_message(self.Selected(self))
 
     def on_key(self, event: "Key") -> None:
         """Handle Enter to toggle expand/collapse."""