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..deb65e3ea0 100644 --- a/backend/chainlit/socket.py +++ b/backend/chainlit/socket.py @@ -272,6 +272,8 @@ async def clear(_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() 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()