ChatDeps implements StateHandler protocol for proper AG-UI state management

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

@@ -1,6 +1,6 @@
 import hashlib
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 from numpy.typing import NDArray
@@ -156,7 +156,10 @@ async def rank_qa_history_by_similarity(
 
 @dataclass
 class ChatDeps:
-    """Dependencies for chat agent."""
+    """Dependencies for chat agent.
+
+    Implements StateHandler protocol for AG-UI state management.
+    """
 
     client: HaikuRAG
     config: AppConfig
@@ -164,6 +167,44 @@ class ChatDeps:
     session_state: ChatSessionState | None = None
     state_key: str | None = None
 
+    @property
+    def state(self) -> dict[str, Any] | None:
+        """Get current state for AG-UI protocol."""
+        if self.session_state is None:
+            return None
+        snapshot = self.session_state.model_dump()
+        if self.state_key:
+            return {self.state_key: snapshot}
+        return snapshot
+
+    @state.setter
+    def state(self, value: dict[str, Any] | None) -> None:
+        """Set state from AG-UI protocol."""
+        if value is None:
+            return
+        # Extract from namespaced key if present
+        state_data: dict[str, Any] = value
+        if self.state_key and self.state_key in value:
+            nested = value[self.state_key]
+            if isinstance(nested, dict):
+                state_data = nested
+        # Update session_state from incoming state
+        if self.session_state is not None:
+            if "qa_history" in state_data:
+                self.session_state.qa_history = [
+                    QAResponse(**qa) if isinstance(qa, dict) else qa
+                    for qa in state_data.get("qa_history", [])
+                ]
+            if "citations" in state_data:
+                self.session_state.citations = [
+                    CitationInfo(**c) if isinstance(c, dict) else c
+                    for c in state_data.get("citations", [])
+                ]
+            if "initial_context" in state_data:
+                self.session_state.initial_context = state_data.get("initial_context")
+            if "session_id" in state_data:
+                self.session_state.session_id = state_data.get("session_id", "")
+
 
 @dataclass
 class SearchDeps:

tests/agents/chat/test_state.py

@@ -267,3 +267,206 @@ def test_chat_session_state_initial_context_defaults_to_none():
 
     state = ChatSessionState(session_id="test-session")
     assert state.initial_context is None
+
+
+def test_chat_deps_state_getter_returns_namespaced_state():
+    """Test ChatDeps.state getter returns state under namespaced key."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    session_state = ChatSessionState(
+        session_id="test-123",
+        qa_history=[
+            QAResponse(question="Q1", answer="A1", confidence=0.9),
+        ],
+        initial_context="Background info",
+    )
+
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=session_state,
+        state_key=AGUI_STATE_KEY,
+    )
+
+    state = deps.state
+    assert state is not None
+    assert AGUI_STATE_KEY in state
+    assert state[AGUI_STATE_KEY]["session_id"] == "test-123"
+    assert len(state[AGUI_STATE_KEY]["qa_history"]) == 1
+    assert state[AGUI_STATE_KEY]["qa_history"][0]["question"] == "Q1"
+    assert state[AGUI_STATE_KEY]["initial_context"] == "Background info"
+
+
+def test_chat_deps_state_getter_without_namespace():
+    """Test ChatDeps.state getter returns flat state when no state_key."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    session_state = ChatSessionState(session_id="test-123")
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=session_state,
+        state_key=None,
+    )
+
+    state = deps.state
+    assert state is not None
+    assert "session_id" in state
+    assert state["session_id"] == "test-123"
+
+
+def test_chat_deps_state_getter_returns_none_without_session():
+    """Test ChatDeps.state getter returns None when no session_state."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import ChatDeps
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=None,
+    )
+
+    assert deps.state is None
+
+
+def test_chat_deps_state_setter_updates_from_namespaced_state():
+    """Test ChatDeps.state setter updates session_state from namespaced incoming state."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    session_state = ChatSessionState(session_id="initial")
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=session_state,
+        state_key=AGUI_STATE_KEY,
+    )
+
+    # Simulate incoming AG-UI state with namespaced key
+    incoming_state = {
+        AGUI_STATE_KEY: {
+            "session_id": "updated-123",
+            "qa_history": [
+                {"question": "Q1", "answer": "A1", "confidence": 0.9, "citations": []}
+            ],
+            "citations": [],
+            "initial_context": "New context",
+        }
+    }
+
+    deps.state = incoming_state
+
+    assert deps.session_state is not None
+    assert deps.session_state.session_id == "updated-123"
+    assert len(deps.session_state.qa_history) == 1
+    assert deps.session_state.qa_history[0].question == "Q1"
+    assert deps.session_state.initial_context == "New context"
+
+
+def test_chat_deps_state_setter_handles_none():
+    """Test ChatDeps.state setter handles None gracefully."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    session_state = ChatSessionState(session_id="original")
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=session_state,
+    )
+
+    # Setting None should not raise and should not change state
+    deps.state = None
+
+    assert deps.session_state is not None
+    assert deps.session_state.session_id == "original"
+
+
+def test_chat_deps_state_setter_without_session_state():
+    """Test ChatDeps.state setter does nothing when session_state is None."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import ChatDeps
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=None,
+    )
+
+    # Should not raise even with valid incoming state
+    deps.state = {"session_id": "test", "qa_history": [], "citations": []}
+
+    assert deps.session_state is None
+
+
+def test_chat_deps_state_setter_with_citation_dicts():
+    """Test ChatDeps.state setter converts citation dicts to CitationInfo."""
+    from unittest.mock import MagicMock
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
+
+    mock_client = MagicMock()
+    mock_config = MagicMock()
+
+    session_state = ChatSessionState(session_id="test")
+    deps = ChatDeps(
+        client=mock_client,
+        config=mock_config,
+        session_state=session_state,
+        state_key=AGUI_STATE_KEY,
+    )
+
+    incoming_state = {
+        AGUI_STATE_KEY: {
+            "session_id": "test",
+            "qa_history": [],
+            "citations": [
+                {
+                    "index": 1,
+                    "document_id": "doc-1",
+                    "chunk_id": "chunk-1",
+                    "document_uri": "test.md",
+                    "document_title": "Test Doc",
+                    "page_numbers": [1, 2],
+                    "headings": ["Intro"],
+                    "content": "Test content",
+                }
+            ],
+            "initial_context": None,
+        }
+    }
+
+    deps.state = incoming_state
+
+    assert deps.session_state is not None
+    assert len(deps.session_state.citations) == 1
+    citation = deps.session_state.citations[0]
+    assert citation.document_id == "doc-1"
+    assert citation.chunk_id == "chunk-1"
+    assert citation.page_numbers == [1, 2]

tests/test_cli.py

@@ -283,6 +283,7 @@ def test_ask():
             cite=False,
             deep=False,
             filter=None,
+            initial_context=None,
         )
 
 
@@ -300,6 +301,7 @@ def test_ask_with_cite():
             cite=True,
             deep=False,
             filter=None,
+            initial_context=None,
         )
 
 
@@ -317,6 +319,7 @@ def test_ask_with_deep():
             cite=False,
             deep=True,
             filter=None,
+            initial_context=None,
         )
 
 
@@ -334,6 +337,7 @@ def test_ask_with_deep_and_cite():
             cite=True,
             deep=True,
             filter=None,
+            initial_context=None,
         )