handle state_key duplication between deps and ToolContext

app/backend/main.py

@@ -13,7 +13,6 @@ from starlette.responses import JSONResponse, Response, StreamingResponse
 from starlette.routing import Route
 
 from haiku.rag.agents.chat import (
-    AGUI_STATE_KEY,
     ChatDeps,
     create_chat_agent,
     prepare_chat_context,
@@ -92,7 +91,6 @@ async def stream_chat(request: Request) -> Response:
         config=Config,
         client=get_client(),
         tool_context=context,
-        state_key=AGUI_STATE_KEY,
     )
 
     adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)

docs/tools.md

@@ -235,10 +235,12 @@ async with HaikuRAG("path/to/db.lancedb") as client:
         print(f"Total search results: {len(search_state.results)}")
 ```
 
-`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, pass a `state_key`:
+`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, set `state_key` on the `ToolContext` (via `prepare_context`):
 
 ```python
-deps = AgentDeps(client=client, tool_context=context, state_key="my_app")
+context = ToolContext()
+prepare_context(context, features=["search", "qa"], state_key="my_app")
+deps = AgentDeps(client=client, tool_context=context)
 ```
 
 Tool functions access `client` and `tool_context` via pydantic-ai's `RunContext.deps`, so toolsets can be created once and reused across requests.
@@ -247,7 +249,7 @@ All toolsets respect session-level document filters when a `SessionState` is reg
 
 ## AG-UI State Management
 
-Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. State is emitted under a namespaced key via `state_key`.
+Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. `ChatDeps` extends `AgentDeps` with chat-specific config and state handling. State is emitted under a namespaced key via `state_key` on the `ToolContext` — set it once via `prepare_context()`.
 
 **Custom agents** use `AgentDeps` + `prepare_context`:
 
@@ -256,14 +258,14 @@ from haiku.rag.tools import AgentDeps, ToolContext, ToolContextCache, prepare_co
 
 context = ToolContext()
 prepare_context(context, features=["search", "qa"], state_key="my_app")
-deps = AgentDeps(client=client, tool_context=context, state_key="my_app")
+deps = AgentDeps(client=client, tool_context=context)
 ```
 
 **Chat agent** uses `ChatDeps` + `prepare_chat_context` (adds chat-specific overrides like background summarization and initial context handling):
 
 ```python
 from haiku.rag.agents.chat import (
-    AGUI_STATE_KEY, ChatDeps, create_chat_agent, prepare_chat_context,
+    ChatDeps, create_chat_agent, prepare_chat_context,
 )
 from haiku.rag.tools import ToolContext, ToolContextCache
 
@@ -272,13 +274,12 @@ agent = create_chat_agent(config)
 # For multi-session apps, cache ToolContext per thread
 cache = ToolContextCache()
 context, _is_new = cache.get_or_create(thread_id)
-prepare_chat_context(context)  # idempotent namespace registration
+prepare_chat_context(context)  # idempotent; sets state_key="haiku.rag.chat"
 
 deps = ChatDeps(
     config=config,
     client=client,
     tool_context=context,
-    state_key=AGUI_STATE_KEY,  # "haiku.rag.chat"
 )
 ```
 

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

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 from pydantic_ai import Agent
@@ -8,9 +8,9 @@ from haiku.rag.agents.chat.context import (
 )
 from haiku.rag.agents.chat.prompts import build_chat_prompt
 from haiku.rag.agents.chat.state import AGUI_STATE_KEY
-from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
 from haiku.rag.tools.context import ToolContext, prepare_context
+from haiku.rag.tools.deps import AgentDeps
 from haiku.rag.tools.document import create_document_toolset
 from haiku.rag.tools.qa import (
     QA_SESSION_NAMESPACE,
@@ -34,41 +34,21 @@ def _on_qa_complete(qa_session_state: QASessionState, config: AppConfig) -> None
 
 
 @dataclass
-class ChatDeps:
+class ChatDeps(AgentDeps):
     """Dependencies for chat agent.
 
-    Implements RAGDeps protocol and StateHandler protocol for AG-UI state management.
+    Extends AgentDeps with chat-specific config and state handling.
     """
 
-    config: AppConfig
+    config: AppConfig = field(default_factory=AppConfig)
-    client: HaikuRAG
-    tool_context: ToolContext
-    state_key: str | None = None
 
-    @property
+    @AgentDeps.state.setter
-    def state(self) -> dict[str, Any]:
-        """Get current state for AG-UI protocol.
-
-        Combines all registered namespace states into a single flat dict,
-        matching the ChatSessionState schema expected by AG-UI clients.
-        """
-        snapshot = self.tool_context.build_state_snapshot()
-        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."""
+        """Set state from AG-UI protocol with chat-specific overrides."""
         if value is None:
             return
 
-        # Extract from namespaced key if present
+        state_data = self._extract_state_data(value)
-        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
 
         # Preserve server's session_context before restore overwrites it
         qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)

haiku_rag_slim/haiku/rag/tools/deps.py

@@ -14,14 +14,14 @@ class AgentDeps:
 
     client: HaikuRAG
     tool_context: ToolContext
-    state_key: str | None = None
 
     @property
     def state(self) -> dict[str, Any]:
         """Get current state for AG-UI protocol."""
         snapshot = self.tool_context.build_state_snapshot()
-        if self.state_key:
+        state_key = self.tool_context.state_key
-            return {self.state_key: snapshot}
+        if state_key:
+            return {state_key: snapshot}
         return snapshot
 
     @state.setter
@@ -29,10 +29,14 @@ class AgentDeps:
         """Set state from AG-UI protocol."""
         if value is None:
             return
-
+        data = self._extract_state_data(value)
-        data: dict[str, Any] = value
-        if self.state_key and self.state_key in value:
-            nested = value[self.state_key]
-            if isinstance(nested, dict):
-                data = nested
         self.tool_context.restore_state_snapshot(data)
+
+    def _extract_state_data(self, value: dict[str, Any]) -> dict[str, Any]:
+        """Extract flat state dict, unwrapping state_key if present."""
+        state_key = self.tool_context.state_key
+        if state_key and state_key in value:
+            nested = value[state_key]
+            if isinstance(nested, dict):
+                return nested
+        return value

tests/agents/chat/test_chat_agent.py

@@ -64,7 +64,17 @@ def test_chat_deps_initialization(temp_db_path):
     assert deps.config is Config
     assert deps.client is client
     assert deps.tool_context is context
-    assert deps.state_key is None
+    client.close()
+
+
+def test_chat_deps_is_agent_deps(temp_db_path):
+    """Test ChatDeps is a subclass of AgentDeps."""
+    from haiku.rag.tools.deps import AgentDeps
+
+    client = HaikuRAG(temp_db_path, create=True)
+    context = ToolContext()
+    deps = ChatDeps(config=Config, client=client, tool_context=context)
+    assert isinstance(deps, AgentDeps)
     client.close()
 
 
@@ -73,42 +83,18 @@ def test_agui_state_key_constant():
     assert AGUI_STATE_KEY == "haiku.rag.chat"
 
 
-def test_chat_deps_with_state_key(temp_db_path):
-    """Test ChatDeps can be initialized with state_key for keyed state emission."""
-    client = HaikuRAG(temp_db_path, create=True)
-    context = ToolContext()
-    deps = ChatDeps(
-        config=Config, client=client, tool_context=context, state_key="my_state"
-    )
-
-    assert deps.config is Config
-    assert deps.state_key == "my_state"
-    client.close()
-
-
-def test_chat_deps_state_key_default_none(temp_db_path):
-    """Test ChatDeps state_key defaults to None."""
-    client = HaikuRAG(temp_db_path, create=True)
-    context = ToolContext()
-    deps = ChatDeps(config=Config, client=client, tool_context=context)
-
-    assert deps.state_key is None
-    client.close()
-
-
 def test_chat_deps_state_setter_handles_initial_context(temp_db_path):
     """Test ChatDeps.state setter transfers initial_context to qa_session_state."""
     from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
 
     client = HaikuRAG(temp_db_path, create=True)
     context = ToolContext()
+    context.state_key = AGUI_STATE_KEY
     # Register QASessionState (normally done by prepare_chat_context)
     context.register(QA_SESSION_NAMESPACE, QASessionState())
     context.register(SESSION_NAMESPACE, SessionState())
 
-    deps = ChatDeps(
+    deps = ChatDeps(config=Config, client=client, tool_context=context)
-        config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY
-    )
 
     # Client sends initial_context with no session_context
     incoming_state = {
@@ -140,12 +126,11 @@ def test_chat_deps_state_setter_parses_session_context_dict(temp_db_path):
 
     client = HaikuRAG(temp_db_path, create=True)
     context = ToolContext()
+    context.state_key = AGUI_STATE_KEY
     context.register(QA_SESSION_NAMESPACE, QASessionState())
     context.register(SESSION_NAMESPACE, SessionState())
 
-    deps = ChatDeps(
+    deps = ChatDeps(config=Config, client=client, tool_context=context)
-        config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY
-    )
 
     # Client sends session_context as a dict (as it comes from JSON)
     incoming_state = {
@@ -178,6 +163,7 @@ def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path):
 
     client = HaikuRAG(temp_db_path, create=True)
     context = ToolContext()
+    context.state_key = AGUI_STATE_KEY
     qa_state = QASessionState()
     qa_state.session_context = SessionContext(
         summary="Fresh summary from background summarizer"
@@ -185,9 +171,7 @@ def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path):
     context.register(QA_SESSION_NAMESPACE, qa_state)
     context.register(SESSION_NAMESPACE, SessionState())
 
-    deps = ChatDeps(
+    deps = ChatDeps(config=Config, client=client, tool_context=context)
-        config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY
-    )
 
     # Client sends stale session_context
     incoming_state = {
@@ -500,7 +484,6 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
             config=Config,
             client=client,
             tool_context=context,
-            state_key=AGUI_STATE_KEY,
         )
 
         # Ask a question that should use the ask tool
@@ -557,7 +540,6 @@ async def test_chat_agent_ask_triggers_background_summarization(
             config=Config,
             client=client,
             tool_context=context,
-            state_key=AGUI_STATE_KEY,
         )
 
         # Patch internal trigger to avoid concurrent HTTP calls during VCR
@@ -623,7 +605,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_
             config=Config,
             client=client,
             tool_context=context,
-            state_key=AGUI_STATE_KEY,
         )
 
         # Set initial state with initial_context (mimicking AG-UI client)
@@ -725,7 +706,6 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
             config=Config,
             client=client,
             tool_context=context,
-            state_key=AGUI_STATE_KEY,
         )
 
         # First ask - establishes qa_history

tests/agents/chat/test_features.py

@@ -95,6 +95,8 @@ def test_no_qa_skips_qa_session_state(temp_db_path):
 
 def test_chat_deps_state_without_qa(temp_db_path):
     """ChatDeps.state getter omits qa_history/session_context when QASessionState absent."""
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
+
     client = HaikuRAG(temp_db_path, create=True)
     context = ToolContext()
     prepare_chat_context(context, features=[FEATURE_SEARCH])
@@ -102,13 +104,17 @@ def test_chat_deps_state_without_qa(temp_db_path):
     deps = ChatDeps(config=Config, client=client, tool_context=context)
     state = deps.state
 
+    # State is wrapped under the AGUI state key
+    assert AGUI_STATE_KEY in state
+    inner = state[AGUI_STATE_KEY]
+
     # SessionState fields should be present
-    assert "document_filter" in state
+    assert "document_filter" in inner
-    assert "citation_registry" in state
+    assert "citation_registry" in inner
-    assert "citations" in state
+    assert "citations" in inner
     # QA fields should NOT be present
-    assert "qa_history" not in state
+    assert "qa_history" not in inner
-    assert "session_context" not in state
+    assert "session_context" not in inner
     client.close()
 
 

tests/tools/test_deps.py

@@ -31,10 +31,11 @@ def test_agent_deps_state_getter_with_session(mock_client):
 
 
 def test_agent_deps_state_getter_with_state_key(mock_client):
-    """state wraps snapshot under state_key when set."""
+    """state wraps snapshot under state_key when set on context."""
     ctx = ToolContext()
+    ctx.state_key = "my_app"
     ctx.register(SESSION_NAMESPACE, SessionState())
-    deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="my_app")
+    deps = AgentDeps(client=mock_client, tool_context=ctx)
     snapshot = deps.state
     assert "my_app" in snapshot
     assert "citations" in snapshot["my_app"]
@@ -54,8 +55,9 @@ def test_agent_deps_state_setter_restores(mock_client):
 def test_agent_deps_state_setter_with_state_key(mock_client):
     """state setter extracts data from namespaced key."""
     ctx = ToolContext()
+    ctx.state_key = "my_app"
     ctx.register(SESSION_NAMESPACE, SessionState())
-    deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="my_app")
+    deps = AgentDeps(client=mock_client, tool_context=ctx)
     deps.state = {"my_app": {"document_filter": ["doc1"]}}
     session = ctx.get(SESSION_NAMESPACE, SessionState)
     assert session is not None
@@ -76,16 +78,18 @@ def test_agent_deps_state_setter_ignores_none(mock_client):
 def test_agent_deps_state_roundtrip(mock_client):
     """Build snapshot then restore produces equivalent state."""
     ctx = ToolContext()
+    ctx.state_key = "app"
     ctx.register(SESSION_NAMESPACE, SessionState(document_filter=["doc1"]))
     ctx.register(QA_SESSION_NAMESPACE, QASessionState())
-    deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="app")
+    deps = AgentDeps(client=mock_client, tool_context=ctx)
 
     snapshot = deps.state
 
     ctx2 = ToolContext()
+    ctx2.state_key = "app"
     ctx2.register(SESSION_NAMESPACE, SessionState())
     ctx2.register(QA_SESSION_NAMESPACE, QASessionState())
-    deps2 = AgentDeps(client=mock_client, tool_context=ctx2, state_key="app")
+    deps2 = AgentDeps(client=mock_client, tool_context=ctx2)
     deps2.state = snapshot
 
     session = ctx2.get(SESSION_NAMESPACE, SessionState)