diff --git a/src/langbot/libs/wecom_ai_bot_api/ws_client.py b/src/langbot/libs/wecom_ai_bot_api/ws_client.py index c749a1c6f..23cbb8d47 100644 --- a/src/langbot/libs/wecom_ai_bot_api/ws_client.py +++ b/src/langbot/libs/wecom_ai_bot_api/ws_client.py @@ -140,6 +140,8 @@ class WecomBotWsClient: self._stream_ids: dict[str, str] = {} # msg_id -> req_id|stream_id # Dedup: skip sending when content hasn't changed self._stream_last_content: dict[str, str] = {} # msg_id -> last content sent + # Latest accepted snapshot, including frames skipped by the throttle. + self._stream_latest_content: dict[str, str] = {} # Throttle: msg_id -> monotonic timestamp of the last frame sent self._stream_last_push_at: dict[str, float] = {} # Preserve completion identity after removing active stream mappings. @@ -180,6 +182,7 @@ class WecomBotWsClient: self._stream_sessions.pop(msg_id, None) self._stream_ids.pop(msg_id, None) self._stream_last_content.pop(msg_id, None) + self._stream_latest_content.pop(msg_id, None) self._stream_last_push_at.pop(msg_id, None) task_id = self._task_id_by_msg.pop(msg_id, None) if task_id: @@ -261,6 +264,7 @@ class WecomBotWsClient: self._callback_tasks.clear() self._stream_ids.clear() self._stream_last_content.clear() + self._stream_latest_content.clear() self._stream_last_push_at.clear() self._stream_sessions.clear() self._feedback_sessions.clear() @@ -446,6 +450,7 @@ class WecomBotWsClient: self._stream_ids.pop(msg_id, None) self._stream_last_content.pop(msg_id, None) + self._stream_latest_content.pop(msg_id, None) self._stream_last_push_at.pop(msg_id, None) self._stream_sessions.pop(msg_id, None) return True, stream_id, None @@ -482,6 +487,7 @@ class WecomBotWsClient: # push_stream_chunk calls for this msg_id become no-ops. self._stream_ids.pop(msg_id, None) self._stream_last_content.pop(msg_id, None) + self._stream_latest_content.pop(msg_id, None) self._stream_last_push_at.pop(msg_id, None) # Keep _stream_sessions so the button callback can still resolve # user/chat context; it gets cleaned up when the click fires. @@ -691,10 +697,17 @@ class WecomBotWsClient: previous_content = self._stream_last_content.get(msg_id, '') # Pipeline delivery provides complete snapshots, including edits # that need not start with the previous text (e.g. reasoning tags). - next_content = content or previous_content + next_content = content or self._stream_latest_content.get(msg_id, previous_content) if len(next_content) > _MAX_STREAM_CONTENT_CHARS: next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:] + # Remember valid snapshots before deduplication: an edit can revert + # a throttled snapshot to the last sent text before an empty final. + import re as _re + + if _re.sub(r'[\s\u200b\u200c\u200d\ufeff]', '', next_content): + self._stream_latest_content[msg_id] = next_content + # Skip sending if content hasn't changed (e.g. during tool call argument streaming) if not is_final and next_content == previous_content: return True @@ -760,6 +773,7 @@ class WecomBotWsClient: self._cap_mapping(self._completed_stream_content, _STREAM_CACHE_MAX) self._stream_ids.pop(msg_id, None) self._stream_last_content.pop(msg_id, None) + self._stream_latest_content.pop(msg_id, None) self._stream_last_push_at.pop(msg_id, None) self._stream_sessions.pop(msg_id, None) return True diff --git a/tests/unit_tests/platform/test_wecombot_stream_throttle.py b/tests/unit_tests/platform/test_wecombot_stream_throttle.py index b5dda912c..30bbb5beb 100644 --- a/tests/unit_tests/platform/test_wecombot_stream_throttle.py +++ b/tests/unit_tests/platform/test_wecombot_stream_throttle.py @@ -166,3 +166,49 @@ async def test_throttle_state_is_released_when_the_stream_finishes(frozen_thrott await client.push_stream_chunk('msg-1', 'ab', is_final=True) assert 'msg-1' not in client._stream_last_push_at + assert 'msg-1' not in client._stream_latest_content + + +@pytest.mark.asyncio +@pytest.mark.parametrize('latest', ['hello world', 'rewritten answer']) +async def test_empty_final_flushes_latest_throttled_snapshot(frozen_throttle, latest): + client = RecordingClient() + client.seed_stream() + + await client.push_stream_chunk('msg-1', 'hello') + await client.push_stream_chunk('msg-1', latest) + assert len(client.stream_frames()) == 1 + await client.push_stream_chunk('msg-1', '', is_final=True) + + frames = client.stream_frames() + assert len(frames) == 2 + assert frames[-1]['stream']['content'] == latest + assert frames[-1]['stream']['finish'] is True + assert 'msg-1' not in client._stream_latest_content + + +@pytest.mark.asyncio +async def test_reverted_snapshot_replaces_throttled_content(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + for text in ['hello', 'hello world', 'hello']: + await client.push_stream_chunk('msg-1', text) + await client.push_stream_chunk('msg-1', '', is_final=True) + + assert len(client.stream_frames()) == 2 + assert client.stream_frames()[-1]['stream']['content'] == 'hello' + + +@pytest.mark.asyncio +async def test_repeated_snapshot_is_sent_after_throttle_expires(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + await client.push_stream_chunk('msg-1', 'hello') + await client.push_stream_chunk('msg-1', 'hello world') + client._stream_last_push_at['msg-1'] = time.monotonic() - 7200 + await client.push_stream_chunk('msg-1', 'hello world') + + assert len(client.stream_frames()) == 2 + assert client.stream_frames()[-1]['stream']['content'] == 'hello world'