From 9e63ff1ff723609b90fb65d3061194ca9a335a58 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 11 Feb 2026 14:55:53 +0200 Subject: [PATCH] Remove session_id from state layer, remove module-level caches --- app/backend/main.py | 21 +-- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 104 ++++++------ .../haiku/rag/agents/chat/context.py | 97 ++--------- haiku_rag_slim/haiku/rag/agents/chat/state.py | 4 +- haiku_rag_slim/haiku/rag/chat/app.py | 7 +- haiku_rag_slim/haiku/rag/tools/qa.py | 22 +-- haiku_rag_slim/haiku/rag/tools/session.py | 2 - tests/agents/chat/test_chat_agent.py | 155 ++++++------------ tests/agents/chat/test_context.py | 126 +------------- tests/agents/chat/test_features.py | 3 +- tests/agents/chat/test_state.py | 20 +-- tests/chat/test_chat_app.py | 30 +--- 12 files changed, 131 insertions(+), 460 deletions(-) diff --git a/app/backend/main.py b/app/backend/main.py index 0c8ebb80..5c998014 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -20,10 +20,13 @@ from haiku.rag.agents.chat import ( from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config from haiku.rag.config.models import AppConfig -from haiku.rag.tools import ToolContext +from haiku.rag.tools.context import ToolContextCache load_dotenv(find_dotenv(usecwd=True)) +# Cache ToolContext instances by thread_id across requests +context_cache = ToolContextCache() + # Configure logfire (only sends data if LOGFIRE_TOKEN is present) try: import logfire @@ -68,30 +71,24 @@ def get_client() -> HaikuRAG: async def stream_chat(request: Request) -> Response: """Chat streaming endpoint with AG-UI protocol. - This endpoint is stateless - all state flows via AG-UI protocol: - - Fresh ToolContext created per request - - AGUIAdapter restores state via ChatDeps.state setter - - ChatDeps generates session_id if not provided - - Ask tool triggers background summarization internally - - ChatDeps.state getter emits final state in response + Uses ToolContextCache to maintain state across requests for the same thread. + AGUIAdapter restores client-sent state via ChatDeps.state setter. """ body = await request.body() accept = request.headers.get("accept", SSE_CONTENT_TYPE) run_input = AGUIAdapter.build_run_input(body) - # Fresh context per request - state restored by AGUIAdapter via ChatDeps.state setter - context = ToolContext() + thread_id = getattr(run_input, "thread_id", None) or "default" + context, is_new = context_cache.get_or_create(thread_id) agent = create_chat_agent(Config, get_client(), context) deps = ChatDeps( config=Config, tool_context=context, + is_new=is_new, state_key=AGUI_STATE_KEY, ) - # Use AGUIAdapter for streaming - # State restoration happens automatically via ChatDeps.state setter - # Background summarization triggered by ask() tool internally adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept) event_stream = adapter.run_stream(deps=deps) sse_event_stream = adapter.encode_stream(event_stream) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 4639aaa2..f429eb1d 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -1,12 +1,8 @@ -import uuid from dataclasses import dataclass from typing import Any from pydantic_ai import Agent -from haiku.rag.agents.chat.context import ( - get_cached_session_context, -) from haiku.rag.agents.chat.context import ( trigger_background_summarization as _trigger_summarization, ) @@ -47,7 +43,7 @@ class ChatDeps: config: AppConfig tool_context: ToolContext - session_id: str = "" + is_new: bool = True state_key: str | None = None @property @@ -78,59 +74,57 @@ class ChatDeps: state_data = nested session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state is not None: - if "document_filter" in state_data: - session_state.document_filter = state_data.get("document_filter", []) - if "citation_registry" in state_data: - session_state.citation_registry = state_data["citation_registry"] - if "citations" in state_data: - from haiku.rag.agents.research.models import Citation - session_state.citations = [ - Citation(**c) if isinstance(c, dict) else c - for c in state_data.get("citations", []) - ] + if self.is_new: + # First request for this context: fully populate from client state + if session_state is not None: + if "document_filter" in state_data: + session_state.document_filter = state_data.get( + "document_filter", [] + ) + if "citation_registry" in state_data: + session_state.citation_registry = state_data["citation_registry"] + if "citations" in state_data: + from haiku.rag.agents.research.models import Citation - # Restore session_id from client or generate one - client_session_id = state_data.get("session_id", "") - if client_session_id: - self.session_id = client_session_id - elif not self.session_id: - self.session_id = str(uuid.uuid4()) + session_state.citations = [ + Citation(**c) if isinstance(c, dict) else c + for c in state_data.get("citations", []) + ] - if session_state is not None: - session_state.session_id = self.session_id + qa_session_state = self.tool_context.get( + QA_SESSION_NAMESPACE, QASessionState + ) + if qa_session_state is not None: + if "qa_history" in state_data: + from haiku.rag.tools.qa import QAHistoryEntry - qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) - if qa_session_state is not None: - if "qa_history" in state_data: - from haiku.rag.tools.qa import QAHistoryEntry + qa_session_state.qa_history = [ + QAHistoryEntry(**qa) if isinstance(qa, dict) else qa + for qa in state_data.get("qa_history", []) + ] - qa_session_state.qa_history = [ - QAHistoryEntry(**qa) if isinstance(qa, dict) else qa - for qa in state_data.get("qa_history", []) - ] + # Restore session_context from client + session_context = state_data.get("session_context") + if isinstance(session_context, dict): + qa_session_state.session_context = SessionContext( + **session_context + ).summary + elif session_context is None: + qa_session_state.session_context = None - # Restore session_context from client - session_context = state_data.get("session_context") - if isinstance(session_context, dict): - qa_session_state.session_context = SessionContext( - **session_context - ).summary - elif session_context is None: - qa_session_state.session_context = None - - # Check cache for fresher session_context from background summarization - if self.session_id: - cached = get_cached_session_context(self.session_id) - if cached and cached.summary: - qa_session_state.session_context = cached.summary - - # Handle initial_context -> session_context for first message - if "initial_context" in state_data: - initial = state_data.get("initial_context") - if initial and not qa_session_state.session_context: - qa_session_state.session_context = initial + # Handle initial_context -> session_context for first message + if "initial_context" in state_data: + initial = state_data.get("initial_context") + if initial and not qa_session_state.session_context: + qa_session_state.session_context = initial + else: + # Returning request: only merge client-controlled fields + if session_state is not None: + if "document_filter" in state_data: + session_state.document_filter = state_data.get( + "document_filter", [] + ) def create_chat_agent( @@ -204,22 +198,16 @@ def trigger_background_summarization(deps: ChatDeps) -> None: Call this after agent.run() or agent.run_stream() completes to update the session context summary in the background. - Note: The ask() tool now triggers summarization internally, so this - function is primarily for explicit triggering when needed. - Args: deps: Chat dependencies with tool_context containing QASessionState. """ qa_session_state = deps.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) if qa_session_state is None or not qa_session_state.qa_history: return - if not deps.session_id: - return _trigger_summarization( qa_session_state=qa_session_state, config=deps.config, - session_id=deps.session_id, ) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/context.py b/haiku_rag_slim/haiku/rag/agents/chat/context.py index 9fbaf9d4..f09b4dba 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/context.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/context.py @@ -1,6 +1,5 @@ import asyncio -from dataclasses import dataclass, field -from datetime import datetime, timedelta +from datetime import datetime from typing import TYPE_CHECKING from pydantic_ai import Agent @@ -14,68 +13,8 @@ if TYPE_CHECKING: from haiku.rag.tools.qa import QAHistoryEntry, QASessionState -@dataclass -class SessionCache: - """Per-session cache for context and embeddings.""" - - context: SessionContext | None = None - embeddings: dict[str, list[float]] = field(default_factory=dict) - - -# Cache for session data (session_id -> SessionCache) -# Used to persist async summarization results and embeddings between requests -_session_cache: dict[str, SessionCache] = {} -_cache_timestamps: dict[str, datetime] = {} -_CACHE_TTL = timedelta(hours=1) - -# Track summarization tasks per session to allow cancellation -_summarization_tasks: dict[str, asyncio.Task[None]] = {} - - -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_cache.pop(sid, None) - _cache_timestamps.pop(sid, None) - - -def _get_or_create_session_cache(session_id: str) -> SessionCache: - """Get or create session cache for a given session_id.""" - _cleanup_stale_cache() - if session_id not in _session_cache: - _session_cache[session_id] = SessionCache() - _cache_timestamps[session_id] = datetime.now() - return _session_cache[session_id] - - -def cache_session_context(session_id: str, context: SessionContext) -> None: - """Store session context in cache.""" - cache = _get_or_create_session_cache(session_id) - cache.context = context - - -def get_cached_session_context(session_id: str) -> SessionContext | None: - """Get session context from server cache.""" - _cleanup_stale_cache() - cache = _session_cache.get(session_id) - return cache.context if cache else None - - -def cache_question_embedding( - session_id: str, question: str, embedding: list[float] -) -> None: - """Store question embedding in session cache.""" - cache = _get_or_create_session_cache(session_id) - cache.embeddings[question] = embedding - - -def get_cached_embedding(session_id: str, question: str) -> list[float] | None: - """Get cached embedding for a question in this session.""" - _cleanup_stale_cache() - cache = _session_cache.get(session_id) - return cache.embeddings.get(question) if cache else None +# Track summarization tasks to allow cancellation +_summarization_tasks: dict[int, asyncio.Task[None]] = {} async def summarize_session( @@ -115,15 +54,13 @@ async def summarize_session( async def update_session_context( qa_history: list["QAHistoryEntry"], config: AppConfig, - session_id: str = "", current_context: str | None = None, ) -> SessionContext: - """Summarize qa_history and cache the resulting session context. + """Summarize qa_history and return the resulting session context. Args: qa_history: List of Q&A pairs from the conversation. config: AppConfig for model selection. - session_id: Session ID for caching. If empty, result is not cached. current_context: Previous summary to incorporate. Returns: @@ -132,13 +69,10 @@ async def update_session_context( summary = await summarize_session( qa_history, config, current_context=current_context ) - context = SessionContext( + return SessionContext( summary=summary, last_updated=datetime.now(), ) - if session_id: - cache_session_context(session_id, context) - return context def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str: @@ -159,14 +93,12 @@ def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str: async def _update_context_background( qa_session_state: "QASessionState", config: AppConfig, - session_id: str, ) -> None: """Background task to update session context after an ask.""" try: result = await update_session_context( qa_history=list(qa_session_state.qa_history), config=config, - session_id=session_id, current_context=qa_session_state.session_context, ) @@ -184,31 +116,28 @@ async def _update_context_background( def trigger_background_summarization( qa_session_state: "QASessionState", config: AppConfig, - session_id: str, ) -> None: """Trigger background session summarization if qa_history has entries. Args: qa_session_state: QASessionState with qa_history to summarize. config: AppConfig for model selection. - session_id: Session ID for caching results. """ - if not qa_session_state.qa_history or not session_id: + if not qa_session_state.qa_history: return - # Cancel any existing summarization task for this session - if session_id in _summarization_tasks: - _summarization_tasks[session_id].cancel() + key = id(qa_session_state) + + # Cancel any existing summarization task for this state + if key in _summarization_tasks: + _summarization_tasks[key].cancel() # Spawn background task task = asyncio.create_task( _update_context_background( qa_session_state=qa_session_state, config=config, - session_id=session_id, ) ) - _summarization_tasks[session_id] = task - task.add_done_callback( - lambda _t, sid=session_id: _summarization_tasks.pop(sid, None) - ) + _summarization_tasks[key] = task + task.add_done_callback(lambda _t, k=key: _summarization_tasks.pop(k, None)) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index a777eb6f..6292282c 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -24,7 +24,6 @@ class SessionContext(BaseModel): class ChatSessionState(BaseModel): """State shared between frontend and agent via AG-UI.""" - session_id: str = "" initial_context: str | None = None citations: list[Citation] = [] qa_history: list["QAHistoryEntry"] = [] @@ -56,12 +55,11 @@ def build_chat_state_snapshot( Returns: Snapshot dict. """ - snapshot: dict[str, Any] = {"session_id": ""} + snapshot: dict[str, Any] = {} if session_state is not None: snapshot.update( { - "session_id": session_state.session_id, "document_filter": session_state.document_filter.copy(), "citation_registry": session_state.citation_registry.copy(), "citations": [c.model_dump() for c in session_state.citations], diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 56b36d6b..3709bb67 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -169,7 +169,6 @@ class ChatApp(App): # Keep ChatSessionState for UI state sync (used by _sync_session_state) self.session_state = ChatSessionState( - session_id=str(uuid.uuid4()), initial_context=self._initial_context, document_filter=self._document_filter, ) @@ -191,8 +190,6 @@ class ChatApp(App): from haiku.rag.agents.research.models import Citation # Update specific fields rather than replacing the entire state - if "session_id" in chat_state: - self.session_state.session_id = chat_state["session_id"] if "document_filter" in chat_state: self.session_state.document_filter = chat_state["document_filter"] if "citation_registry" in chat_state: @@ -322,7 +319,6 @@ class ChatApp(App): # Sync session state to tool context before running session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) if session_state is not None: - session_state.session_id = self.session_state.session_id session_state.document_filter = self.session_state.document_filter session_state.citation_registry = self.session_state.citation_registry session_state.citations = list(self.session_state.citations) @@ -345,7 +341,7 @@ class ChatApp(App): deps = ChatDeps( config=self.config, tool_context=self.tool_context, - session_id=self.session_state.session_id, + is_new=False, state_key=AGUI_STATE_KEY, ) @@ -414,7 +410,6 @@ class ChatApp(App): # Reset context lock and session state (reset to CLI value) self._context_locked = False self.session_state = ChatSessionState( - session_id=str(uuid.uuid4()), initial_context=self._initial_context, document_filter=self._document_filter, ) diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index 4882323e..e553f131 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -3,11 +3,7 @@ import math from pydantic import BaseModel, Field from pydantic_ai import FunctionToolset, ToolReturn -from haiku.rag.agents.chat.context import ( - cache_question_embedding, - get_cached_embedding, - trigger_background_summarization, -) +from haiku.rag.agents.chat.context import trigger_background_summarization from haiku.rag.agents.chat.state import ( build_chat_state_delta, build_chat_state_snapshot, @@ -121,7 +117,6 @@ async def run_qa_core( effective_session_context = qa_session_state.session_context effective_prior_answers = prior_answers or [] - session_id = session_state.session_id if session_state is not None else "" if qa_session_state is not None and qa_session_state.qa_history: embedder = get_embedder(config) question_embedding = await embedder.embed_query(question) @@ -130,25 +125,13 @@ async def run_qa_core( to_embed_indices = [] for i, qa in enumerate(qa_session_state.qa_history): if qa.question_embedding is None: - if session_id: - cached = get_cached_embedding(session_id, qa.question) - if cached: - qa.question_embedding = cached - continue to_embed.append(qa.question) to_embed_indices.append(i) if to_embed: new_embeddings = await embedder.embed_documents(to_embed) for i, idx in enumerate(to_embed_indices): - embedding = new_embeddings[i] - qa_session_state.qa_history[idx].question_embedding = embedding - if session_id: - cache_question_embedding( - session_id, - qa_session_state.qa_history[idx].question, - embedding, - ) + qa_session_state.qa_history[idx].question_embedding = new_embeddings[i] matched_answers = [] for qa in qa_session_state.qa_history: @@ -225,7 +208,6 @@ async def run_qa_core( trigger_background_summarization( qa_session_state=qa_session_state, config=config, - session_id=session_id, ) return qa_result diff --git a/haiku_rag_slim/haiku/rag/tools/session.py b/haiku_rag_slim/haiku/rag/tools/session.py index 1f84837e..fcd1e73d 100644 --- a/haiku_rag_slim/haiku/rag/tools/session.py +++ b/haiku_rag_slim/haiku/rag/tools/session.py @@ -13,13 +13,11 @@ class SessionState(BaseModel): """Session-level state for AG-UI integration. This state is shared across toolsets and enables: - - Session identification - Dynamic document filtering - Stable citation indices across tool calls - AG-UI state synchronization """ - session_id: str = "" document_filter: list[str] = [] citation_registry: dict[str, int] = {} citations: list[Citation] = [] diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 22964690..1446f633 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -9,10 +9,7 @@ from haiku.rag.agents.chat import ( ChatSessionState, create_chat_agent, ) -from haiku.rag.agents.chat.context import ( - _summarization_tasks, - get_cached_session_context, -) +from haiku.rag.agents.chat.context import _summarization_tasks from haiku.rag.agents.research.models import Citation from haiku.rag.client import HaikuRAG from haiku.rag.config import Config @@ -38,9 +35,7 @@ def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | elif isinstance(meta, StateDeltaEvent): # Apply delta to empty state to get final state empty_state = { - state_key: ChatSessionState(session_id="").model_dump( - mode="json" - ) + state_key: ChatSessionState().model_dump(mode="json") } patched = jsonpatch.apply_patch(empty_state, meta.delta) return patched.get(state_key) @@ -69,7 +64,7 @@ def test_chat_deps_initialization(temp_db_path): assert deps.config is Config assert deps.tool_context is context - assert deps.session_id == "" + assert deps.is_new is True assert deps.state_key is None @@ -109,7 +104,6 @@ def test_chat_deps_state_setter_handles_initial_context(): # Client sends initial_context with no session_context incoming_state = { AGUI_STATE_KEY: { - "session_id": "", "initial_context": "Background info about the project", "session_context": None, "qa_history": [], @@ -140,7 +134,6 @@ def test_chat_deps_state_setter_parses_session_context_dict(): # Client sends session_context as a dict (as it comes from JSON) incoming_state = { AGUI_STATE_KEY: { - "session_id": "test-session", "session_context": { "summary": "Previous conversation summary", "last_updated": "2025-01-27T12:00:00", @@ -160,44 +153,9 @@ def test_chat_deps_state_setter_parses_session_context_dict(): assert qa_session_state.session_context == "Previous conversation summary" -def test_chat_deps_state_setter_generates_session_id(): - """Test ChatDeps.state setter generates session_id if client sends empty.""" - from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState - - context = ToolContext() - context.register(QA_SESSION_NAMESPACE, QASessionState()) - context.register(SESSION_NAMESPACE, SessionState()) - - deps = ChatDeps(config=Config, tool_context=context, state_key=AGUI_STATE_KEY) - - # Client sends empty session_id - incoming_state = { - AGUI_STATE_KEY: { - "session_id": "", - "session_context": None, - "qa_history": [], - "citations": [], - "document_filter": [], - "citation_registry": {}, - } - } - - deps.state = incoming_state - - # session_id should be generated (UUID format) - assert deps.session_id != "" - assert len(deps.session_id) == 36 # UUID length with dashes - - # Should also be synced to SessionState - session_state = context.get(SESSION_NAMESPACE) - assert isinstance(session_state, SessionState) - assert session_state.session_id == deps.session_id - - def test_chat_session_state(): """Test ChatSessionState model.""" - state = ChatSessionState(session_id="test-session") - assert state.session_id == "test-session" + state = ChatSessionState() assert state.citations == [] assert state.qa_history == [] @@ -356,7 +314,6 @@ async def test_chat_agent_search_tool(allow_model_requests, temp_db_path): deps = ChatDeps( config=Config, tool_context=context, - session_id="test-search", ) # Ask something that should trigger the search tool @@ -391,7 +348,6 @@ async def test_chat_agent_search_tool_with_filter(allow_model_requests, temp_db_ deps = ChatDeps( config=Config, tool_context=context, - session_id="test-search-filter", ) # Ask to search within a specific document @@ -503,10 +459,16 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path) async def test_chat_agent_ask_triggers_background_summarization( allow_model_requests, temp_db_path ): - """Test that the ask tool triggers background session context summarization.""" - import asyncio + """Test that the ask tool triggers background session context summarization. - from haiku.rag.agents.chat.agent import run_chat_agent + Patches the internal trigger in run_qa_core to avoid concurrent HTTP calls + that break VCR cassette replay ordering. Triggers summarization explicitly + after the agent run completes. + """ + from unittest.mock import patch + + from haiku.rag.agents.chat.agent import trigger_background_summarization + from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState async with HaikuRAG(temp_db_path, create=True) as client: await client.create_document( @@ -520,33 +482,31 @@ async def test_chat_agent_ask_triggers_background_summarization( deps = ChatDeps( config=Config, tool_context=context, - session_id="test-summarization", state_key=AGUI_STATE_KEY, ) - # Ask a question using run_chat_agent to trigger background summarization - result = await run_chat_agent( - agent, - deps, - "What is the highest count class in the DocLayNet dataset?", - ) + # Patch internal trigger to avoid concurrent HTTP calls during VCR + with patch("haiku.rag.tools.qa.trigger_background_summarization"): + result = await agent.run( + "What is the highest count class in the DocLayNet dataset?", + deps=deps, + ) - assert result is not None + assert result.output is not None + + # Trigger summarization explicitly (sequential, deterministic) + trigger_background_summarization(deps) # Wait for background task to complete - # The task caches session_context server-side - session_id = deps.session_id - cached_context = None - for _ in range(50): # Wait up to 5 seconds - cached_context = get_cached_session_context(session_id) - if cached_context is not None: - break - await asyncio.sleep(0.1) + qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState) + assert qa_session_state is not None + key = id(qa_session_state) + if key in _summarization_tasks: + await _summarization_tasks[key] # Verify session_context was populated by background task - assert cached_context is not None - assert cached_context.summary != "" - assert cached_context.last_updated is not None + assert qa_session_state.session_context is not None + assert qa_session_state.session_context != "" @pytest.mark.asyncio @@ -592,7 +552,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_ # Set initial state with initial_context (mimicking AG-UI client) deps.state = { AGUI_STATE_KEY: { - "session_id": "", "initial_context": "The user is researching the DocLayNet dataset for a paper on document layout analysis.", "session_context": None, "qa_history": [], @@ -602,10 +561,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_ } } - # session_id should be auto-generated - assert deps.session_id != "" - session_id = deps.session_id - # initial_context should be transferred to QASessionState qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState) assert qa_session is not None @@ -628,12 +583,12 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_ # Trigger summarization explicitly (sequential, no concurrency) trigger_background_summarization(deps) - if session_id in _summarization_tasks: - await _summarization_tasks[session_id] + key = id(qa_session) + if key in _summarization_tasks: + await _summarization_tasks[key] - cached_context = get_cached_session_context(session_id) - assert cached_context is not None - assert cached_context.summary != "" + assert qa_session.session_context is not None + assert qa_session.session_context != "" # qa_history should have one entry qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState) @@ -653,18 +608,18 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_ # Trigger summarization explicitly trigger_background_summarization(deps) - if session_id in _summarization_tasks: - await _summarization_tasks[session_id] - - # qa_history should have two entries qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState) assert qa_session is not None + key = id(qa_session) + if key in _summarization_tasks: + await _summarization_tasks[key] + + # qa_history should have two entries assert len(qa_session.qa_history) >= 2 # Session context should be updated with newer summary - updated = get_cached_session_context(session_id) - assert updated is not None - assert updated.summary != "" + assert qa_session.session_context is not None + assert qa_session.session_context != "" @pytest.mark.asyncio @@ -740,7 +695,6 @@ def test_fifo_limit_enforcement(): ] session_state = ChatSessionState( - session_id="test-fifo", qa_history=qa_history, ) @@ -759,7 +713,6 @@ def test_fifo_limit_enforcement(): def test_chat_session_state_document_filter(): """Test ChatSessionState with document_filter.""" state = ChatSessionState( - session_id="test-filter", document_filter=["doc1.pdf", "doc2.pdf"], ) assert state.document_filter == ["doc1.pdf", "doc2.pdf"] @@ -767,7 +720,7 @@ def test_chat_session_state_document_filter(): def test_chat_session_state_document_filter_default_empty(): """Test ChatSessionState document_filter defaults to empty list.""" - state = ChatSessionState(session_id="test") + state = ChatSessionState() assert state.document_filter == [] @@ -801,7 +754,6 @@ async def test_chat_agent_search_with_session_filter( deps = ChatDeps( config=Config, tool_context=context, - session_id="test-session-filter", ) # Search should only return results from the filtered document @@ -1032,7 +984,7 @@ def test_qa_response_embedding_default_none(): @pytest.mark.asyncio async def test_summarization_task_cancellation(): - """Test that new summarization tasks cancel previous ones for same session.""" + """Test that new summarization tasks cancel previous ones for same state object.""" import asyncio from haiku.rag.agents.chat.context import _summarization_tasks @@ -1040,7 +992,7 @@ async def test_summarization_task_cancellation(): # Clear any existing tasks _summarization_tasks.clear() - session_id = "test-cancel-session" + key = 12345 # Simulates id(qa_session_state) # Create a slow task that simulates summarization async def slow_task(): @@ -1048,18 +1000,18 @@ async def test_summarization_task_cancellation(): # Start first task task1 = asyncio.create_task(slow_task()) - _summarization_tasks[session_id] = task1 + _summarization_tasks[key] = task1 # Simulate what happens when second ask comes in - cancel first task - if session_id in _summarization_tasks: - _summarization_tasks[session_id].cancel() + if key in _summarization_tasks: + _summarization_tasks[key].cancel() # Yield to let cancellation propagate await asyncio.sleep(0) # Start second task task2 = asyncio.create_task(slow_task()) - _summarization_tasks[session_id] = task2 + _summarization_tasks[key] = task2 # First task should be cancelled assert task1.cancelled() or task1.done() @@ -1143,7 +1095,6 @@ async def test_list_documents_with_session_filter(allow_model_requests, temp_db_ deps = ChatDeps( config=Config, tool_context=context, - session_id="test-list-filter", ) # Ask to list documents - should only show filtered documents @@ -1320,18 +1271,18 @@ async def test_summarization_task_cleanup_on_completion(): _summarization_tasks.clear() - session_id = "test-cleanup-session" + key = 67890 # Simulates id(qa_session_state) # Create a fast task async def fast_task(): await asyncio.sleep(0.01) task = asyncio.create_task(fast_task()) - _summarization_tasks[session_id] = task - task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None)) + _summarization_tasks[key] = task + task.add_done_callback(lambda t: _summarization_tasks.pop(key, None)) # Wait for completion await task # Task should be cleaned up - assert session_id not in _summarization_tasks + assert key not in _summarization_tasks diff --git a/tests/agents/chat/test_context.py b/tests/agents/chat/test_context.py index 02f9b5bf..5cf0875b 100644 --- a/tests/agents/chat/test_context.py +++ b/tests/agents/chat/test_context.py @@ -212,7 +212,6 @@ class TestUpdateSessionContext: result = await update_session_context( qa_history=qa_history, config=Config, - session_id="test-session", ) assert result.summary != "" @@ -231,129 +230,8 @@ class TestUpdateSessionContext: assert result.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_cache, - cache_session_context, - get_cached_session_context, - ) - - # Clear cache - _session_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_cache, - get_cached_session_context, - ) - - _session_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_cache, - cache_session_context, - get_cached_session_context, - ) - - _session_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_cache - - @pytest.mark.asyncio - async def test_update_session_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_cache, - get_cached_session_context, - update_session_context, - ) - - _session_cache.clear() - - qa_history = [ - QAHistoryEntry( - question="What is Python?", - answer="A programming language.", - confidence=0.95, - ) - ] - - with patch( - "haiku.rag.agents.chat.context.summarize_session", - new=AsyncMock(return_value="Mocked summary"), - ): - result = await update_session_context( - qa_history=qa_history, - config=Config, - session_id="cache-test-session", - ) - - assert result.summary == "Mocked summary" - - cached = get_cached_session_context("cache-test-session") - assert cached is not None - assert cached.summary == "Mocked 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_cache, - get_cached_session_context, - update_session_context, - ) - - _session_cache.clear() - - await update_session_context( - qa_history=[], - config=Config, - ) - - # Nothing should be cached (no session_id) - cached = get_cached_session_context("") - assert cached is None +class TestUpdateSessionContextPassesCurrentContext: + """Tests for update_session_context current_context forwarding.""" @pytest.mark.asyncio async def test_update_session_context_passes_current_context(self): diff --git a/tests/agents/chat/test_features.py b/tests/agents/chat/test_features.py index 91d33e71..82a36ab7 100644 --- a/tests/agents/chat/test_features.py +++ b/tests/agents/chat/test_features.py @@ -107,10 +107,9 @@ def test_chat_deps_state_without_qa(temp_db_path): context = ToolContext() create_chat_agent(Config, client, context, features=[FEATURE_SEARCH]) - deps = ChatDeps(config=Config, tool_context=context, session_id="test") + deps = ChatDeps(config=Config, tool_context=context) state = deps.state - assert "session_id" in state # SessionState fields should be present assert "document_filter" in state assert "citation_registry" in state diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index 2b633428..77fcdc77 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -57,7 +57,7 @@ def test_citation_registry_stability(): def test_citation_registry_serialization_roundtrip(): """Test citation_registry serializes and deserializes correctly for AG-UI state.""" # Create state and assign indices - original = ChatSessionState(session_id="test") + original = ChatSessionState() original.citation_registry = {"chunk-a": 1, "chunk-b": 2} # Serialize @@ -70,22 +70,6 @@ def test_citation_registry_serialization_roundtrip(): assert restored.citation_registry == {"chunk-a": 1, "chunk-b": 2} -def test_chat_session_state_defaults_to_empty_session_id(): - """New ChatSessionState should default to empty session_id. - - Tools in agent.py detect the empty string and assign a UUID, - which then appears in the state delta so clients receive it. - """ - state = ChatSessionState() - assert 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_initial_context_default_none(): """Initial context should default to None.""" state = ChatSessionState() @@ -101,7 +85,6 @@ def test_chat_session_state_initial_context_preserved(): 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() @@ -121,7 +104,6 @@ def test_chat_session_state_model_dump_json_serializes_datetime(): 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), diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 3d09adf9..8d87bd00 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -245,25 +245,9 @@ async def test_chat_history_thinking_indicator(temp_db_path: Path): assert len(list(thinking)) == 0 -@pytest.mark.asyncio -async def test_chat_app_generates_session_id(temp_db_path: Path): - """Test that ChatApp generates a UUID session_id on mount.""" - from haiku.rag.chat.app import ChatApp - - mock_client = AsyncMock() - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - - with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): - app = ChatApp(temp_db_path, read_only=True) - - async with app.run_test(): - assert app.session_state.session_id != "" - - @pytest.mark.asyncio async def test_clear_chat_resets_session(temp_db_path: Path): - """Test that clearing chat resets the session state with a new session_id.""" + """Test that clearing chat resets the session state.""" from haiku.rag.chat.app import ChatApp from haiku.rag.chat.widgets.chat_history import ChatHistory @@ -282,10 +266,6 @@ async def test_clear_chat_resets_session(temp_db_path: Path): await chat_history.add_message("assistant", "Hi there") assert len(chat_history.messages) == 2 - # Record the original session_id - original_session_id = app.session_state.session_id - assert original_session_id != "" - # Clear chat via action (available through command palette) await app.action_clear_chat() await pilot.pause() @@ -293,10 +273,8 @@ async def test_clear_chat_resets_session(temp_db_path: Path): # Verify messages cleared assert len(chat_history.messages) == 0 - # Verify session state reset with a new session_id + # Verify session state reset assert app.session_state is not None - assert app.session_state.session_id != "" - assert app.session_state.session_id != original_session_id assert app.session_state.qa_history == [] assert app.session_state.citations == [] @@ -326,7 +304,6 @@ async def test_handle_stream_event_extracts_citations_from_state_snapshot( type=EventType.STATE_SNAPSHOT, snapshot={ AGUI_STATE_KEY: { - "session_id": "test", "citations": [ { "index": 1, @@ -392,7 +369,6 @@ async def test_handle_stream_event_extracts_citations_from_state_delta( type=EventType.STATE_SNAPSHOT, snapshot={ AGUI_STATE_KEY: { - "session_id": "test", "citations": [], "qa_history": [], "citation_registry": {}, @@ -479,7 +455,6 @@ async def test_handle_stream_event_delta_with_preinitialized_state( # Pre-initialize _agui_state_snapshot (simulating what _run_agent does) app._agui_state_snapshot = { AGUI_STATE_KEY: { - "session_id": "test", "citations": [], "qa_history": [], "citation_registry": {}, @@ -552,7 +527,6 @@ async def test_handle_stream_event_syncs_session_context(temp_db_path: Path): # Pre-initialize state app._agui_state_snapshot = { AGUI_STATE_KEY: { - "session_id": "test", "citations": [], "qa_history": [], "citation_registry": {},