From 09aa94bfe5678b51378386eaaf105dd89dbc3640 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 27 Jan 2026 10:52:23 +0200
Subject: [PATCH] Fix datetime JSON serialization in AG-UI StateSnapshotEvent

---
 CHANGELOG.md                                  |  5 ++++
 haiku_rag_slim/haiku/rag/agents/chat/agent.py |  4 +--
 tests/agents/chat/test_state.py               | 26 +++++++++++++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7fcd2e2c..958b91f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,11 @@
   - Eliminates need for clients to import and call internal cache functions (`cache_session_context`, `get_cached_session_context`)
   - `session_id` now auto-generates a UUID if not provided (previously defaulted to empty string)
 
+### Fixed
+
+- **AG-UI StateSnapshotEvent JSON Serialization**: Chat agent tools now use `model_dump(mode="json")` when creating `StateSnapshotEvent`
+  - Fixes `TypeError: Object of type datetime is not JSON serializable` when external clients persist AG-UI state to database JSON columns
+
 ## [0.27.0] - 2026-01-26
 
 ### Added
diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py
index 8e3b1c52..c3cf1ea5 100644
--- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py
+++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py
@@ -170,7 +170,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             line += f"\n    {snippet}"
             result_lines.append(line)
 
-        snapshot = new_state.model_dump()
+        snapshot = new_state.model_dump(mode="json")
         if ctx.deps.state_key:
             snapshot = {ctx.deps.state_key: snapshot}
 
@@ -355,7 +355,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             citation_refs = " ".join(f"[{c.index}]" for c in citation_infos)
             answer_text = f"{answer_text}\n\nSources: {citation_refs}"
 
-        snapshot = new_state.model_dump()
+        snapshot = new_state.model_dump(mode="json")
         if ctx.deps.state_key:
             snapshot = {ctx.deps.state_key: snapshot}
 
diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py
index 12d22d40..2ed78e7a 100644
--- a/tests/agents/chat/test_state.py
+++ b/tests/agents/chat/test_state.py
@@ -4,6 +4,7 @@ from haiku.rag.agents.chat.state import (
     MAX_QA_HISTORY,
     ChatSessionState,
     QAResponse,
+    SessionContext,
     build_document_filter,
     build_multi_document_filter,
     combine_filters,
@@ -615,3 +616,28 @@ def test_chat_session_state_initial_context_serialization():
 
     restored = ChatSessionState.model_validate(state_dict)
     assert restored.initial_context == "User is working on authentication"
+
+
+def test_chat_session_state_model_dump_json_serializes_datetime():
+    """model_dump(mode='json') should serialize datetime to ISO string.
+
+    Agent tools use model_dump(mode='json') when creating StateSnapshotEvent
+    to ensure datetime fields are JSON-serializable for external clients
+    persisting AG-UI state to database JSON columns.
+    """
+    from datetime import datetime
+
+    session_state = ChatSessionState(
+        session_id="test",
+        session_context=SessionContext(
+            summary="Test summary",
+            last_updated=datetime(2025, 1, 27, 12, 0, 0),
+        ),
+    )
+
+    # This is how agent.py creates snapshots for StateSnapshotEvent
+    snapshot = session_state.model_dump(mode="json")
+
+    # datetime should be serialized as ISO string, not datetime object
+    assert isinstance(snapshot["session_context"]["last_updated"], str)
+    assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00"