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

backend/chainlit/chat_context.py

@@ -48,6 +48,10 @@ def clear(self) -> None:
         if context.session and context.session.id in chat_contexts:
             chat_contexts[context.session.id] = []
 
+    def delete_session(self, session_id: str) -> None:
+        """Remove a session's chat context to free memory."""
+        chat_contexts.pop(session_id, None)
+
     def to_openai(self):
         messages = []
         for message in self.get():

backend/chainlit/socket.py

@@ -261,29 +261,32 @@ async def disconnect(sid):
 
     init_ws_context(session)
 
-    if config.code.on_chat_end:
-        await config.code.on_chat_end()
-
-    if session.thread_id and session.has_first_interaction:
-        await persist_user_session(session.thread_id, session.to_persistable())
-
     async def clear(_sid):
         if session := WebsocketSession.get(_sid):
             # Clean up the user session
             if session.id in user_sessions:
                 user_sessions.pop(session.id)
+            # Clean up the chat context
+            chat_context.delete_session(session.id)
             # Clean up the session
             await session.delete()
 
-    if session.to_clear:
-        await clear(sid)
-    else:
+    try:
+        if config.code.on_chat_end:
+            await config.code.on_chat_end()
+
+        if session.thread_id and session.has_first_interaction:
+            await persist_user_session(session.thread_id, session.to_persistable())
+    finally:
+        if session.to_clear:
+            await clear(sid)
+        else:
 
-        async def clear_on_timeout(_sid):
-            await asyncio.sleep(config.project.session_timeout)
-            await clear(_sid)
+            async def clear_on_timeout(_sid):
+                await asyncio.sleep(config.project.session_timeout)
+                await clear(_sid)
 
-        asyncio.ensure_future(clear_on_timeout(sid))
+            asyncio.ensure_future(clear_on_timeout(sid))
 
 
 @sio.on("stop")  # pyright: ignore [reportOptionalCall]

backend/tests/test_chat_context.py

@@ -225,6 +225,34 @@ def test_clear_existing_context(self):
             assert "session_123" in chat_contexts
             assert chat_contexts["session_123"] == []
 
+    def test_delete_session_removes_entry(self):
+        """Test delete_session removes the session context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.add(mock_message)
+
+        chat_context.delete_session("session_123")
+
+        assert "session_123" not in chat_contexts
+
+    def test_delete_session_nonexistent(self):
+        """Test delete_session does nothing for an unknown session."""
+        chat_context.delete_session("unknown_session")
+
+        assert chat_contexts == {}
+
+    def test_delete_session_without_session(self):
+        """Test delete_session works with an explicit ID and no context."""
+        chat_contexts["session_123"] = [Mock()]
+
+        with mock_chainlit_context(session=None):
+            chat_context.delete_session("session_123")
+
+        assert "session_123" not in chat_contexts
+
     def test_to_openai_with_assistant_message(self):
         """Test to_openai converts assistant messages correctly."""
         mock_session = Mock()

backend/tests/test_socket.py

@@ -10,6 +10,7 @@
     _get_token_from_cookie,
     clean_session,
     connection_successful,
+    disconnect,
     load_user_env,
     persist_user_session,
     restore_existing_session,
@@ -445,6 +446,47 @@ async def test_clean_session_without_session(self):
             await clean_session("socket_123")
 
 
+class TestDisconnectCleanup:
+    @pytest.mark.asyncio
+    async def test_cleanup_runs_when_persisting_user_session_fails(self):
+        """Disconnect cleanup must run even when persistence raises."""
+        from chainlit.chat_context import chat_contexts
+        from chainlit.user_session import user_sessions
+
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.id = "session_123"
+        mock_session.thread_id = "thread_123"
+        mock_session.has_first_interaction = True
+        mock_session.to_clear = True
+        mock_session.to_persistable.return_value = {}
+        mock_session.delete = AsyncMock()
+
+        user_sessions["session_123"] = {"key": "value"}
+        chat_contexts["session_123"] = [Mock()]
+
+        try:
+            with (
+                patch.object(WebsocketSession, "get", return_value=mock_session),
+                patch("chainlit.socket.init_ws_context"),
+                patch("chainlit.socket.config") as mock_config,
+                patch(
+                    "chainlit.socket.persist_user_session",
+                    new=AsyncMock(side_effect=RuntimeError("database unavailable")),
+                ),
+            ):
+                mock_config.code.on_chat_end = None
+
+                with pytest.raises(RuntimeError, match="database unavailable"):
+                    await disconnect("socket_123")
+
+            assert "session_123" not in user_sessions
+            assert "session_123" not in chat_contexts
+            mock_session.delete.assert_awaited_once_with()
+        finally:
+            user_sessions.pop("session_123", None)
+            chat_contexts.pop("session_123", None)
+
+
 class TestSocketEdgeCases:
     """Test suite for socket edge cases."""