Merge pull request #261 from ggozad/chore/chat-agent-improvements

Initial Context for Chat Sessions & fixes

CHANGELOG.md

@@ -1,6 +1,19 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- **Initial Context for Chat Sessions**: New `initial_context` field on `ChatSessionState` allows external clients to seed sessions with background context
+  - Static context set once at session creation, used as fallback when no cached session context exists
+  - Incorporated into first summarization, after which evolved `session_context` takes precedence
+  - 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

app/backend/main.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import uuid
 from pathlib import Path
 
 from dotenv import find_dotenv, load_dotenv
@@ -94,18 +93,13 @@ async def stream_chat(request: Request) -> Response:
         session_id = chat_state.get("session_id")
         document_filter = chat_state.get("document_filter", [])
 
-    # Determine session_id: prefer state, fall back to thread_id, generate UUID if neither
-    thread_id = getattr(run_input, "thread_id", None)
-    if not session_id:
-        session_id = thread_id or str(uuid.uuid4())
-
     deps = ChatDeps(
         client=get_client(db_path),
         config=Config,
         session_state=ChatSessionState(
-            session_id=session_id,
             qa_history=initial_qa_history,
             document_filter=document_filter,
+            **({"session_id": session_id} if session_id else {}),
         ),
         state_key=AGUI_STATE_KEY,
     )

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}
 
@@ -217,12 +217,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         graph = build_conversational_graph(config=ctx.deps.config)
         session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
 
-        # Get session context from server cache for planning
+        # Get session context from server cache for planning, fallback to initial_context
         cached_context = get_cached_session_context(session_id) if session_id else None
         session_context = (
             cached_context.render_markdown()
             if cached_context and cached_context.summary
-            else None
+            else (
+                ctx.deps.session_state.initial_context
+                if ctx.deps.session_state
+                else None
+            )
         )
 
         # Find relevant prior answers from qa_history
@@ -351,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}
 

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

@@ -82,10 +82,12 @@ async def update_session_context(
         config: AppConfig for model selection.
         session_state: The session state to update.
     """
-    # Use existing session_context summary if available
+    # Use existing session_context summary if available, else initial_context
     current_context: str | None = None
     if session_state.session_context and session_state.session_context.summary:
         current_context = session_state.session_context.summary
+    elif session_state.initial_context:
+        current_context = session_state.initial_context
 
     summary = await summarize_session(
         qa_history, config, current_context=current_context

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

@@ -1,3 +1,4 @@
+import uuid
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any
@@ -55,7 +56,8 @@ class SessionContext(BaseModel):
 class ChatSessionState(BaseModel):
     """State shared between frontend and agent via AG-UI."""
 
-    session_id: str = ""
+    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    initial_context: str | None = None
     citations: list[Citation] = []
     qa_history: list[QAResponse] = []
     session_context: SessionContext | None = None

haiku_rag_slim/haiku/rag/agents/research/graph.py

@@ -87,7 +87,7 @@ async def _plan_step_logic(
         else plan_prompt
     )
 
-    plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent(
+    plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent(  # type: ignore[assignment]
         model=get_model(model_config, config),
         output_type=ResearchPlan,
         instructions=effective_plan_prompt,
@@ -154,7 +154,7 @@ async def _search_one_step_logic(
         deps.semaphore = asyncio.Semaphore(state.max_concurrency)
 
     async with deps.semaphore:
-        agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(
+        agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
             output_type=ToolOutput(RawSearchAnswer, max_retries=3),
             instructions=search_prompt,
@@ -280,7 +280,7 @@ def build_research_graph(
         state = ctx.state
         deps = ctx.deps
 
-        agent: Agent[ResearchDependencies, EvaluationResult] = Agent(
+        agent: Agent[ResearchDependencies, EvaluationResult] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
             output_type=EvaluationResult,
             instructions=decision_prompt,
@@ -341,7 +341,7 @@ def build_research_graph(
         state = ctx.state
         deps = ctx.deps
 
-        agent: Agent[ResearchDependencies, ResearchReport] = Agent(
+        agent: Agent[ResearchDependencies, ResearchReport] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
             output_type=ResearchReport,
             instructions=synthesis_prompt,
@@ -479,7 +479,7 @@ def build_conversational_graph(
         state = ctx.state
         deps = ctx.deps
 
-        agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent(
+        agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent(  # type: ignore[assignment]
             model=get_model(config.research.model, config),
             output_type=ConversationalAnswer,
             instructions=conversational_prompt,

haiku_rag_slim/haiku/rag/chat/app.py

@@ -153,7 +153,6 @@ class ChatApp(App):
         # Create agent and session state
         self.agent = create_chat_agent(self.config)
         self.session_state = ChatSessionState(
-            session_id=str(uuid.uuid4()),
             document_filter=self._document_filter,
         )
 
@@ -300,7 +299,6 @@ class ChatApp(App):
         self._message_history.clear()
         # Reset session state for fresh conversation (preserve document filter)
         self.session_state = ChatSessionState(
-            session_id=str(uuid.uuid4()),
             document_filter=self._document_filter,
         )
 

tests/agents/chat/test_context.py

@@ -388,3 +388,103 @@ class TestSessionContextCache:
         # Nothing should be cached (empty session_id)
         cached = get_cached_session_context("")
         assert cached is None
+
+    @pytest.mark.asyncio
+    async def test_update_session_context_uses_initial_context_as_fallback(self):
+        """Test update_session_context uses initial_context when no session_context exists."""
+        from unittest.mock import patch
+
+        from haiku.rag.agents.chat.context import (
+            _session_context_cache,
+            update_session_context,
+        )
+        from haiku.rag.agents.chat.state import ChatSessionState
+
+        _session_context_cache.clear()
+
+        # Create session_state with initial_context but no session_context
+        session_state = ChatSessionState(
+            session_id="initial-context-test",
+            initial_context="User is working on a Python web application with FastAPI.",
+        )
+
+        qa_history = [
+            QAResponse(
+                question="What is JWT?",
+                answer="JSON Web Token for authentication.",
+                confidence=0.95,
+            )
+        ]
+
+        # Mock summarize_session to capture what gets passed as current_context
+        captured_current_context = []
+
+        async def mock_summarize(qa_history, config, current_context=None):
+            captured_current_context.append(current_context)
+            return "Mocked summary"
+
+        with patch(
+            "haiku.rag.agents.chat.context.summarize_session",
+            new=mock_summarize,
+        ):
+            await update_session_context(
+                qa_history=qa_history,
+                config=Config,
+                session_state=session_state,
+            )
+
+        # initial_context should have been passed as current_context
+        assert len(captured_current_context) == 1
+        assert (
+            captured_current_context[0]
+            == "User is working on a Python web application with FastAPI."
+        )
+
+    @pytest.mark.asyncio
+    async def test_update_session_context_session_context_takes_precedence(self):
+        """Test session_context.summary takes precedence over initial_context."""
+        from unittest.mock import patch
+
+        from haiku.rag.agents.chat.context import (
+            _session_context_cache,
+            update_session_context,
+        )
+        from haiku.rag.agents.chat.state import ChatSessionState, SessionContext
+
+        _session_context_cache.clear()
+
+        # Create session_state with BOTH initial_context and session_context
+        session_state = ChatSessionState(
+            session_id="precedence-test",
+            initial_context="Initial background info",
+            session_context=SessionContext(summary="Evolved session summary"),
+        )
+
+        qa_history = [
+            QAResponse(
+                question="What is JWT?",
+                answer="JSON Web Token.",
+                confidence=0.95,
+            )
+        ]
+
+        # Mock summarize_session to capture what gets passed as current_context
+        captured_current_context = []
+
+        async def mock_summarize(qa_history, config, current_context=None):
+            captured_current_context.append(current_context)
+            return "Mocked summary"
+
+        with patch(
+            "haiku.rag.agents.chat.context.summarize_session",
+            new=mock_summarize,
+        ):
+            await update_session_context(
+                qa_history=qa_history,
+                config=Config,
+                session_state=session_state,
+            )
+
+        # session_context.summary should take precedence over initial_context
+        assert len(captured_current_context) == 1
+        assert captured_current_context[0] == "Evolved session summary"

tests/agents/chat/test_state.py

@@ -1,6 +1,10 @@
+import uuid
+
 from haiku.rag.agents.chat.state import (
     MAX_QA_HISTORY,
+    ChatSessionState,
     QAResponse,
+    SessionContext,
     build_document_filter,
     build_multi_document_filter,
     combine_filters,
@@ -565,3 +569,75 @@ def test_chat_deps_state_getter_includes_document_filter():
     assert state is not None
     assert AGUI_STATE_KEY in state
     assert state[AGUI_STATE_KEY]["document_filter"] == ["doc1.pdf", "doc2.pdf"]
+
+
+def test_chat_session_state_auto_generates_session_id():
+    """New ChatSessionState should have a valid UUID session_id."""
+    state = ChatSessionState()
+    assert state.session_id
+    assert len(state.session_id) == 36  # UUID format
+    # Verify it's a valid UUID
+    uuid.UUID(state.session_id)
+
+
+def test_chat_session_state_preserves_explicit_session_id():
+    """Explicit session_id should be preserved."""
+    state = ChatSessionState(session_id="my-custom-id")
+    assert state.session_id == "my-custom-id"
+
+
+def test_chat_session_state_each_instance_gets_unique_id():
+    """Each new instance should get a unique session_id."""
+    state1 = ChatSessionState()
+    state2 = ChatSessionState()
+    assert state1.session_id != state2.session_id
+
+
+def test_chat_session_state_initial_context_default_none():
+    """Initial context should default to None."""
+    state = ChatSessionState()
+    assert state.initial_context is None
+
+
+def test_chat_session_state_initial_context_preserved():
+    """Explicit initial_context should be preserved."""
+    state = ChatSessionState(initial_context="Background info about the project")
+    assert state.initial_context == "Background info about the project"
+
+
+def test_chat_session_state_initial_context_serialization():
+    """initial_context should serialize and deserialize correctly."""
+    state = ChatSessionState(
+        session_id="test-123",
+        initial_context="User is working on authentication",
+    )
+    state_dict = state.model_dump()
+    assert state_dict["initial_context"] == "User is working on authentication"
+
+    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"