fix(embed): harden reset and reconnect recovery

Author: RockChinQ

src/langbot/pkg/platform/sources/websocket_adapter.py

@@ -536,7 +536,7 @@ def get_websocket_messages(
         """Return history for one pipeline/client conversation."""
         conversation_key = self._conversation_key(pipeline_uuid, session_id)
         session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session
-        return [message.model_dump() for message in session.get_message_list(conversation_key)]
+        return [message.model_dump() for message in session.message_lists.get(conversation_key, [])]
 
     def reset_session(
         self,
@@ -551,3 +551,23 @@ def reset_session(
             session.message_lists[conversation_key] = []
         if conversation_key in session.stream_message_indexes:
             session.stream_message_indexes[conversation_key] = {}
+
+        if session_id:
+            launcher_id = (
+                f'websocketgroup_{pipeline_uuid}:{session_id}'
+                if session_type == 'group'
+                else f'websocket_{pipeline_uuid}:{session_id}'
+            )
+            self.ap.sess_mgr.session_list = [
+                candidate_session
+                for candidate_session in self.ap.sess_mgr.session_list
+                if not (
+                    str(
+                        candidate_session.launcher_type.value
+                        if hasattr(candidate_session.launcher_type, 'value')
+                        else candidate_session.launcher_type
+                    )
+                    == session_type
+                    and str(candidate_session.launcher_id) == launcher_id
+                )
+            ]

src/langbot/templates/embed/widget.js

@@ -364,6 +364,9 @@
     isStreaming: false,
     streamingMsgId: null,
     historyLoaded: false,
+    hasConnected: false,
+    messageVersion: 0,
+    historyReloadTimer: null,
     pendingImage: null,
     feedbackState: {},
   };
@@ -571,6 +574,8 @@
       case "connected":
         state.isConnected = true;
         state.connectionId = data.connection_id;
+        if (state.hasConnected) loadHistory(true);
+        state.hasConnected = true;
         updateStatusDot();
         updateSendBtn();
         break;
@@ -638,6 +643,7 @@
 
     if (existingIdx >= 0) {
       state.messages[existingIdx] = msg;
+      state.messageVersion++;
       updateMessageEl(existingIdx, msg);
     } else {
       addMessage(msg);
@@ -707,9 +713,18 @@
   }
 
   // ========== Message History ==========
-  function loadHistory() {
-    if (state.historyLoaded) return;
+  function scheduleHistoryReload() {
+    if (state.historyReloadTimer) clearTimeout(state.historyReloadTimer);
+    state.historyReloadTimer = setTimeout(function () {
+      state.historyReloadTimer = null;
+      loadHistory(true);
+    }, 100);
+  }
+
+  function loadHistory(force) {
+    if (state.historyLoaded && !force) return;
     state.historyLoaded = true;
+    var messageVersion = state.messageVersion;
 
     var url =
       CONFIG.baseUrl +
@@ -728,6 +743,16 @@
       })
       .then(function (json) {
         if (json.code === 0 && json.data && json.data.messages) {
+          if (force && messageVersion !== state.messageVersion) {
+            scheduleHistoryReload();
+            return;
+          }
+          if (force) {
+            state.messages = [];
+            state.isStreaming = false;
+            state.streamingMsgId = null;
+            renderMessages();
+          }
           var msgs = json.data.messages;
           for (var i = 0; i < msgs.length; i++) {
             addMessage(msgs[i], true);
@@ -755,6 +780,7 @@
     fetch(url, { method: "POST", headers: headers })
       .then(function () {
         state.messages = [];
+        state.messageVersion++;
         state.isStreaming = false;
         state.streamingMsgId = null;
         state.historyLoaded = true;
@@ -768,6 +794,7 @@
   // ========== UI Rendering ==========
   function addMessage(msg, silent) {
     state.messages.push(msg);
+    if (!silent) state.messageVersion++;
     var el = createMessageEl(msg);
     if (els.welcome) {
       els.welcome.style.display = "none";

tests/unit_tests/platform/test_websocket_session_isolation.py

@@ -170,9 +170,36 @@ def test_session_ids_must_be_canonical_random_uuids():
     assert not is_valid_session_id('00000000-0000-0000-0000-000000000000')
 
 
+def test_history_read_does_not_allocate_unknown_session():
+    adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=AsyncMock())
+    adapter.websocket_person_session = WebSocketSession(id='person')
+    adapter.websocket_group_session = WebSocketSession(id='group')
+
+    assert adapter.get_websocket_messages('pipeline-1', 'person', 'missing-session') == []
+    assert adapter.websocket_person_session.message_lists == {}
+
+
 def test_history_and_reset_are_scoped_to_browser_session():
+    matching_provider_session = Mock(
+        launcher_type=Mock(value='person'),
+        launcher_id='websocket_pipeline-1:session-a',
+    )
+    matching_group_provider_session = Mock(
+        launcher_type=Mock(value='group'),
+        launcher_id='websocketgroup_pipeline-1:session-a',
+    )
+    other_session = Mock(
+        launcher_type=Mock(value='person'),
+        launcher_id='websocket_pipeline-1:session-b',
+    )
+    ap = Mock()
+    ap.sess_mgr.session_list = [
+        matching_provider_session,
+        matching_group_provider_session,
+        other_session,
+    ]
     adapter = WebSocketAdapter.model_construct(
-        ap=Mock(),
+        ap=ap,
         logger=AsyncMock(),
     )
     adapter.websocket_person_session = Mock()
@@ -207,6 +234,8 @@ def test_history_and_reset_are_scoped_to_browser_session():
     adapter.websocket_person_session.get_message_list.side_effect = histories.__getitem__
     adapter.websocket_person_session.message_lists = histories
     adapter.websocket_person_session.stream_message_indexes = stream_indexes
+    adapter.websocket_group_session.message_lists = {}
+    adapter.websocket_group_session.stream_message_indexes = {}
 
     assert adapter.get_websocket_messages('pipeline-1', 'person', 'session-a')[0]['content'] == 'private-a'
     assert adapter.get_websocket_messages('pipeline-1', 'person', 'session-b')[0]['content'] == 'private-b'
@@ -217,6 +246,11 @@ def test_history_and_reset_are_scoped_to_browser_session():
     assert stream_indexes['pipeline-1:session-a'] == {}
     assert histories['pipeline-1:session-b'] == session_b
     assert stream_indexes['pipeline-1:session-b'] == {'response-b': 0}
+    assert ap.sess_mgr.session_list == [matching_group_provider_session, other_session]
+
+    adapter.reset_session('pipeline-1', 'group', 'session-a')
+
+    assert ap.sess_mgr.session_list == [other_session]
 
 
 def test_widget_sends_stable_session_id_to_all_conversation_endpoints():
@@ -228,3 +262,6 @@ def test_widget_sends_stable_session_id_to_all_conversation_endpoints():
     assert 'window.localStorage' not in widget
     assert 'session_id=' in widget
     assert widget.count('encodeURIComponent(state.sessionId)') >= 3
+    assert 'loadHistory(true)' in widget
+    assert 'messageVersion !== state.messageVersion' in widget
+    assert 'scheduleHistoryReload();' in widget