Move session context to agent, make AGUI work

app/backend/main.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import uuid
 from pathlib import Path
 
 from dotenv import find_dotenv, load_dotenv
@@ -17,7 +18,6 @@ from haiku.rag.agents.chat import (
     ChatDeps,
     ChatSessionState,
     QAResponse,
-    SessionContext,
     create_chat_agent,
 )
 from haiku.rag.client import HaikuRAG
@@ -75,7 +75,6 @@ def get_client(effective_db_path: Path) -> HaikuRAG:
 async def stream_chat(request: Request) -> Response:
     """Chat streaming endpoint with AG-UI protocol."""
     body = await request.body()
-    logger.info("Received chat request")
 
     # Parse request to build run_input
     accept = request.headers.get("accept", SSE_CONTENT_TYPE)
@@ -84,7 +83,7 @@ async def stream_chat(request: Request) -> Response:
     # Restore session state from incoming AG-UI state (look under namespaced key)
     initial_qa_history: list[QAResponse] = []
     background_context: str | None = None
-    session_context: SessionContext | None = None
+    session_id: str | None = None
     state = getattr(run_input, "state", None)
     if state:
         chat_state = state.get(AGUI_STATE_KEY, state)
@@ -93,20 +92,21 @@ async def stream_chat(request: Request) -> Response:
                 QAResponse(**qa) for qa in chat_state.get("qa_history", [])
             ]
         background_context = chat_state.get("background_context")
-        ctx_data = chat_state.get("session_context")
+        session_id = chat_state.get("session_id")
-        if ctx_data and isinstance(ctx_data, dict):
-            session_context = SessionContext(**ctx_data)
 
-    # Build deps with session state
+    # 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=thread_id or "",
+            session_id=session_id,
             qa_history=initial_qa_history,
             background_context=background_context,
-            session_context=session_context,
+            # session_context intentionally NOT set - agent will fetch from cache
         ),
         state_key=AGUI_STATE_KEY,
     )

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

@@ -1,5 +1,8 @@
 from haiku.rag.agents.chat.agent import create_chat_agent
-from haiku.rag.agents.chat.context import summarize_session, update_session_context
+from haiku.rag.agents.chat.context import (
+    summarize_session,
+    update_session_context,
+)
 from haiku.rag.agents.chat.search import SearchAgent
 from haiku.rag.agents.chat.state import (
     AGUI_STATE_KEY,

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

@@ -1,10 +1,12 @@
 import asyncio
-import logging
 
 from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic_ai import Agent, RunContext, ToolReturn
 
-from haiku.rag.agents.chat.context import update_session_context
+from haiku.rag.agents.chat.context import (
+    get_cached_session_context,
+    update_session_context,
+)
 from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
 from haiku.rag.agents.chat.search import SearchAgent
 from haiku.rag.agents.chat.state import (
@@ -21,8 +23,6 @@ from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.config.models import AppConfig
 from haiku.rag.utils import get_model
 
-logger = logging.getLogger(__name__)
-
 # Track summarization tasks per session to allow cancellation
 _summarization_tasks: dict[str, asyncio.Task[None]] = {}
 
@@ -39,11 +39,10 @@ async def _update_context_background(
             config=config,
             session_state=session_state,
         )
-        logger.debug("Session context updated")
     except asyncio.CancelledError:
-        logger.debug("Session context summarization cancelled")
+        pass
     except Exception:
-        logger.exception("Failed to update session context")
+        pass
 
 
 def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
@@ -111,10 +110,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         ]
 
         # Build new state with citations
+        session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
         new_state = ChatSessionState(
-            session_id=(
+            session_id=session_id,
-                ctx.deps.session_state.session_id if ctx.deps.session_state else ""
-            ),
             citations=citation_infos,
             qa_history=(
                 ctx.deps.session_state.qa_history if ctx.deps.session_state else []
@@ -124,11 +122,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
                 if ctx.deps.session_state
                 else None
             ),
-            session_context=(
+            session_context=get_cached_session_context(session_id)
-                ctx.deps.session_state.session_context
+            if session_id
-                if ctx.deps.session_state
+            else None,
-                else None
-            ),
         )
 
         # Return detailed results for the agent to present
@@ -183,19 +179,20 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         graph = build_conversational_graph(config=ctx.deps.config)
 
         # Determine context strategy:
-        # 1. If session_context exists, use compressed summary (skip raw qa_history)
+        # 1. Read from server cache (ignoring client state)
-        # 2. Otherwise, fall back to explicit background_context
+        # 2. Fall back to explicit background_context (first request)
         background_context: str | None = None
+        session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
         if ctx.deps.session_state:
-            if (
+            cached_context = (
-                ctx.deps.session_state.session_context
+                get_cached_session_context(session_id) if session_id else None
-                and ctx.deps.session_state.session_context.summary
+            )
-            ):
+
-                # Use compressed SessionContext - no need for raw qa_history
+            if cached_context and cached_context.summary:
-                background_context = (
+                # Use cached SessionContext from previous summarization
-                    ctx.deps.session_state.session_context.render_markdown()
+                background_context = cached_context.render_markdown()
-                )
             elif ctx.deps.session_state.background_context:
+                # Fall back to explicit background_context (first request)
                 background_context = ctx.deps.session_state.background_context
 
         context = ResearchContext(
@@ -247,7 +244,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
 
             # Spawn background task to update session context
             # Cancel any previous summarization for this session
-            session_id = ctx.deps.session_state.session_id
             if session_id in _summarization_tasks:
                 _summarization_tasks[session_id].cancel()
 
@@ -263,9 +259,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
 
         # Build new state with citations AND accumulated qa_history
         new_state = ChatSessionState(
-            session_id=(
+            session_id=session_id,
-                ctx.deps.session_state.session_id if ctx.deps.session_state else ""
-            ),
             citations=citation_infos,
             qa_history=(
                 ctx.deps.session_state.qa_history if ctx.deps.session_state else []
@@ -275,11 +269,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
                 if ctx.deps.session_state
                 else None
             ),
-            session_context=(
+            session_context=get_cached_session_context(session_id)
-                ctx.deps.session_state.session_context
+            if session_id
-                if ctx.deps.session_state
+            else None,
-                else None
-            ),
         )
 
         # Format answer with citation references and confidence

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

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 
 from pydantic_ai import Agent
 
@@ -7,6 +7,41 @@ from haiku.rag.agents.chat.state import ChatSessionState, QAResponse, SessionCon
 from haiku.rag.config.models import AppConfig
 from haiku.rag.utils import get_model
 
+# Cache for session contexts (session_id -> SessionContext)
+# Used to persist async summarization results between requests
+_session_context_cache: dict[str, SessionContext] = {}
+_cache_timestamps: dict[str, datetime] = {}
+_CACHE_TTL = timedelta(hours=1)
+
+
+def _cleanup_stale_cache() -> None:
+    """Remove cache entries older than TTL."""
+    now = datetime.now()
+    stale = [sid for sid, ts in _cache_timestamps.items() if now - ts > _CACHE_TTL]
+    for sid in stale:
+        _session_context_cache.pop(sid, None)
+        _cache_timestamps.pop(sid, None)
+
+
+def cache_session_context(session_id: str, context: SessionContext) -> None:
+    """Store session context in cache."""
+    _cleanup_stale_cache()
+    _session_context_cache[session_id] = context
+    _cache_timestamps[session_id] = datetime.now()
+
+
+def get_cached_session_context(session_id: str) -> SessionContext | None:
+    """Get session context from server cache.
+
+    Args:
+        session_id: The session identifier.
+
+    Returns:
+        Cached SessionContext, or None if not cached.
+    """
+    _cleanup_stale_cache()
+    return _session_context_cache.get(session_id)
+
 
 async def summarize_session(
     qa_history: list[QAResponse],
@@ -70,6 +105,9 @@ async def update_session_context(
         summary=summary,
         last_updated=datetime.now(),
     )
+    # Also cache for next-run delivery in stateless contexts
+    if session_state.session_id:
+        cache_session_context(session_state.session_id, session_state.session_context)
 
 
 def _format_qa_history(qa_history: list[QAResponse]) -> str:

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

@@ -101,16 +101,10 @@ class ChatDeps:
                 self.session_state.background_context = state_data.get(
                     "background_context"
                 )
-            if "session_id" in state_data:
+            if state_data.get("session_id"):
-                self.session_state.session_id = state_data.get("session_id", "")
+                self.session_state.session_id = state_data["session_id"]
-            if "session_context" in state_data:
+            # NOTE: session_context intentionally NOT updated from client
-                ctx_data = state_data.get("session_context")
+            # The agent owns this via server-side cache
-                if ctx_data is None:
-                    self.session_state.session_context = None
-                elif isinstance(ctx_data, dict):
-                    self.session_state.session_context = SessionContext(**ctx_data)
-                else:
-                    self.session_state.session_context = ctx_data
 
 
 @dataclass

haiku_rag_slim/haiku/rag/chat/app.py

@@ -21,6 +21,7 @@ from haiku.rag.agents.chat.state import (
     AGUI_STATE_KEY,
     ChatDeps,
     ChatSessionState,
+    SessionContext,
 )
 from haiku.rag.agents.research.models import Citation
 from haiku.rag.client import HaikuRAG
@@ -128,9 +129,15 @@ class ChatApp(App):
 
         # Create agent and session state
         self.agent = create_chat_agent(self.config)
+        initial_context = (
+            SessionContext(summary=self.background_context, last_updated=datetime.now())
+            if self.background_context
+            else None
+        )
         self.session_state = ChatSessionState(
             session_id=str(uuid.uuid4()),
             background_context=self.background_context,
+            session_context=initial_context,
         )
 
         # Focus the input field
@@ -277,9 +284,15 @@ class ChatApp(App):
         self._selected_citation_idx = None
         self._message_history.clear()
         # Reset session state for fresh conversation (preserve background_context)
+        initial_context = (
+            SessionContext(summary=self.background_context, last_updated=datetime.now())
+            if self.background_context
+            else None
+        )
         self.session_state = ChatSessionState(
             session_id=str(uuid.uuid4()),
             background_context=self.background_context,
+            session_context=initial_context,
         )
 
     def action_focus_input(self) -> None:

tests/agents/chat/test_context.py

@@ -254,3 +254,137 @@ class TestUpdateSessionContext:
         # session_context should exist but have empty summary
         assert session_state.session_context is not None
         assert session_state.session_context.summary == ""
+
+
+class TestSessionContextCache:
+    """Tests for server-side session context caching."""
+
+    def test_cache_and_retrieve_session_context(self):
+        """Test caching and retrieving a session context."""
+        from haiku.rag.agents.chat.context import (
+            _session_context_cache,
+            cache_session_context,
+            get_cached_session_context,
+        )
+
+        # Clear cache
+        _session_context_cache.clear()
+
+        now = datetime.now()
+        ctx = SessionContext(summary="Test summary", last_updated=now)
+
+        cache_session_context("session-1", ctx)
+        result = get_cached_session_context("session-1")
+
+        assert result is not None
+        assert result.summary == "Test summary"
+        assert result.last_updated == now
+
+    def test_get_cached_session_context_returns_none_when_not_cached(self):
+        """Test get_cached_session_context returns None when nothing cached."""
+        from haiku.rag.agents.chat.context import (
+            _session_context_cache,
+            get_cached_session_context,
+        )
+
+        _session_context_cache.clear()
+
+        result = get_cached_session_context("nonexistent-session")
+        assert result is None
+
+    def test_cache_ttl_cleanup_removes_stale_entries(self):
+        """Test that stale cache entries are cleaned up."""
+        from datetime import timedelta
+
+        from haiku.rag.agents.chat.context import (
+            _CACHE_TTL,
+            _cache_timestamps,
+            _session_context_cache,
+            cache_session_context,
+            get_cached_session_context,
+        )
+
+        _session_context_cache.clear()
+        _cache_timestamps.clear()
+
+        # Add an entry
+        ctx = SessionContext(summary="Old summary", last_updated=datetime.now())
+        cache_session_context("stale-session", ctx)
+
+        # Make the entry stale by backdating its timestamp
+        _cache_timestamps["stale-session"] = (
+            datetime.now() - _CACHE_TTL - timedelta(seconds=1)
+        )
+
+        # Getting session context should trigger cleanup
+        result = get_cached_session_context("stale-session")
+
+        # Should be None because the entry was cleaned up
+        assert result is None
+        assert "stale-session" not in _session_context_cache
+
+    @pytest.mark.asyncio
+    async def test_update_session_context_caches_result(self):
+        """Test update_session_context stores result in cache."""
+        from unittest.mock import AsyncMock, patch
+
+        from haiku.rag.agents.chat.context import (
+            _session_context_cache,
+            get_cached_session_context,
+            update_session_context,
+        )
+        from haiku.rag.agents.chat.state import ChatSessionState
+
+        _session_context_cache.clear()
+
+        session_state = ChatSessionState(session_id="cache-test-session")
+
+        qa_history = [
+            QAResponse(
+                question="What is Python?",
+                answer="A programming language.",
+                confidence=0.95,
+            )
+        ]
+
+        # Mock summarize_session to avoid LLM call (we're testing caching, not summarization)
+        with patch(
+            "haiku.rag.agents.chat.context.summarize_session",
+            new=AsyncMock(return_value="Mocked summary"),
+        ):
+            await update_session_context(
+                qa_history=qa_history,
+                config=Config,
+                session_state=session_state,
+            )
+
+        # Should be cached
+        cached = get_cached_session_context("cache-test-session")
+        assert cached is not None
+        assert cached.summary == "Mocked summary"
+        assert cached.summary == session_state.session_context.summary
+
+    @pytest.mark.asyncio
+    async def test_update_session_context_no_cache_without_session_id(self):
+        """Test update_session_context doesn't cache without session_id."""
+        from haiku.rag.agents.chat.context import (
+            _session_context_cache,
+            get_cached_session_context,
+            update_session_context,
+        )
+        from haiku.rag.agents.chat.state import ChatSessionState
+
+        _session_context_cache.clear()
+
+        # No session_id
+        session_state = ChatSessionState()
+
+        await update_session_context(
+            qa_history=[],
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Nothing should be cached (empty session_id)
+        cached = get_cached_session_context("")
+        assert cached is None

tests/agents/chat/test_state.py

@@ -297,45 +297,12 @@ def test_chat_deps_state_getter_includes_session_context():
     )
 
 
-def test_chat_deps_state_setter_restores_session_context():
+def test_chat_deps_state_setter_ignores_session_context():
-    """Test ChatDeps.state setter restores session_context from incoming state."""
+    """Test ChatDeps.state setter ignores session_context from client.
-    from unittest.mock import MagicMock
 
-    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
+    The agent owns session_context via server-side cache, so client-provided
-
+    session_context should be ignored to prevent stale state overwriting.
-    mock_client = MagicMock()
+    """
-    mock_config = MagicMock()
-
-    session_state = ChatSessionState(session_id="initial")
-    deps = ChatDeps(
-        client=mock_client,
-        config=mock_config,
-        session_state=session_state,
-        state_key=AGUI_STATE_KEY,
-    )
-
-    incoming_state = {
-        AGUI_STATE_KEY: {
-            "session_id": "test-123",
-            "qa_history": [],
-            "citations": [],
-            "background_context": None,
-            "session_context": {
-                "summary": "Restored context summary.",
-                "last_updated": "2025-01-15T10:30:00",
-            },
-        }
-    }
-
-    deps.state = incoming_state
-
-    assert deps.session_state is not None
-    assert deps.session_state.session_context is not None
-    assert deps.session_state.session_context.summary == "Restored context summary."
-
-
-def test_chat_deps_state_setter_handles_null_session_context():
-    """Test ChatDeps.state setter handles null session_context."""
     from unittest.mock import MagicMock
 
     from haiku.rag.agents.chat.state import (
@@ -348,10 +315,10 @@ def test_chat_deps_state_setter_handles_null_session_context():
     mock_client = MagicMock()
     mock_config = MagicMock()
 
-    # Start with a session_context
+    # Start with a session_context (e.g., from cache)
     session_state = ChatSessionState(
         session_id="test",
-        session_context=SessionContext(summary="Initial summary"),
+        session_context=SessionContext(summary="Server-side context"),
     )
     deps = ChatDeps(
         client=mock_client,
@@ -360,17 +327,22 @@ def test_chat_deps_state_setter_handles_null_session_context():
         state_key=AGUI_STATE_KEY,
     )
 
-    # Send null to clear it
+    # Client sends different session_context (stale)
     incoming_state = {
         AGUI_STATE_KEY: {
             "session_id": "test",
             "qa_history": [],
             "citations": [],
-            "session_context": None,
+            "session_context": {
+                "summary": "Client-provided stale context",
+                "last_updated": "2025-01-15T10:30:00",
+            },
         }
     }
 
     deps.state = incoming_state
 
+    # session_context should NOT be overwritten
     assert deps.session_state is not None
-    assert deps.session_state.session_context is None
+    assert deps.session_state.session_context is not None
+    assert deps.session_state.session_context.summary == "Server-side context"