mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-29 20:05:25 +08:00
Merge pull request #2564 from Smilewithoutfalling/fix/wecombot-throttle-stream-push
fix(wecombot): throttle intermediate stream frames in WS mode
This commit is contained in:
Guanchao Wang
@@ -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:
|
||||
|
||||
@@ -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,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'
|
||||
@@ -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