fix(chat-context): clean up session entries on disconnect - #2959
fix(chat-context): clean up session entries on disconnect#2959axelray-dev wants to merge 5 commits into
Conversation
Co-Authored-By: Codex <codex@openai.com>
|
@codex review |
|
@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. |
There was a problem hiding this comment.
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.
| if session.id in user_sessions: | ||
| user_sessions.pop(session.id) | ||
| # Clean up the chat context | ||
| chat_context.delete_session(session.id) |
dokterbob
left a comment
There was a problem hiding this comment.
Please address the copilot feedback thoroughly; e.g. ensure cleanup even when there's errors during the chat, preferably demonstrated in failing unit tests.
Fixes #2951
Problem
The global
chat_contextsdict inbackend/chainlit/chat_context.pyaccumulates a message list persession_idbut never removes entries when sessions disconnect. Over time this causes unbounded memory growth and eventual OOM.Root Cause
The
disconnecthandler insocket.pycleans upuser_sessionsand callssession.delete(), but does not remove the session's entry fromchat_contexts.Fix
ChatContext.delete_session(session_id)to remove a session's entry fromchat_contexts.clear()function inside thedisconnecthandler, alongside the existinguser_sessionscleanup.Tests
3 new tests in
backend/tests/test_chat_context.py:test_delete_session_removes_entry: verifies entry is removed after adding a messagetest_delete_session_nonexistent: verifies no error for unknown session IDtest_delete_session_without_session: verifies deletion works without an active Chainlit contextAll 31 tests pass:
PYTHONPATH=. python -m pytest tests/test_chat_context.py -vSummary by cubic
Removes session entries from
chat_contextson disconnect to prevent memory leaks. Previously entries persisted; nowsocket.clear()deletes them, stopping unbounded growth and OOM.ChatContext.delete_session(session_id)and call it insocket.clear()afteruser_sessionscleanup.Written for commit 91a90ab. Summary will update on new commits.