danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove session_id from state layer, remove module-level caches

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-11 14:55:53 +02:00
parent d88f2f003a
commit 9e63ff1ff7
No known key found for this signature in database
12 changed files with 131 additions and 460 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
app/backend/main.py
View file

@ -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)

104
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -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,
)

97
haiku_rag_slim/haiku/rag/agents/chat/context.py
View file

@ -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))

4
haiku_rag_slim/haiku/rag/agents/chat/state.py
View file

@ -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],

7
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -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,
)

22
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -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

2
haiku_rag_slim/haiku/rag/tools/session.py
View file

@ -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] = []

155
tests/agents/chat/test_chat_agent.py
View file

@ -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

126
tests/agents/chat/test_context.py
View file

@ -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):

3
tests/agents/chat/test_features.py
View file

@ -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

20
tests/agents/chat/test_state.py
View file

@ -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),

30
tests/chat/test_chat_app.py
View file

@ -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": {},