fix(chat-context): clean up session entries on disconnect

[axelray-dev]

Fixes #2951

Problem

The global chat_contexts dict in backend/chainlit/chat_context.py accumulates a message list per session_id but never removes entries when sessions disconnect. Over time this causes unbounded memory growth and eventual OOM.

Root Cause

The disconnect handler in socket.py cleans up user_sessions and calls session.delete(), but does not remove the session's entry from chat_contexts.

Fix

1. Added ChatContext.delete_session(session_id) to remove a session's entry from chat_contexts.
2. Called it from the clear() function inside the disconnect handler, alongside the existing user_sessions cleanup.

Tests

3 new tests in backend/tests/test_chat_context.py:

• test_delete_session_removes_entry: verifies entry is removed after adding a message
• test_delete_session_nonexistent: verifies no error for unknown session ID
• test_delete_session_without_session: verifies deletion works without an active Chainlit context

All 31 tests pass: PYTHONPATH=. python -m pytest tests/test_chat_context.py -v

---

Summary by cubic

Fixes #2951 by removing session entries from chat_contexts on disconnect to prevent unbounded memory growth. Previously entries persisted after disconnect; now socket.clear() deletes them, and cleanup runs even when session persistence fails.

• Bug Fixes
  • Added ChatContext.delete_session(session_id) and call it in socket.clear() after user_sessions cleanup.
  • Moved cleanup into a finally block so it always runs, even if on_chat_end or persist_user_session raises.
  • Added tests for deletion, unknown IDs, deletion without an active context, and cleanup when persistence fails.

^{Written for commit 595f48e. Summary will update on new commits.}

[cubic-dev-ai]

No issues found across 3 files

_{Re-trigger cubic}

[dokterbob]

@codex review

[dokterbob]

@axelray-dev This looks good to me but I would love it if at least one other user/dev could validate that this doesn't have unintended side-effects. Unfortunately, I lack the time to do this investigation myself.

[Copilot]

Pull request overview

Fixes chat-context memory growth by removing session entries after WebSocket disconnect cleanup.

Changes:

• Adds explicit chat-context session deletion.
• Integrates deletion into disconnect cleanup.
• Adds unit tests for deletion behavior.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
backend/chainlit/chat_context.py	Adds session-entry deletion.
backend/chainlit/socket.py	Deletes chat context during disconnect cleanup.
backend/tests/test_chat_context.py	Tests deletion scenarios.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[dokterbob]

Please address the copilot feedback thoroughly; e.g. ensure cleanup even when there's errors during the chat, preferably demonstrated in failing unit tests.

[axelray-dev]

Thanks for the feedback. I moved both the chat-end callback and session persistence inside a try/finally block, so the existing immediate cleanup or timeout cleanup is still scheduled when either operation raises. I also added a regression test that makes persist_user_session() fail and verifies that user-session state, chat context, and the WebSocket session are still cleaned up. Targeted backend tests pass (67 passed) and Ruff is clean. I left the review thread unresolved for verification.