From 11f7f022624ffa07ed602e486cda86983ebeac14 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 14 Jan 2026 14:21:08 +0200
Subject: [PATCH] Add option to namespace state key as a "feature" for ag-ui
 apps that maintain their own state

---
 CHANGELOG.md                                  |  3 +++
 haiku_rag_slim/haiku/rag/agents/chat/agent.py | 12 ++++++++--
 haiku_rag_slim/haiku/rag/agents/chat/state.py |  1 +
 tests/agents/chat/test_chat_agent.py          | 22 +++++++++++++++++++
 4 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 052b764b..9471aad7 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py
index ffae5408..636eb08e 100644
--- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py
+++ b/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,
                 )
             ],
         )
diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py
index 4a0425a5..44d2c4e8 100644
--- a/haiku_rag_slim/haiku/rag/agents/chat/state.py
+++ b/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
diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py
index b022aac2..e7f58484 100644
--- a/tests/agents/chat/test_chat_agent.py
+++ b/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")