From 91a73f8c2879e1903281e980fdff4d90f4b75630 Mon Sep 17 00:00:00 2001
From: axelray-dev <110029405+axelray-dev@users.noreply.github.com>
Date: Fri, 12 Jun 2026 14:57:17 +0800
Subject: [PATCH] fix(chat-context): clean up session entries on disconnect

Co-Authored-By: Codex <codex@openai.com>
---
 backend/chainlit/chat_context.py   |  4 ++++
 backend/chainlit/socket.py         |  2 ++
 backend/tests/test_chat_context.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 34 insertions(+)

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()