Tests, widget to show context in tui

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -25,6 +25,9 @@ from haiku.rag.utils import get_model
 
 logger = logging.getLogger(__name__)
 
+# Track background tasks to prevent garbage collection
+_background_tasks: set[asyncio.Task[None]] = set()
+
 
 async def _update_context_background(
     qa_history: list[QAResponse],
@@ -38,7 +41,7 @@ async def _update_context_background(
             config=config,
             session_state=session_state,
         )
-        logger.debug("Session context updated successfully")
+        logger.debug("Session context updated")
     except Exception:
         logger.exception("Failed to update session context")
 
@@ -286,13 +289,15 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
                 ]
 
             # Spawn background task to update session context
-            asyncio.create_task(
+            task = asyncio.create_task(
                 _update_context_background(
                     qa_history=list(ctx.deps.session_state.qa_history),
                     config=ctx.deps.config,
                     session_state=ctx.deps.session_state,
                 )
             )
+            _background_tasks.add(task)
+            task.add_done_callback(_background_tasks.discard)
 
         # Build new state with citations AND accumulated qa_history
         new_state = ChatSessionState(

haiku_rag_slim/haiku/rag/chat/app.py

@@ -33,6 +33,7 @@ try:
     import logfire
 
     logfire.configure(send_to_logfire="if-token-present", console=False)
+    logfire.instrument_pydantic_ai()
 except ImportError:
     pass
 
@@ -81,6 +82,7 @@ class ChatApp(App):
         Binding("ctrl+l", "clear_chat", "Clear", show=True),
         Binding("ctrl+g", "show_visual", "Visual", show=True),
         Binding("ctrl+i", "show_info", "Info", show=True),
+        Binding("ctrl+o", "show_context", "Context", show=True),
         Binding("escape", "focus_input", "Focus Input", show=False),
     ]
 
@@ -326,6 +328,12 @@ class ChatApp(App):
 
         await self.push_screen(InfoModal(self.client, self.db_path))
 
+    async def action_show_context(self) -> None:
+        """Show current session context in a modal."""
+        from haiku.rag.chat.widgets.context_modal import ContextModal
+
+        await self.push_screen(ContextModal(self.session_state))
+
     def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
         """Handle citation selection."""
         chat_history = self.query_one(ChatHistory)

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

@@ -0,0 +1,76 @@
+from typing import TYPE_CHECKING
+
+from textual.app import ComposeResult
+from textual.binding import Binding
+from textual.containers import Vertical, VerticalScroll
+from textual.screen import ModalScreen
+from textual.widgets import Markdown, Static
+
+if TYPE_CHECKING:
+    from haiku.rag.agents.chat.state import ChatSessionState
+
+
+class ContextModal(ModalScreen):  # pragma: no cover
+    """Modal screen for displaying session context."""
+
+    BINDINGS = [
+        Binding("escape", "dismiss", "Close", show=True),
+        Binding("ctrl+o", "dismiss", "Close", show=True),
+    ]
+
+    CSS = """
+    ContextModal {
+        align: center middle;
+        background: rgba(0, 0, 0, 0.5);
+    }
+
+    #context-container {
+        width: auto;
+        min-width: 40;
+        max-width: 80;
+        height: auto;
+        max-height: 20;
+        background: $surface;
+        border: tall $primary;
+        padding: 1 2;
+    }
+
+    #context-header {
+        height: auto;
+        margin-bottom: 1;
+    }
+
+    #context-content {
+        height: 1fr;
+        scrollbar-gutter: stable;
+    }
+    """
+
+    def __init__(self, session_state: "ChatSessionState | None"):
+        super().__init__()
+        self.session_state = session_state
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="context-container"):
+            yield Static("[bold]Session Context[/bold]", id="context-header")
+            with VerticalScroll(id="context-content"):
+                yield Markdown(self._get_content())
+
+    def _get_content(self) -> str:
+        if not self.session_state:
+            return "*No session state.*"
+
+        if not self.session_state.session_context:
+            return "*No session context yet. Ask a question first.*"
+
+        ctx = self.session_state.session_context
+        updated = (
+            ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S")
+            if ctx.last_updated
+            else "unknown"
+        )
+
+        return f"**Last updated:** {updated}\n\n---\n\n{ctx.summary}"
+
+    async def action_dismiss(self, result=None) -> None:
+        self.app.pop_screen()

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

@@ -25,14 +25,17 @@ class InfoModal(ModalScreen):  # pragma: no cover
     CSS = """
     InfoModal {
         align: center middle;
+        background: rgba(0, 0, 0, 0.5);
     }
 
     #info-container {
-        width: 80;
+        width: auto;
+        min-width: 40;
+        max-width: 80;
         height: auto;
-        max-height: 80%;
+        max-height: 20;
         background: $surface;
-        border: solid $primary;
+        border: tall $primary;
         padding: 1 2;
     }
 
@@ -42,8 +45,8 @@ class InfoModal(ModalScreen):  # pragma: no cover
     }
 
     #info-content {
-        height: auto;
-        max-height: 100%;
+        height: 1fr;
+        scrollbar-gutter: stable;
     }
     """
 

tests/agents/chat/test_chat_agent.py

@@ -584,6 +584,54 @@ async def test_chat_agent_ask_with_state_key(allow_model_requests, temp_db_path)
         assert len(session_state.qa_history) >= 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_ask_triggers_background_summarization(
+    allow_model_requests, temp_db_path
+):
+    """Test that the ask tool triggers background session context summarization."""
+    import asyncio
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+
+        agent = create_chat_agent(Config)
+        session_state = ChatSessionState(session_id="test-summarization")
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Initially no session_context
+        assert session_state.session_context is None
+
+        # Ask a question
+        result = await agent.run(
+            "What is the highest count class in the DocLayNet dataset?",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        assert len(session_state.qa_history) >= 1
+
+        # Wait for background task to complete
+        # The task should update session_state.session_context
+        for _ in range(50):  # Wait up to 5 seconds
+            if session_state.session_context is not None:
+                break
+            await asyncio.sleep(0.1)
+
+        # Verify session_context was populated by background task
+        assert session_state.session_context is not None
+        assert session_state.session_context.summary != ""
+        assert session_state.session_context.last_updated is not None
+
+
 def test_fifo_limit_enforcement():
     """Test that FIFO limit enforcement logic works correctly.