mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-29 20:05:25 +08:00
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:
@@ -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'] = {}
|
||||
|
||||
Reference in New Issue
Block a user