fix(streaming): preserve recovered segments across channels
This commit is contained in:
@@ -613,6 +613,23 @@ class AgentRunner:
|
|||||||
await hook.after_iteration(context)
|
await hook.after_iteration(context)
|
||||||
continue
|
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
|
assistant_message: dict[str, Any] | None = None
|
||||||
if response.finish_reason != "error" and not is_blank_text(clean):
|
if response.finish_reason != "error" and not is_blank_text(clean):
|
||||||
assistant_message = build_assistant_message(
|
assistant_message = build_assistant_message(
|
||||||
|
|||||||
@@ -497,6 +497,10 @@ class DiscordChannel(BaseChannel):
|
|||||||
self.logger.warning("client not ready; dropping stream delta")
|
self.logger.warning("client not ready; dropping stream delta")
|
||||||
return
|
return
|
||||||
|
|
||||||
|
if stream_end and merge_next:
|
||||||
|
if not delta:
|
||||||
|
return
|
||||||
|
stream_end = False
|
||||||
if stream_end:
|
if stream_end:
|
||||||
buf = self._stream_bufs.get(chat_id)
|
buf = self._stream_bufs.get(chat_id)
|
||||||
if not buf or buf.message is None or not buf.text:
|
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 == {}
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_send_delta_stream_end_splits_oversized_reply(monkeypatch) -> None:
|
async def test_send_delta_stream_end_splits_oversized_reply(monkeypatch) -> None:
|
||||||
owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus())
|
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"
|
rid_type = "chat_id" if chat_id.startswith("oc_") else "open_id"
|
||||||
|
|
||||||
# --- stream end: final update or fallback ---
|
# --- stream end: final update or fallback ---
|
||||||
|
if stream_end and merge_next:
|
||||||
|
if not delta:
|
||||||
|
return
|
||||||
|
stream_end = False
|
||||||
if stream_end:
|
if stream_end:
|
||||||
message_id = meta.get("message_id")
|
message_id = meta.get("message_id")
|
||||||
# Only finalize the OnIt -> DONE reaction transition on the truly
|
# 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]
|
settings_call = ch._client.cardkit.v1.card.settings.call_args[0][0]
|
||||||
assert settings_call.body.sequence == 5 # after final content seq 4
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_stream_end_fallback_when_no_card_id(self):
|
async def test_stream_end_fallback_when_no_card_id(self):
|
||||||
"""If card creation failed, stream_end falls back to a plain card message."""
|
"""If card creation failed, stream_end falls back to a plain card message."""
|
||||||
|
|||||||
@@ -602,6 +602,10 @@ class MatrixChannel(BaseChannel):
|
|||||||
) -> None:
|
) -> None:
|
||||||
relates_to = self._build_thread_relates_to(metadata)
|
relates_to = self._build_thread_relates_to(metadata)
|
||||||
|
|
||||||
|
if stream_end and merge_next:
|
||||||
|
if not delta:
|
||||||
|
return
|
||||||
|
stream_end = False
|
||||||
if stream_end:
|
if stream_end:
|
||||||
stream_key = _matrix_stream_key(chat_id, stream_id)
|
stream_key = _matrix_stream_key(chat_id, stream_id)
|
||||||
buf = self._stream_bufs.pop(stream_key, None)
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_send_delta_keeps_same_room_stream_ids_independent(monkeypatch) -> None:
|
async def test_send_delta_keeps_same_room_stream_ids_independent(monkeypatch) -> None:
|
||||||
channel = MatrixChannel(_make_config(), MessageBus())
|
channel = MatrixChannel(_make_config(), MessageBus())
|
||||||
|
|||||||
@@ -931,6 +931,10 @@ class TelegramChannel(BaseChannel):
|
|||||||
meta = metadata or {}
|
meta = metadata or {}
|
||||||
int_chat_id = int(chat_id)
|
int_chat_id = int(chat_id)
|
||||||
|
|
||||||
|
if stream_end and merge_next:
|
||||||
|
if not delta:
|
||||||
|
return
|
||||||
|
stream_end = False
|
||||||
if stream_end:
|
if stream_end:
|
||||||
buf = self._stream_bufs.get(chat_id)
|
buf = self._stream_bufs.get(chat_id)
|
||||||
if not buf or not buf.message_id or not buf.text:
|
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
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_send_delta_stream_end_treats_not_modified_as_success() -> None:
|
async def test_send_delta_stream_end_treats_not_modified_as_success() -> None:
|
||||||
from telegram.error import BadRequest
|
from telegram.error import BadRequest
|
||||||
|
|||||||
@@ -1257,6 +1257,10 @@ class WeixinChannel(BaseChannel):
|
|||||||
return
|
return
|
||||||
is_end = stream_end or bool(meta.get("_stream_end"))
|
is_end = stream_end or bool(meta.get("_stream_end"))
|
||||||
buffer_key = stream_id or chat_id
|
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
|
# Accumulate intermediate deltas. The stream_end message's own content
|
||||||
# (present when the manager coalesces deltas into the end message) is
|
# (present when the manager coalesces deltas into the end message) is
|
||||||
# folded into `full` below instead of appended here, so a send retry
|
# 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
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_stream_end_send_failure_keeps_buffer_for_retry() -> None:
|
async def test_stream_end_send_failure_keeps_buffer_for_retry() -> None:
|
||||||
channel, _bus = _make_channel()
|
channel, _bus = _make_channel()
|
||||||
|
|||||||
@@ -580,6 +580,53 @@ class TestToolEventProgress:
|
|||||||
assert [event.merge_next for event in endings] == [True, False]
|
assert [event.merge_next for event in endings] == [True, False]
|
||||||
assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id}
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_length_recovery_at_max_iterations_streams_only_missing_tail(
|
async def test_length_recovery_at_max_iterations_streams_only_missing_tail(
|
||||||
self,
|
self,
|
||||||
|
|||||||
Reference in New Issue
Block a user