fix(chat): restore message history on reconnect; fix event loop blocking

Two related issues caused the screen to clear after each answer:

1. Blocking event loop: the synchronous llm.stream_complete() for-loop
   was running directly in the async on_message handler, blocking the
   event loop for the full LLM response duration. This caused the
   WebSocket to time out and Chainlit to reconnect after each message.

   Fix: replace stream_complete() with asyncio.to_thread(llm.complete)
   so the network-bound LLM call runs in a thread pool and the event
   loop (and WebSocket) stay alive throughout.

2. Reconnect resets history: on_chat_start always created a fresh empty
   state and sent the welcome message, even on WebSocket reconnections
   where cl.user_session still held the existing conversation.

   Fix: if cl.user_session["state"] is already present, replay the
   stored message history into the new thread instead of starting fresh.

https://claude.ai/code/session_01SUWzMzFvSfWiHXA2p6rPg9
This commit is contained in:
Claude
2026-02-22 08:30:52 +00:00
parent 808be0185f
commit 981948c106
+29 -9
View File
@@ -18,6 +18,7 @@ Environment variables (see .env.example):
DATA_DIR Directory for completed registration JSON (default: data/) DATA_DIR Directory for completed registration JSON (default: data/)
""" """
import asyncio
import logging import logging
import uuid import uuid
from datetime import datetime, timezone from datetime import datetime, timezone
@@ -77,11 +78,30 @@ _WELCOME_DE = (
@cl.on_chat_start @cl.on_chat_start
async def on_chat_start() -> None: async def on_chat_start() -> None:
"""Initialise a fresh conversation state and greet the parent.""" """Initialise a fresh conversation state and greet the parent.
If cl.user_session already holds state (WebSocket reconnect after a
network drop), replay the existing message history so the parent sees
the full conversation rather than a blank screen.
"""
existing = cl.user_session.get("state")
if existing:
# Reconnected — restore visual history from our stored state
state = ConversationState.from_dict(existing)
logger.info(
"Session reconnected: %s (%d messages)",
state.conversation_id,
len(state.messages),
)
for msg in state.messages:
author = "Spielgruppe Pumuckl" if msg.role == "assistant" else "Du / You"
await cl.Message(content=msg.content, author=author).send()
return
# Brand new session
session_id = str(uuid.uuid4()) session_id = str(uuid.uuid4())
state = ConversationState(conversation_id=session_id) state = ConversationState(conversation_id=session_id)
cl.user_session.set("state", state.to_dict()) cl.user_session.set("state", state.to_dict())
logger.info("Chat session started: %s", session_id) logger.info("Chat session started: %s", session_id)
await cl.Message(content=_WELCOME_DE).send() await cl.Message(content=_WELCOME_DE).send()
@@ -100,14 +120,14 @@ async def on_message(message: cl.Message) -> None:
# --- Build system prompt --- # --- Build system prompt ---
system = build_system_prompt(_kb, state) system = build_system_prompt(_kb, state)
# --- Collect LLM response (JSON), then display only the reply field --- # --- Call LLM in a thread so the async event loop (and WebSocket) stay alive ---
# The LLM returns a structured JSON object; we must not stream raw tokens # litellm.completion is synchronous; running it directly in an async handler
# to the user because they would see the JSON wrapper, not the reply text. # blocks the event loop for the full response duration and causes WebSocket
full_content = "" # timeouts that trigger on_chat_start again (clearing the screen).
try: try:
for chunk in llm.stream_complete(_config.ai_model, system, state.messages): full_content = await asyncio.to_thread(
full_content += chunk llm.complete, _config.ai_model, system, state.messages
)
except Exception: except Exception:
logger.exception("LLM call failed for session %s", state.conversation_id) logger.exception("LLM call failed for session %s", state.conversation_id)
error_text = fallback_message(state.language) error_text = fallback_message(state.language)