danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #264 from ggozad/fix/tui

Browse source 
Fix TUI session handling
This commit is contained in:
Yiorgis Gozadinos 2026-01-29 12:26:10 +02:00 committed by GitHub
commit 51e6be0d68
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 91 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -22,6 +22,9 @@
- First request still sends full snapshot; subsequent requests send only changes
- Backend logging shows incoming/outgoing state events for debugging
### Fixed
- **Chat TUI Session State Sync**: TUI now syncs full session state from AG-UI events
## [0.27.1] - 2026-01-27

19
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -23,7 +23,6 @@ from haiku.rag.agents.chat.state import (
ChatDeps,
ChatSessionState,
)
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
@ -103,7 +102,6 @@ class ChatApp(App):
self.session_state = ChatSessionState()
self._is_processing = False
self._tool_call_widgets: dict[str, Any] = {}
self._last_citations: list[Citation] = []
self._current_worker: Worker[None] | None = None
self._message_history: list[ModelMessage] = []
self._document_filter: list[str] = []
@ -170,6 +168,10 @@ class ChatApp(App):
if self.client:
await self.client.__aexit__(None, None, None)
def _sync_session_state(self, chat_state: dict[str, Any]) -> None:
"""Sync session_state from AG-UI state."""
self.session_state = ChatSessionState.model_validate(chat_state)
async def _handle_stream_event(self, event: AgentStreamEvent) -> None:
"""Handle streaming events from the agent."""
chat_history = self.query_one(ChatHistory)
@ -199,8 +201,7 @@ class ChatApp(App):
snapshot = getattr(meta_event, "snapshot", {})
self._agui_state_snapshot = snapshot
chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
citations = chat_state.get("citations", [])
self._last_citations = [Citation(**c) for c in citations]
self._sync_session_state(chat_state)
elif meta_event.type == EventType.STATE_DELTA:
delta = getattr(meta_event, "delta", [])
@ -212,8 +213,7 @@ class ChatApp(App):
chat_state = self._agui_state_snapshot.get(
AGUI_STATE_KEY, self._agui_state_snapshot
)
citations = chat_state.get("citations", [])
self._last_citations = [Citation(**c) for c in citations]
self._sync_session_state(chat_state)
async def _event_stream_handler(
self,
@ -247,7 +247,7 @@ class ChatApp(App):
# Clear for new query
self._tool_call_widgets.clear()
self._last_citations.clear()
self.session_state.citations.clear()
# Run agent in a worker to keep UI responsive
self._is_processing = True
@ -303,8 +303,8 @@ class ChatApp(App):
self._message_history = stream.all_messages()
# Add citations captured from tool metadata
if self._last_citations:
await chat_history.add_citations(self._last_citations)
if self.session_state.citations:
await chat_history.add_citations(self.session_state.citations)
except asyncio.CancelledError:
chat_history.hide_thinking()
@ -323,7 +323,6 @@ class ChatApp(App):
"""Clear the chat history and reset session."""
chat_history = self.query_one(ChatHistory)
await chat_history.clear_messages()
self._last_citations.clear()
self._message_history.clear()
self._agui_state_snapshot = {}
# Reset context lock and session state (reset to CLI value)

91
tests/chat/test_chat_app.py
View file

@ -337,9 +337,10 @@ async def test_handle_stream_event_extracts_citations_from_state_snapshot(
# Handle the event
await app._handle_stream_event(event)
# Citations should be extracted
assert len(app._last_citations) == 1
assert app._last_citations[0].chunk_id == "chunk1"
# Session state should be synced
assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citation_registry == {"chunk1": 1}
@pytest.mark.asyncio
@ -387,7 +388,7 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
)
event1 = FunctionToolResultEvent(result=tool_return1)
await app._handle_stream_event(event1)
assert len(app._last_citations) == 0
assert len(app.session_state.citations) == 0
# Now handle a STATE_DELTA event that adds citations
delta_event = StateDeltaEvent(
@ -425,10 +426,10 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
# Handle the delta event
await app._handle_stream_event(event2)
# Citations should be extracted from the delta
assert len(app._last_citations) == 1
assert app._last_citations[0].chunk_id == "chunk1"
assert app._last_citations[0].content == "Test content from delta"
# Session state should be synced
assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citations[0].content == "Test content from delta"
@pytest.mark.asyncio
@ -505,10 +506,76 @@ async def test_handle_stream_event_delta_with_preinitialized_state(
# Handle the delta event
await app._handle_stream_event(event)
# Citations should be extracted from the delta
assert len(app._last_citations) == 1
assert app._last_citations[0].chunk_id == "chunk1"
assert app._last_citations[0].content == "Content from first delta"
# Session state should be synced
assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citations[0].content == "Content from first delta"
@pytest.mark.asyncio
async def test_handle_stream_event_syncs_session_context(temp_db_path: Path):
"""Test that _handle_stream_event syncs session_context to session_state."""
from ag_ui.core import EventType, StateDeltaEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.chat.app import ChatApp
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test():
# Pre-initialize state
app._agui_state_snapshot = {
AGUI_STATE_KEY: {
"session_id": "test",
"citations": [],
"qa_history": [],
"citation_registry": {},
"document_filter": [],
"initial_context": None,
"session_context": None,
}
}
# Verify session_context starts as None
assert app.session_state.session_context is None
# Handle a delta that adds session_context
delta_event = StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=[
{
"op": "replace",
"path": f"/{AGUI_STATE_KEY}/session_context",
"value": {
"summary": "User asked about Python async patterns.",
"last_updated": "2025-01-15T10:30:00",
},
},
],
)
tool_return = ToolReturnPart(
tool_name="ask",
content="Answer about async",
tool_call_id="test-call-1",
metadata=[delta_event],
)
event = FunctionToolResultEvent(result=tool_return)
await app._handle_stream_event(event)
# Session context should be synced to session_state
assert app.session_state.session_context is not None
assert (
app.session_state.session_context.summary
== "User asked about Python async patterns."
)
@pytest.mark.asyncio