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..23cbb8d47 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,10 @@ 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. # A runner may emit both a final delta and a completed message snapshot. self._completed_stream_content: dict[str, str] = {} @@ -173,6 +182,8 @@ 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: self._pending_forms_by_task.pop(task_id, None) @@ -253,6 +264,8 @@ 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() self._msg_feedback_ids.clear() @@ -437,6 +450,8 @@ 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 @@ -472,6 +487,8 @@ 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. @@ -680,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 @@ -713,6 +737,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 +773,8 @@ 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 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..30bbb5beb --- /dev/null +++ b/tests/unit_tests/platform/test_wecombot_stream_throttle.py @@ -0,0 +1,214 @@ +"""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 + 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' 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'] = {}