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.
This commit is contained in:
Smilewithoutfalling
2026-10-28 12:26:37 +08:00
parent f837083781
commit 76873c85e7
4 changed files with 210 additions and 2 deletions
@@ -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:
@@ -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 (
@@ -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
@@ -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'] = {}