fix(streaming): preserve recovered segments across channels
This commit is contained in:
@@ -613,6 +613,23 @@ class AgentRunner:
|
||||
await hook.after_iteration(context)
|
||||
continue
|
||||
|
||||
# Some streaming providers recover with a complete response but no
|
||||
# content deltas. When an earlier length segment is already visible,
|
||||
# emit this terminal segment into the same stream; otherwise the
|
||||
# regular full response would duplicate the visible prefix.
|
||||
if (
|
||||
length_recovery_parts
|
||||
and hook.wants_streaming()
|
||||
and not context.streamed_content
|
||||
and response.finish_reason != "error"
|
||||
and not is_blank_text(clean)
|
||||
):
|
||||
await hook.on_stream(
|
||||
context,
|
||||
_restore_outer_whitespace(clean, original_content),
|
||||
)
|
||||
context.streamed_content = True
|
||||
|
||||
assistant_message: dict[str, Any] | None = None
|
||||
if response.finish_reason != "error" and not is_blank_text(clean):
|
||||
assistant_message = build_assistant_message(
|
||||
|
||||
@@ -497,6 +497,10 @@ class DiscordChannel(BaseChannel):
|
||||
self.logger.warning("client not ready; dropping stream delta")
|
||||
return
|
||||
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
buf = self._stream_bufs.get(chat_id)
|
||||
if not buf or buf.message is None or not buf.text:
|
||||
|
||||
@@ -754,6 +754,36 @@ async def test_send_delta_streams_by_editing_message(monkeypatch) -> None:
|
||||
assert owner._stream_bufs == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_merge_next_keeps_one_message(monkeypatch) -> None:
|
||||
owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus())
|
||||
client = _FakeDiscordClient(owner, intents=None)
|
||||
owner._client = client
|
||||
owner._running = True
|
||||
target = _FakeChannel(channel_id=123)
|
||||
client.channels[123] = target
|
||||
|
||||
times = iter([1.0, 3.0, 5.0])
|
||||
monkeypatch.setattr("nanobot.channels.discord.runtime.time.monotonic", lambda: next(times, 5.0))
|
||||
|
||||
await owner.send_delta(
|
||||
"123",
|
||||
"first-",
|
||||
stream_id="s1",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
await owner.send_delta("123", "second", stream_id="s1")
|
||||
await owner.send_delta("123", "", stream_id="s1", stream_end=True)
|
||||
|
||||
assert target.sent_payloads == [{"content": "first-"}]
|
||||
assert target.sent_messages[0].edits == [
|
||||
{"content": "first-second"},
|
||||
{"content": "first-second"},
|
||||
]
|
||||
assert owner._stream_bufs == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_stream_end_splits_oversized_reply(monkeypatch) -> None:
|
||||
owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus())
|
||||
|
||||
@@ -2232,6 +2232,10 @@ class FeishuChannel(BaseChannel):
|
||||
rid_type = "chat_id" if chat_id.startswith("oc_") else "open_id"
|
||||
|
||||
# --- stream end: final update or fallback ---
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
message_id = meta.get("message_id")
|
||||
# Only finalize the OnIt -> DONE reaction transition on the truly
|
||||
|
||||
@@ -285,6 +285,27 @@ class TestSendDelta:
|
||||
settings_call = ch._client.cardkit.v1.card.settings.call_args[0][0]
|
||||
assert settings_call.body.sequence == 5 # after final content seq 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_merge_next_preserves_buffer(self):
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="first-",
|
||||
card_id="card_1",
|
||||
sequence=3,
|
||||
last_edit=time.monotonic(),
|
||||
)
|
||||
|
||||
await ch.send_delta(
|
||||
"oc_chat1",
|
||||
"boundary",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert ch._stream_bufs["oc_chat1"].text == "first-boundary"
|
||||
ch._client.cardkit.v1.card_element.content.assert_not_called()
|
||||
ch._client.cardkit.v1.card.settings.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_fallback_when_no_card_id(self):
|
||||
"""If card creation failed, stream_end falls back to a plain card message."""
|
||||
|
||||
@@ -602,6 +602,10 @@ class MatrixChannel(BaseChannel):
|
||||
) -> None:
|
||||
relates_to = self._build_thread_relates_to(metadata)
|
||||
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
stream_key = _matrix_stream_key(chat_id, stream_id)
|
||||
buf = self._stream_bufs.pop(stream_key, None)
|
||||
|
||||
@@ -1937,6 +1937,29 @@ async def test_send_delta_stream_end_replaces_existing_message() -> None:
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_merge_next_preserves_buffer() -> None:
|
||||
channel = MatrixChannel(_make_config(), MessageBus())
|
||||
client = _FakeAsyncClient("", "", "", None)
|
||||
channel.client = client
|
||||
channel._stream_bufs["!room:matrix.org"] = matrix_module._StreamBuf(
|
||||
text="first-",
|
||||
event_id="event-1",
|
||||
last_edit=100.0,
|
||||
)
|
||||
channel.monotonic_time = lambda: 100.1
|
||||
|
||||
await channel.send_delta(
|
||||
"!room:matrix.org",
|
||||
"boundary",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert channel._stream_bufs["!room:matrix.org"].text == "first-boundary"
|
||||
assert client.room_send_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_keeps_same_room_stream_ids_independent(monkeypatch) -> None:
|
||||
channel = MatrixChannel(_make_config(), MessageBus())
|
||||
|
||||
@@ -931,6 +931,10 @@ class TelegramChannel(BaseChannel):
|
||||
meta = metadata or {}
|
||||
int_chat_id = int(chat_id)
|
||||
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
buf = self._stream_bufs.get(chat_id)
|
||||
if not buf or not buf.message_id or not buf.text:
|
||||
|
||||
@@ -675,6 +675,33 @@ async def test_send_delta_stream_end_raises_and_keeps_buffer_on_failure() -> Non
|
||||
assert "123" in channel._stream_bufs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_merge_next_preserves_buffer() -> None:
|
||||
channel = TelegramChannel(
|
||||
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"]),
|
||||
MessageBus(),
|
||||
)
|
||||
channel._app = _FakeApp(lambda: None)
|
||||
channel._app.bot.edit_message_text = AsyncMock()
|
||||
channel._stream_bufs["123"] = _StreamBuf(
|
||||
text="first-",
|
||||
message_id=7,
|
||||
last_edit=float("inf"),
|
||||
stream_id="s:0",
|
||||
)
|
||||
|
||||
await channel.send_delta(
|
||||
"123",
|
||||
"boundary",
|
||||
stream_id="s:0",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert channel._stream_bufs["123"].text == "first-boundary"
|
||||
channel._app.bot.edit_message_text.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_stream_end_treats_not_modified_as_success() -> None:
|
||||
from telegram.error import BadRequest
|
||||
|
||||
@@ -1257,6 +1257,10 @@ class WeixinChannel(BaseChannel):
|
||||
return
|
||||
is_end = stream_end or bool(meta.get("_stream_end"))
|
||||
buffer_key = stream_id or chat_id
|
||||
if is_end and merge_next:
|
||||
if delta:
|
||||
self._stream_buffers.setdefault(buffer_key, []).append(delta)
|
||||
return
|
||||
# Accumulate intermediate deltas. The stream_end message's own content
|
||||
# (present when the manager coalesces deltas into the end message) is
|
||||
# folded into `full` below instead of appended here, so a send retry
|
||||
|
||||
@@ -1824,6 +1824,29 @@ async def test_stream_end_flushes_buffered_answer() -> None:
|
||||
assert "wx-user" not in channel._stream_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_merge_next_preserves_buffer_until_final_end() -> None:
|
||||
channel, _bus = _make_channel()
|
||||
channel._client = object()
|
||||
channel._token = "token"
|
||||
channel._context_tokens["wx-user"] = "ctx-1"
|
||||
channel._context_token_at["wx-user"] = time.time()
|
||||
channel._send_text = AsyncMock()
|
||||
|
||||
await channel.send_delta(
|
||||
"wx-user",
|
||||
"first-",
|
||||
stream_id="s1",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
await channel.send_delta("wx-user", "second", stream_id="s1")
|
||||
await channel.send_delta("wx-user", "", stream_id="s1", stream_end=True)
|
||||
|
||||
channel._send_text.assert_awaited_once_with("wx-user", "first-second", "ctx-1")
|
||||
assert "s1" not in channel._stream_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_send_failure_keeps_buffer_for_retry() -> None:
|
||||
channel, _bus = _make_channel()
|
||||
|
||||
@@ -580,6 +580,53 @@ class TestToolEventProgress:
|
||||
assert [event.merge_next for event in endings] == [True, False]
|
||||
assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_length_recovery_streams_non_delta_terminal_segment(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.supports_progress_deltas = True
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
call_count = 0
|
||||
|
||||
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
await on_content_delta("first-")
|
||||
return LLMResponse(content="first-", finish_reason="length")
|
||||
return LLMResponse(content="second", finish_reason="stop")
|
||||
|
||||
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||
provider.chat_with_retry = AsyncMock()
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
_attach_webui_runtime_events(loop, bus)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop._dispatch(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u1",
|
||||
chat_id="chat1",
|
||||
content="give a long answer",
|
||||
metadata={"_wants_stream": True},
|
||||
))
|
||||
|
||||
outbound = []
|
||||
while bus.outbound_size > 0:
|
||||
outbound.append(await bus.consume_outbound())
|
||||
|
||||
deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)]
|
||||
endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)]
|
||||
final = [m for m in outbound if m.content == "first-second"]
|
||||
|
||||
assert [event.content for event in deltas] == ["first-", "second"]
|
||||
assert [event.merge_next for event in endings] == [True, False]
|
||||
assert len(final) == 1
|
||||
assert isinstance(final[0].event, StreamedResponseEvent)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_length_recovery_at_max_iterations_streams_only_missing_tail(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user