Add option to namespace state key as a "feature" for ag-ui apps that maintain their own state

CHANGELOG.md

@@ -3,6 +3,9 @@
 
 ### Added
 
+- **Keyed State Emission for Chat Agent**: New `state_key` parameter in `ChatDeps` for namespaced AG-UI state snapshots
+  - When set, tools emit `{state_key: snapshot}` instead of bare state, enabling state merging when multiple agents share state
+  - Default `None` preserves backwards compatibility (bare state emission)
 - **Page Image Generation Control**: New `generate_page_images` option in `ConversionOptions` to control PDF page image extraction
   - `generate_page_images: bool = True` - Enable/disable rendered page images (used by `visualize_chunk()`)
   - Works with both `docling-local` and `docling-serve` converters

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

@@ -102,13 +102,17 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             line += f"\n    {snippet}"
             result_lines.append(line)
 
+        snapshot = new_state.model_dump()
+        if ctx.deps.state_key:
+            snapshot = {ctx.deps.state_key: snapshot}
+
         return ToolReturn(
             return_value=f"Found {len(results)} results:\n\n"
             + "\n\n".join(result_lines),
             metadata=[
                 StateSnapshotEvent(
                     type=EventType.STATE_SNAPSHOT,
-                    snapshot=new_state.model_dump(),
+                    snapshot=snapshot,
                 )
             ],
         )
@@ -241,12 +245,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citation_infos)))
             answer_text = f"{answer_text}\n\nSources: {citation_refs}"
 
+        snapshot = new_state.model_dump()
+        if ctx.deps.state_key:
+            snapshot = {ctx.deps.state_key: snapshot}
+
         return ToolReturn(
             return_value=answer_text,
             metadata=[
                 StateSnapshotEvent(
                     type=EventType.STATE_SNAPSHOT,
-                    snapshot=new_state.model_dump(),
+                    snapshot=snapshot,
                 )
             ],
         )

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

@@ -159,6 +159,7 @@ class ChatDeps:
     config: AppConfig
     search_results: list[SearchResult] | None = None
     session_state: ChatSessionState | None = None
+    state_key: str | None = None
 
 
 @dataclass

tests/agents/chat/test_chat_agent.py

@@ -40,6 +40,28 @@ def test_chat_deps_initialization(temp_db_path):
     client.close()
 
 
+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)
+    deps = ChatDeps(client=client, config=Config, state_key="my_state")
+
+    assert deps.client is client
+    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)
+    deps = ChatDeps(client=client, config=Config)
+
+    assert deps.state_key is None
+
+    client.close()
+
+
 def test_chat_session_state():
     """Test ChatSessionState model."""
     state = ChatSessionState(session_id="test-session")