From 76873c85e7a789d7f7447ca863095ab91de7db34 Mon Sep 17 00:00:00 2001 From: Smilewithoutfalling <91455578+Smilewithoutfalling@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:26:37 +0800 Subject: [PATCH 1/2] fix(wecombot): throttle intermediate stream frames in WS mode Drop intermediate stream frames that land inside the pacing window so a fast runner cannot flood the WeCom wire. The first frame of a stream is released immediately because there is no prior timestamp to compare against, and final frames always bypass the throttle. Frames carry complete snapshots since a2b7f6b7, so dropping an intermediate frame loses no content: the next frame already carries everything. --- .../libs/wecom_ai_bot_api/ws_client.py | 27 +++ .../test_wecombot_stream_empty_final.py | 9 +- .../platform/test_wecombot_stream_throttle.py | 168 ++++++++++++++++++ .../platform/test_wecombot_template_card.py | 8 +- 4 files changed, 210 insertions(+), 2 deletions(-) create mode 100644 tests/unit_tests/platform/test_wecombot_stream_throttle.py 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 9daf15900..c749a1c6f 100644 --- a/src/langbot/libs/wecom_ai_bot_api/ws_client.py +++ b/src/langbot/libs/wecom_ai_bot_api/ws_client.py @@ -64,6 +64,11 @@ _FEEDBACK_CACHE_MAX = 4096 _PENDING_FORM_MAX = 1024 _PENDING_FORM_TTL_SECONDS = 1800 _MAX_STREAM_CONTENT_CHARS = 200000 +# Minimum spacing between two intermediate stream frames. WeCom replaces the +# whole bubble on every frame and every frame must carry the complete snapshot, +# so forwarding each runner chunk re-uploads the growing answer again and +# again. Final chunks bypass the throttle, so the reply still ends promptly. +_STREAM_PUSH_MIN_INTERVAL = 0.3 _MAX_CALLBACK_TASKS = 100 _MAX_REPLY_WORKERS = 100 _MAX_REPLY_QUEUE_SIZE = 100 @@ -135,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 + # 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. # A runner may emit both a final delta and a completed message snapshot. self._completed_stream_content: dict[str, str] = {} @@ -173,6 +180,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_last_push_at.pop(msg_id, None) task_id = self._task_id_by_msg.pop(msg_id, None) if task_id: self._pending_forms_by_task.pop(task_id, None) @@ -253,6 +261,7 @@ class WecomBotWsClient: self._callback_tasks.clear() self._stream_ids.clear() self._stream_last_content.clear() + self._stream_last_push_at.clear() self._stream_sessions.clear() self._feedback_sessions.clear() self._msg_feedback_ids.clear() @@ -437,6 +446,7 @@ class WecomBotWsClient: self._stream_ids.pop(msg_id, None) self._stream_last_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 @@ -472,6 +482,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_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. @@ -713,6 +724,21 @@ class WecomBotWsClient: if not _re.sub(r'[\s\u200b\u200c\u200d\ufeff]', '', next_content): return True + # Throttle intermediate frames. Every frame carries the complete + # snapshot, so dropping one loses nothing; the next frame already + # contains everything that was skipped. The final chunk is never + # throttled, so a reply always terminates promptly. + if not is_final: + now = time.monotonic() + last_push_at = self._stream_last_push_at.get(msg_id) + # A missing entry means "nothing sent yet", not "sent at 0.0": + # time.monotonic() has an arbitrary origin, so on a freshly + # booted host a raw reading can be smaller than the interval + # and the first frame of a message would be dropped. + if last_push_at is not None and now - last_push_at < _STREAM_PUSH_MIN_INTERVAL: + return True + self._stream_last_push_at[msg_id] = now + # Generate feedback_id for final chunk feedback_id = '' if is_final: @@ -734,6 +760,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_last_push_at.pop(msg_id, None) self._stream_sessions.pop(msg_id, None) return True except Exception: diff --git a/tests/unit_tests/platform/test_wecombot_stream_empty_final.py b/tests/unit_tests/platform/test_wecombot_stream_empty_final.py index 8f18fa866..9a3b3c65e 100644 --- a/tests/unit_tests/platform/test_wecombot_stream_empty_final.py +++ b/tests/unit_tests/platform/test_wecombot_stream_empty_final.py @@ -17,6 +17,7 @@ These tests lock in the fixed behaviour: import pytest import langbot.pkg.core.app # noqa: F401 +from langbot.libs.wecom_ai_bot_api import ws_client from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient @@ -131,7 +132,13 @@ async def test_duplicate_final_is_acknowledged_without_sending_again(): @pytest.mark.asyncio -async def test_runner_cumulative_snapshots_reach_wire_once_without_concatenation(): +async def test_runner_cumulative_snapshots_reach_wire_once_without_concatenation(monkeypatch): + # Frame pacing has its own coverage in test_wecombot_stream_throttle.py. + # Disable it here so this test keeps asserting what it is named for: that + # each runner snapshot reaches the wire as a complete snapshot rather than + # being concatenated onto the previous one. + monkeypatch.setattr(ws_client, '_STREAM_PUSH_MIN_INTERVAL', 0.0) + from langbot.pkg.pipeline.process.stream_results import coalesce_stream_results from langbot_plugin.api.entities.builtin.provider.message import Message, MessageChunk from tests.unit_tests.platform.test_wecombot_eba_adapter import ( diff --git a/tests/unit_tests/platform/test_wecombot_stream_throttle.py b/tests/unit_tests/platform/test_wecombot_stream_throttle.py new file mode 100644 index 000000000..b5dda912c --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_stream_throttle.py @@ -0,0 +1,168 @@ +"""Regression tests for WeCom AI Bot stream frame throttling. + +A WeCom stream frame replaces the whole bubble and must carry the complete +snapshot, so forwarding every runner chunk re-uploads the growing answer over +and over on long replies. ``push_stream_chunk`` now spaces intermediate frames +out; these tests lock in that behaviour without touching a socket. + +Covered here: + +* the first frame of a message always goes out; +* intermediate frames inside the interval are dropped; +* a later frame goes out once the interval has elapsed; +* the final chunk is never throttled, so a reply always terminates; +* dropped frames lose no content, because the runner sends cumulative + snapshots and the closing frame still carries the whole answer; +* throttle state is tracked per message and released when the stream finishes. +""" + +import time + +import pytest + +import langbot.pkg.core.app # noqa: F401 +from langbot.libs.wecom_ai_bot_api import ws_client +from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient + + +class Logger: + def __init__(self): + self.warnings = [] + self.errors = [] + + async def warning(self, message): + self.warnings.append(message) + + async def error(self, message): + self.errors.append(message) + + async def info(self, message): + return None + + +class RecordingClient(WecomBotWsClient): + """A client that records reply frames instead of touching a socket.""" + + def __init__(self): + super().__init__(bot_id='bot', secret='secret', logger=Logger()) + self.replies = [] + + async def _send_reply(self, req_id, body, cmd='aibot_respond_msg'): + self.replies.append((cmd, req_id, body)) + return {'errcode': 0} + + def seed_stream(self, msg_id='msg-1', req_id='req-1', stream_id='stream-1'): + self._stream_ids[msg_id] = f'{req_id}|{stream_id}' + self._stream_sessions[msg_id] = { + 'req_id': req_id, + 'stream_id': stream_id, + 'msg_id': msg_id, + 'user_id': 'user-1', + 'chat_id': 'chat-1', + 'chat_type': 'single', + } + + def stream_frames(self): + return [body for _cmd, _req, body in self.replies if body.get('msgtype') == 'stream'] + + +@pytest.fixture +def frozen_throttle(monkeypatch): + """Widen the interval so every intermediate frame falls inside it.""" + monkeypatch.setattr(ws_client, '_STREAM_PUSH_MIN_INTERVAL', 3600.0) + + +@pytest.mark.asyncio +async def test_first_chunk_of_a_message_is_always_pushed(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + await client.push_stream_chunk('msg-1', 'hello', is_final=False) + + assert len(client.stream_frames()) == 1 + + +@pytest.mark.asyncio +async def test_intermediate_chunks_inside_the_interval_are_dropped(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + for i in range(1, 31): + await client.push_stream_chunk('msg-1', 'a' * i, is_final=False) + + assert len(client.stream_frames()) == 1 + + +@pytest.mark.asyncio +async def test_chunk_is_pushed_once_the_interval_has_elapsed(monkeypatch): + monkeypatch.setattr(ws_client, '_STREAM_PUSH_MIN_INTERVAL', 0.5) + client = RecordingClient() + client.seed_stream() + + await client.push_stream_chunk('msg-1', 'a', is_final=False) + # Inside the interval: dropped. Without the throttle this frame would go + # out too, so the frame count below is what separates the two behaviours. + await client.push_stream_chunk('msg-1', 'ab', is_final=False) + client._stream_last_push_at['msg-1'] = time.monotonic() - 10.0 + await client.push_stream_chunk('msg-1', 'abc', is_final=False) + + frames = client.stream_frames() + assert len(frames) == 2 + assert frames[-1]['stream']['content'] == 'abc' + + +@pytest.mark.asyncio +async def test_final_chunk_bypasses_the_throttle(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + await client.push_stream_chunk('msg-1', 'a', is_final=False) + # Inside the interval and therefore dropped; a throttle that also held the + # final frame back would leave the reply open, so the count matters. + await client.push_stream_chunk('msg-1', 'ab', is_final=False) + await client.push_stream_chunk('msg-1', 'abc', is_final=True) + + frames = client.stream_frames() + assert len(frames) == 2 + assert frames[-1]['stream']['content'] == 'abc' + assert frames[-1]['stream']['finish'] is True + + +@pytest.mark.asyncio +async def test_dropped_frames_do_not_lose_content(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + for i in range(1, 31): + await client.push_stream_chunk('msg-1', 'a' * i, is_final=False) + await client.push_stream_chunk('msg-1', 'a' * 30, is_final=True) + + frames = client.stream_frames() + assert len(frames) == 2 + assert frames[-1]['stream']['content'] == 'a' * 30 + + +@pytest.mark.asyncio +async def test_throttle_is_tracked_per_message(frozen_throttle): + client = RecordingClient() + client.seed_stream('msg-1', req_id='req-1', stream_id='stream-1') + client.seed_stream('msg-2', req_id='req-2', stream_id='stream-2') + + await client.push_stream_chunk('msg-1', 'a', is_final=False) + await client.push_stream_chunk('msg-2', 'b', is_final=False) + + frames = client.stream_frames() + assert len(frames) == 2 + assert 'msg-1' in client._stream_last_push_at + assert 'msg-2' in client._stream_last_push_at + + +@pytest.mark.asyncio +async def test_throttle_state_is_released_when_the_stream_finishes(frozen_throttle): + client = RecordingClient() + client.seed_stream() + + await client.push_stream_chunk('msg-1', 'a', is_final=False) + await client.push_stream_chunk('msg-1', 'ab', is_final=True) + + assert 'msg-1' not in client._stream_last_push_at diff --git a/tests/unit_tests/platform/test_wecombot_template_card.py b/tests/unit_tests/platform/test_wecombot_template_card.py index d27a6650c..410ad621a 100644 --- a/tests/unit_tests/platform/test_wecombot_template_card.py +++ b/tests/unit_tests/platform/test_wecombot_template_card.py @@ -22,6 +22,7 @@ from langbot.libs.wecom_ai_bot_api.api import ( # noqa: E402 build_human_input_text_prompt, parse_select_button_action, ) +from langbot.libs.wecom_ai_bot_api import ws_client # noqa: E402 from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient # noqa: E402 @@ -355,7 +356,12 @@ async def test_ws_push_form_pause_sends_text_prompt_without_empty_card(): @pytest.mark.asyncio -async def test_ws_stream_sends_cumulative_snapshots_to_wecom(): +async def test_ws_stream_sends_cumulative_snapshots_to_wecom(monkeypatch): + # Frame pacing has its own coverage in test_wecombot_stream_throttle.py. + # Disable it here so this test keeps asserting what it is named for: that + # every frame carries the cumulative snapshot rather than a delta. + monkeypatch.setattr(ws_client, '_STREAM_PUSH_MIN_INTERVAL', 0.0) + client = WecomBotWsClient('bot-id', 'secret', object()) client._stream_ids['msg-1'] = 'req-1|stream-1' client._stream_sessions['msg-1'] = {} From bab86430e66a2d3ec9d175f24699f5234adfc710 Mon Sep 17 00:00:00 2001 From: WangCham <651122857@qq.com> Date: Wed, 28 Oct 2026 20:46:15 +0800 Subject: [PATCH 2/2] fix(wecombot): preserve throttled snapshots for empty final frames --- .../libs/wecom_ai_bot_api/ws_client.py | 16 ++++++- .../platform/test_wecombot_stream_throttle.py | 46 +++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) 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'