diff --git a/backend/chainlit/chat_context.py b/backend/chainlit/chat_context.py index 81bf66b3d2..436ea9dd9c 100644 --- a/backend/chainlit/chat_context.py +++ b/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(): diff --git a/backend/chainlit/socket.py b/backend/chainlit/socket.py index 74306005ff..c55b71c9e8 100644 --- a/backend/chainlit/socket.py +++ b/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] diff --git a/backend/tests/test_chat_context.py b/backend/tests/test_chat_context.py index dae92ac385..431c96ba35 100644 --- a/backend/tests/test_chat_context.py +++ b/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() diff --git a/backend/tests/test_socket.py b/backend/tests/test_socket.py index e20959c6d6..11ca6994d3 100644 --- a/backend/tests/test_socket.py +++ b/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."""