Introduce ToolContextCache

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

@@ -180,9 +180,9 @@ def create_chat_agent(
 
     existing = context.get(SESSION_NAMESPACE, SessionState)
     if existing is None:
-        context.register(SESSION_NAMESPACE, SessionState(state_key=AGUI_STATE_KEY))
-    elif existing.state_key is None:
-        existing.state_key = AGUI_STATE_KEY
+        context.register(SESSION_NAMESPACE, SessionState())
+    if context.state_key is None:
+        context.state_key = AGUI_STATE_KEY
 
     if FEATURE_QA in features:
         if context.get(QA_SESSION_NAMESPACE, QASessionState) is None:

haiku_rag_slim/haiku/rag/tools/__init__.py

@@ -1,5 +1,5 @@
 from haiku.rag.tools.analysis import create_analysis_toolset
-from haiku.rag.tools.context import ToolContext
+from haiku.rag.tools.context import ToolContext, ToolContextCache
 from haiku.rag.tools.document import (
     DocumentInfo,
     DocumentListResponse,
@@ -30,6 +30,7 @@ from haiku.rag.tools.session import (
 
 __all__ = [
     "ToolContext",
+    "ToolContextCache",
     "QAResult",
     "AnalysisResult",
     "build_document_filter",

haiku_rag_slim/haiku/rag/tools/context.py

@@ -1,3 +1,4 @@
+from datetime import datetime, timedelta
 from typing import Any, TypeVar, overload
 
 from pydantic import BaseModel, PrivateAttr
@@ -44,6 +45,7 @@ class ToolContext(BaseModel):
         ns_data = context.dump_namespaces()
     """
 
+    state_key: str | None = None
     _namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
 
     def register(self, namespace: str, state: BaseModel) -> None:
@@ -126,3 +128,48 @@ class ToolContext(BaseModel):
         state = state_type.model_validate(data)
         self._namespaces[namespace] = state
         return state
+
+
+class ToolContextCache:
+    """In-memory cache for ToolContext instances, keyed by external session/thread ID."""
+
+    def __init__(self, ttl: timedelta = timedelta(hours=1)) -> None:
+        self._cache: dict[str, ToolContext] = {}
+        self._timestamps: dict[str, datetime] = {}
+        self._ttl = ttl
+
+    def get_or_create(self, key: str) -> tuple[ToolContext, bool]:
+        """Get an existing context or create a new one.
+
+        Returns:
+            Tuple of (context, is_new) where is_new is True if a new context was created.
+        """
+        self._cleanup()
+        if key in self._cache:
+            self._timestamps[key] = datetime.now()
+            return self._cache[key], False
+
+        context = ToolContext()
+        self._cache[key] = context
+        self._timestamps[key] = datetime.now()
+        return context, True
+
+    def remove(self, key: str) -> None:
+        """Remove a specific key from the cache."""
+        self._cache.pop(key, None)
+        self._timestamps.pop(key, None)
+
+    def clear(self) -> None:
+        """Clear all entries."""
+        self._cache.clear()
+        self._timestamps.clear()
+
+    def _cleanup(self) -> None:
+        """Remove entries older than TTL."""
+        now = datetime.now()
+        expired = [
+            key for key, ts in self._timestamps.items() if (now - ts) >= self._ttl
+        ]
+        for key in expired:
+            self._cache.pop(key, None)
+            self._timestamps.pop(key, None)

haiku_rag_slim/haiku/rag/tools/qa.py

@@ -281,10 +281,12 @@ def create_qa_toolset(
         session_state: SessionState | None = None
         qa_session_state: QASessionState | None = None
         old_state_snapshot: dict | None = None
+        state_key: str | None = None
 
         if context is not None:
             session_state = context.get(SESSION_NAMESPACE, SessionState)
             qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
+            state_key = context.state_key
 
             # Use incoming values (what client sent) so delta shows server-side updates
             if session_state is not None:
@@ -315,7 +317,7 @@ def create_qa_toolset(
             state_event = build_chat_state_delta(
                 old_state_snapshot,
                 new_state_snapshot,
-                state_key=session_state.state_key,
+                state_key=state_key,
             )
 
             answer_text = qa_result.answer

haiku_rag_slim/haiku/rag/tools/search.py

@@ -68,8 +68,10 @@ def create_search_toolset(
         """
         session_state: SessionState | None = None
         old_session_state: SessionState | None = None
+        state_key: str | None = None
         if context is not None:
             session_state = context.get(SESSION_NAMESPACE, SessionState)
+            state_key = context.state_key
             if session_state is not None:
                 old_session_state = session_state.model_copy(deep=True)
 
@@ -130,7 +132,11 @@ def create_search_toolset(
             formatted = f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines)
 
             if old_session_state is not None:
-                state_event = compute_state_delta(old_session_state, session_state)
+                state_event = compute_state_delta(
+                    old_session_state,
+                    session_state,
+                    state_key=state_key,
+                )
                 if state_event is not None:
                     return ToolReturn(
                         return_value=formatted,

haiku_rag_slim/haiku/rag/tools/session.py

@@ -24,7 +24,6 @@ class SessionState(BaseModel):
     document_filter: list[str] = []
     citation_registry: dict[str, int] = {}
     citations: list[Citation] = []
-    state_key: str | None = Field(default=None, exclude=True)
 
     def get_or_assign_index(self, chunk_id: str) -> int:
         """Get or assign a stable citation index for a chunk_id.
@@ -44,16 +43,16 @@ class SessionState(BaseModel):
 def compute_state_delta(
     old_state: SessionState,
     new_state: SessionState,
+    state_key: str | None = None,
 ) -> StateDeltaEvent | None:
     """Compute state delta between old and new session state.
 
     Returns a StateDeltaEvent if there are changes, None otherwise.
-    The state_key from new_state is used for namespacing.
     """
     return compute_combined_state_delta(
         old_state.model_dump(mode="json"),
         new_state.model_dump(mode="json"),
-        state_key=new_state.state_key,
+        state_key=state_key,
     )
 
 

tests/tools/test_context.py

@@ -201,3 +201,114 @@ def test_get_without_type():
 
     result = ctx.get("ns")
     assert result is state
+
+
+def test_tool_context_state_key_default_none():
+    """Test ToolContext state_key defaults to None."""
+    ctx = ToolContext()
+    assert ctx.state_key is None
+
+
+def test_tool_context_state_key_set():
+    """Test ToolContext state_key can be set."""
+    ctx = ToolContext()
+    ctx.state_key = "haiku.rag.chat"
+    assert ctx.state_key == "haiku.rag.chat"
+
+
+# =============================================================================
+# ToolContextCache Tests
+# =============================================================================
+
+
+def test_tool_context_cache_get_or_create_new():
+    """Test get_or_create returns a new context with is_new=True."""
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache()
+    context, is_new = cache.get_or_create("thread-1")
+
+    assert isinstance(context, ToolContext)
+    assert is_new is True
+
+
+def test_tool_context_cache_get_or_create_existing():
+    """Test get_or_create returns existing context with is_new=False."""
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache()
+    ctx1, is_new1 = cache.get_or_create("thread-1")
+    ctx2, is_new2 = cache.get_or_create("thread-1")
+
+    assert ctx2 is ctx1
+    assert is_new1 is True
+    assert is_new2 is False
+
+
+def test_tool_context_cache_different_keys():
+    """Test get_or_create returns different contexts for different keys."""
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache()
+    ctx1, _ = cache.get_or_create("thread-1")
+    ctx2, _ = cache.get_or_create("thread-2")
+
+    assert ctx1 is not ctx2
+
+
+def test_tool_context_cache_ttl_expiry():
+    """Test that contexts are evicted after TTL expires."""
+    from datetime import timedelta
+
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache(ttl=timedelta(seconds=0))
+    ctx1, _ = cache.get_or_create("thread-1")
+
+    # With zero TTL, next access should create a new context
+    ctx2, is_new = cache.get_or_create("thread-1")
+
+    assert ctx2 is not ctx1
+    assert is_new is True
+
+
+def test_tool_context_cache_remove():
+    """Test remove deletes a specific key."""
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache()
+    cache.get_or_create("thread-1")
+    cache.get_or_create("thread-2")
+
+    cache.remove("thread-1")
+
+    ctx, is_new = cache.get_or_create("thread-1")
+    assert is_new is True
+
+    # thread-2 should still exist
+    ctx2, is_new2 = cache.get_or_create("thread-2")
+    assert is_new2 is False
+
+
+def test_tool_context_cache_remove_nonexistent():
+    """Test remove handles nonexistent key gracefully."""
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache()
+    cache.remove("nonexistent")  # Should not raise
+
+
+def test_tool_context_cache_clear():
+    """Test clear removes all entries."""
+    from haiku.rag.tools.context import ToolContextCache
+
+    cache = ToolContextCache()
+    cache.get_or_create("thread-1")
+    cache.get_or_create("thread-2")
+
+    cache.clear()
+
+    ctx1, is_new1 = cache.get_or_create("thread-1")
+    ctx2, is_new2 = cache.get_or_create("thread-2")
+    assert is_new1 is True
+    assert is_new2 is True