fix(webui): merge length recovery stream segments
This commit is contained in:
@@ -25,6 +25,7 @@ class AgentHookContext:
|
||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||
streamed_content: bool = False
|
||||
streamed_reasoning: bool = False
|
||||
stream_continues_current_message: bool = False
|
||||
final_content: str | None = None
|
||||
stop_reason: str | None = None
|
||||
error: str | None = None
|
||||
|
||||
+25
-4
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import inspect
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
@@ -860,9 +861,9 @@ class AgentLoop:
|
||||
"""Run the agent iteration loop.
|
||||
|
||||
*on_stream*: called with each content delta during streaming.
|
||||
*on_stream_end(resuming)*: called when a streaming session finishes.
|
||||
``resuming=True`` means tool calls follow (spinner should restart);
|
||||
``resuming=False`` means this is the final response.
|
||||
*on_stream_end(resuming, merge_next)*: called when a streaming session finishes.
|
||||
``resuming=True`` means the active turn continues. ``merge_next=True`` means
|
||||
the next text segment belongs to the same user-visible assistant message.
|
||||
|
||||
Returns (final_content, tools_used, messages, stop_reason, had_injections).
|
||||
"""
|
||||
@@ -1385,6 +1386,19 @@ class AgentLoop:
|
||||
if ctx.on_stream is not None:
|
||||
stream_callback = ctx.on_stream
|
||||
stream_end_callback = ctx.on_stream_end
|
||||
stream_end_accepts_merge_next = False
|
||||
if stream_end_callback is not None:
|
||||
try:
|
||||
stream_end_signature = inspect.signature(stream_end_callback)
|
||||
stream_end_accepts_merge_next = (
|
||||
"merge_next" in stream_end_signature.parameters
|
||||
or any(
|
||||
parameter.kind is inspect.Parameter.VAR_KEYWORD
|
||||
for parameter in stream_end_signature.parameters.values()
|
||||
)
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
segment_streamed_content = False
|
||||
|
||||
async def _tracked_stream(delta: str) -> None:
|
||||
@@ -1393,11 +1407,18 @@ class AgentLoop:
|
||||
segment_streamed_content = True
|
||||
await stream_callback(delta)
|
||||
|
||||
async def _tracked_stream_end(*, resuming: bool = False) -> None:
|
||||
async def _tracked_stream_end(
|
||||
*,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
nonlocal segment_streamed_content
|
||||
ctx.streamed_content = segment_streamed_content
|
||||
segment_streamed_content = False
|
||||
if stream_end_callback is not None:
|
||||
if merge_next and stream_end_accepts_merge_next:
|
||||
await stream_end_callback(resuming=resuming, merge_next=True)
|
||||
else:
|
||||
await stream_end_callback(resuming=resuming)
|
||||
|
||||
ctx.on_stream = _tracked_stream
|
||||
|
||||
@@ -85,7 +85,13 @@ class AgentProgressHook(AgentHook):
|
||||
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
|
||||
await self.emit_reasoning_end()
|
||||
if self._on_stream_end:
|
||||
await self._on_stream_end(resuming=resuming)
|
||||
kwargs: dict[str, bool] = {"resuming": resuming}
|
||||
if (
|
||||
context.stream_continues_current_message
|
||||
and self._on_progress_accepts(self._on_stream_end, "merge_next")
|
||||
):
|
||||
kwargs["merge_next"] = True
|
||||
await self._on_stream_end(**kwargs)
|
||||
self._stream_buf = ""
|
||||
self._think_extractor.reset()
|
||||
|
||||
|
||||
@@ -599,6 +599,7 @@ class AgentRunner:
|
||||
_MAX_LENGTH_RECOVERIES,
|
||||
)
|
||||
if hook.wants_streaming():
|
||||
context.stream_continues_current_message = True
|
||||
await hook.on_stream_end(context, resuming=True)
|
||||
messages.append(build_assistant_message(
|
||||
clean,
|
||||
|
||||
@@ -285,7 +285,12 @@ class TurnDelivery:
|
||||
)
|
||||
)
|
||||
|
||||
async def _publish_stream_end(self, *, resuming: bool = False) -> None:
|
||||
async def _publish_stream_end(
|
||||
self,
|
||||
*,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
await self.bus.publish_outbound(
|
||||
outbound_message_for_event(
|
||||
channel=self.delivery_message.channel,
|
||||
@@ -293,8 +298,10 @@ class TurnDelivery:
|
||||
event=StreamEndEvent(
|
||||
stream_id=self._stream_id(),
|
||||
resuming=resuming,
|
||||
merge_next=merge_next,
|
||||
),
|
||||
metadata=self.delivery_message.metadata,
|
||||
)
|
||||
)
|
||||
if not merge_next:
|
||||
self._stream_segment += 1
|
||||
|
||||
@@ -46,6 +46,7 @@ class StreamEndEvent(OutboundEvent):
|
||||
content: str = ""
|
||||
stream_id: str | None = None
|
||||
resuming: bool = False
|
||||
merge_next: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -176,6 +177,7 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None:
|
||||
content=msg.content,
|
||||
stream_id=_metadata_str(meta, "_stream_id"),
|
||||
resuming=bool(meta.get("_resuming")),
|
||||
merge_next=bool(meta.get("_merge_next")),
|
||||
)
|
||||
if meta.get("_stream_delta"):
|
||||
return StreamDeltaEvent(
|
||||
|
||||
@@ -110,6 +110,7 @@ class BaseChannel(ABC):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Deliver a streaming text chunk.
|
||||
|
||||
@@ -118,6 +119,9 @@ class BaseChannel(ABC):
|
||||
|
||||
Stateful implementations should key buffers by ``stream_id`` rather
|
||||
than only by ``chat_id`` when it is provided.
|
||||
|
||||
``merge_next`` marks a resumable provider boundary whose next text
|
||||
segment belongs to the same user-visible message.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@@ -489,6 +489,7 @@ class DiscordChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Progressive Discord delivery: send once, then edit until the stream ends."""
|
||||
client = self._client
|
||||
|
||||
@@ -2216,6 +2216,7 @@ class FeishuChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Progressive streaming via CardKit: create card on first delta, stream-update on subsequent.
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import inspect
|
||||
from collections.abc import Callable, Iterable
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
@@ -763,13 +764,29 @@ class ChannelManager:
|
||||
msg: OutboundMessage,
|
||||
event: StreamDeltaEvent | StreamEndEvent,
|
||||
) -> None:
|
||||
kwargs: dict[str, Any] = {
|
||||
"stream_id": event.stream_id,
|
||||
"stream_end": isinstance(event, StreamEndEvent),
|
||||
"resuming": event.resuming if isinstance(event, StreamEndEvent) else False,
|
||||
}
|
||||
if isinstance(event, StreamEndEvent) and event.merge_next:
|
||||
try:
|
||||
signature = inspect.signature(channel.send_delta)
|
||||
if (
|
||||
"merge_next" in signature.parameters
|
||||
or any(
|
||||
parameter.kind is inspect.Parameter.VAR_KEYWORD
|
||||
for parameter in signature.parameters.values()
|
||||
)
|
||||
):
|
||||
kwargs["merge_next"] = True
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
await channel.send_delta(
|
||||
msg.chat_id,
|
||||
msg.content,
|
||||
msg.metadata,
|
||||
stream_id=event.stream_id,
|
||||
stream_end=isinstance(event, StreamEndEvent),
|
||||
resuming=event.resuming if isinstance(event, StreamEndEvent) else False,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@@ -850,6 +867,7 @@ class ChannelManager:
|
||||
final_event = StreamEndEvent(
|
||||
stream_id=next_stream_id,
|
||||
resuming=next_event.resuming,
|
||||
merge_next=next_event.merge_next,
|
||||
)
|
||||
# Stream ended - stop coalescing this stream
|
||||
break
|
||||
|
||||
@@ -598,6 +598,7 @@ class MatrixChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
relates_to = self._build_thread_relates_to(metadata)
|
||||
|
||||
|
||||
@@ -515,6 +515,7 @@ class MattermostChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
if not self._http_client:
|
||||
return
|
||||
|
||||
@@ -923,6 +923,7 @@ class TelegramChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Progressive message editing: send on first delta, edit on subsequent ones."""
|
||||
if not self._app:
|
||||
|
||||
@@ -995,13 +995,18 @@ class WebSocketChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
conns = list(self._subs.get(chat_id, ()))
|
||||
meta = metadata or {}
|
||||
stream_key = (chat_id, str(stream_id or ""))
|
||||
if stream_end:
|
||||
body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id}
|
||||
buffered = self._stream_text_buffers.pop(stream_key, [])
|
||||
buffered = (
|
||||
self._stream_text_buffers.setdefault(stream_key, [])
|
||||
if merge_next
|
||||
else self._stream_text_buffers.pop(stream_key, [])
|
||||
)
|
||||
if delta:
|
||||
buffered.append(delta)
|
||||
full_text = "".join(buffered)
|
||||
@@ -1019,6 +1024,8 @@ class WebSocketChannel(BaseChannel):
|
||||
body["stream_id"] = stream_id
|
||||
if stream_end and resuming:
|
||||
body["resuming"] = True
|
||||
if stream_end and merge_next:
|
||||
body["merge_next"] = True
|
||||
self._transcripts.prepare_and_append(
|
||||
chat_id,
|
||||
body,
|
||||
|
||||
@@ -1350,6 +1350,39 @@ async def test_send_delta_marks_resuming_stream_end() -> None:
|
||||
assert payload["resuming"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_keeps_buffer_across_merged_stream_boundary() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel(
|
||||
{"enabled": True, "allowFrom": ["*"], "streaming": True},
|
||||
bus,
|
||||
gateway=_basic_handler(bus),
|
||||
)
|
||||
mock_ws = AsyncMock()
|
||||
channel._attach(mock_ws, "chat-1")
|
||||
|
||||
await channel.send_delta("chat-1", "first ", stream_id="sid")
|
||||
await channel.send_delta(
|
||||
"chat-1",
|
||||
"",
|
||||
stream_id="sid",
|
||||
stream_end=True,
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
)
|
||||
await channel.send_delta("chat-1", "second", stream_id="sid")
|
||||
await channel.send_delta("chat-1", "", stream_id="sid", stream_end=True)
|
||||
|
||||
payloads = [json.loads(call.args[0]) for call in mock_ws.send.await_args_list]
|
||||
assert payloads[1]["merge_next"] is True
|
||||
assert payloads[1]["resuming"] is True
|
||||
assert [payload["text"] for payload in payloads if payload["event"] == "delta"] == [
|
||||
"first ",
|
||||
"second",
|
||||
]
|
||||
assert ("chat-1", "sid") not in channel._stream_text_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_stream_end_includes_inline_final_text() -> None:
|
||||
bus = MagicMock()
|
||||
|
||||
@@ -1243,6 +1243,7 @@ class WeixinChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Deliver a streamed reply to WeChat.
|
||||
|
||||
|
||||
@@ -1770,6 +1770,7 @@ def replay_transcript_to_ui_messages(
|
||||
buffer_message_id = None
|
||||
buffer_parts = []
|
||||
continue
|
||||
merge_next = rec.get("resuming") is True and rec.get("merge_next") is True
|
||||
final_text = rec.get("text")
|
||||
if isinstance(final_text, str):
|
||||
if buffer_message_id is None:
|
||||
@@ -1794,6 +1795,9 @@ def replay_transcript_to_ui_messages(
|
||||
**_turn_fields(rec, "answer"),
|
||||
}
|
||||
break
|
||||
if merge_next:
|
||||
buffer_parts = [final_text]
|
||||
if not merge_next:
|
||||
buffer_message_id = None
|
||||
buffer_parts = []
|
||||
continue
|
||||
|
||||
@@ -534,6 +534,52 @@ class TestToolEventProgress:
|
||||
assert turn_end_msgs[0].content == ""
|
||||
provider.chat_with_retry.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_length_recovery_keeps_one_user_visible_stream(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.supports_progress_deltas = True
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
responses = iter([
|
||||
LLMResponse(content="first-", finish_reason="length"),
|
||||
LLMResponse(content="second", finish_reason="stop"),
|
||||
])
|
||||
|
||||
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
|
||||
response = next(responses)
|
||||
await on_content_delta(response.content or "")
|
||||
return response
|
||||
|
||||
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)]
|
||||
|
||||
assert [event.content for event in deltas] == ["first-", "second"]
|
||||
assert [event.resuming 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}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streamed_finalization_is_delivered_as_regular_message(
|
||||
self,
|
||||
|
||||
@@ -151,6 +151,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
streamed: list[str] = []
|
||||
endings: list[bool] = []
|
||||
merge_next: list[bool] = []
|
||||
responses = iter([
|
||||
LLMResponse(content="first ", finish_reason="length"),
|
||||
LLMResponse(content="second", finish_reason="stop"),
|
||||
@@ -175,6 +176,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont
|
||||
|
||||
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
|
||||
endings.append(resuming)
|
||||
merge_next.append(context.stream_continues_current_message)
|
||||
|
||||
runner = AgentRunner()
|
||||
result = await runner.run(make_run_spec(provider,
|
||||
@@ -189,6 +191,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont
|
||||
assert result.final_content == "first second"
|
||||
assert streamed == ["first ", "second"]
|
||||
assert endings == [True, False]
|
||||
assert merge_next == [True, False]
|
||||
provider.chat_with_retry.assert_not_awaited()
|
||||
|
||||
|
||||
|
||||
@@ -94,7 +94,12 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None:
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
content="",
|
||||
metadata={"_stream_end": True, "_stream_id": "s1", "_resuming": True},
|
||||
metadata={
|
||||
"_stream_end": True,
|
||||
"_stream_id": "s1",
|
||||
"_resuming": True,
|
||||
"_merge_next": True,
|
||||
},
|
||||
)
|
||||
|
||||
delta_event = outbound_event_from_message(delta)
|
||||
@@ -106,6 +111,7 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None:
|
||||
assert isinstance(end_event, StreamEndEvent)
|
||||
assert end_event.stream_id == "s1"
|
||||
assert end_event.resuming is True
|
||||
assert end_event.merge_next is True
|
||||
|
||||
|
||||
def test_legacy_webui_runtime_metadata_flags_create_runtime_events() -> None:
|
||||
@@ -221,7 +227,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None:
|
||||
|
||||
updated = replace_outbound_event(
|
||||
msg,
|
||||
StreamEndEvent(stream_id="s1", resuming=True),
|
||||
StreamEndEvent(stream_id="s1", resuming=True, merge_next=True),
|
||||
content="hello world",
|
||||
)
|
||||
|
||||
@@ -230,6 +236,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None:
|
||||
assert isinstance(updated.event, StreamEndEvent)
|
||||
assert updated.event.stream_id == "s1"
|
||||
assert updated.event.resuming is True
|
||||
assert updated.event.merge_next is True
|
||||
|
||||
|
||||
def test_streamed_response_event_keeps_final_content_outside_event_payload() -> None:
|
||||
|
||||
@@ -49,6 +49,7 @@ class MockChannel(BaseChannel):
|
||||
stream_id=None,
|
||||
stream_end=False,
|
||||
resuming=False,
|
||||
merge_next=False,
|
||||
):
|
||||
return await self._send_delta_mock(
|
||||
chat_id,
|
||||
@@ -57,6 +58,7 @@ class MockChannel(BaseChannel):
|
||||
stream_id=stream_id,
|
||||
stream_end=stream_end,
|
||||
resuming=resuming,
|
||||
merge_next=merge_next,
|
||||
)
|
||||
|
||||
|
||||
@@ -92,11 +94,17 @@ def _end(
|
||||
chat_id: str = "chat1",
|
||||
stream_id: str | None = None,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
):
|
||||
return outbound_message_for_event(
|
||||
channel="mock",
|
||||
chat_id=chat_id,
|
||||
event=StreamEndEvent(content=content, stream_id=stream_id, resuming=resuming),
|
||||
event=StreamEndEvent(
|
||||
content=content,
|
||||
stream_id=stream_id,
|
||||
resuming=resuming,
|
||||
merge_next=merge_next,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -137,6 +145,7 @@ class TestDeltaCoalescing:
|
||||
stream_id=None,
|
||||
stream_end=False,
|
||||
resuming=False,
|
||||
merge_next=False,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -184,13 +193,19 @@ class TestDeltaCoalescing:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_terminates_coalescing(self, manager, bus):
|
||||
await bus.publish_outbound(_delta("Hello"))
|
||||
await bus.publish_outbound(_end(" world"))
|
||||
await bus.publish_outbound(_end(
|
||||
" world",
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
))
|
||||
|
||||
first_msg = await bus.consume_outbound()
|
||||
merged, pending = manager._coalesce_stream_deltas(first_msg)
|
||||
|
||||
assert merged.content == "Hello world"
|
||||
assert isinstance(merged.event, StreamEndEvent)
|
||||
assert merged.event.resuming is True
|
||||
assert merged.event.merge_next is True
|
||||
assert len(pending) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -2818,7 +2818,7 @@ async def test_send_with_retry_no_retry_when_max_is_zero():
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_with_retry_calls_send_delta():
|
||||
"""_send_with_retry should call send_delta for stream delta events."""
|
||||
calls: list[tuple[str, str, str | None, bool, bool]] = []
|
||||
calls: list[tuple[str, str, str | None, bool, bool, bool]] = []
|
||||
|
||||
class _StreamingChannel(BaseChannel):
|
||||
name = "streaming"
|
||||
@@ -2842,8 +2842,9 @@ async def test_send_with_retry_calls_send_delta():
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
calls.append((chat_id, delta, stream_id, stream_end, resuming))
|
||||
calls.append((chat_id, delta, stream_id, stream_end, resuming, merge_next))
|
||||
|
||||
fake_config = SimpleNamespace(
|
||||
channels=ChannelsConfig(send_max_retries=3),
|
||||
@@ -2865,13 +2866,18 @@ async def test_send_with_retry_calls_send_delta():
|
||||
end = outbound_message_for_event(
|
||||
channel="streaming",
|
||||
chat_id="123",
|
||||
event=StreamEndEvent(content="", stream_id="s1", resuming=True),
|
||||
event=StreamEndEvent(
|
||||
content="",
|
||||
stream_id="s1",
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
),
|
||||
)
|
||||
await mgr._send_with_retry(mgr.channels["streaming"], end)
|
||||
|
||||
assert calls == [
|
||||
("123", "test delta", "s1", False, False),
|
||||
("123", "", "s1", True, True),
|
||||
("123", "test delta", "s1", False, False, False),
|
||||
("123", "", "s1", True, True, True),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -1121,6 +1121,26 @@ def test_replay_keeps_interrupted_pre_tool_text_in_activity() -> None:
|
||||
assert msgs[2]["content"] == "Done. Open index.html to play."
|
||||
|
||||
|
||||
def test_replay_merges_length_recovery_segments_into_one_assistant_message() -> None:
|
||||
msgs = replay_transcript_to_ui_messages([
|
||||
{"event": "delta", "chat_id": "t-stream", "text": "first "},
|
||||
{
|
||||
"event": "stream_end",
|
||||
"chat_id": "t-stream",
|
||||
"text": "first ",
|
||||
"resuming": True,
|
||||
"merge_next": True,
|
||||
},
|
||||
{"event": "delta", "chat_id": "t-stream", "text": "second"},
|
||||
{"event": "stream_end", "chat_id": "t-stream"},
|
||||
{"event": "turn_end", "chat_id": "t-stream"},
|
||||
])
|
||||
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["role"] == "assistant"
|
||||
assert msgs[0]["content"] == "first second"
|
||||
|
||||
|
||||
def test_replay_tool_events_dedupes_finish_after_start() -> None:
|
||||
msgs = replay_transcript_to_ui_messages([
|
||||
{
|
||||
|
||||
@@ -26,7 +26,7 @@ import type {
|
||||
} from "@/lib/types";
|
||||
|
||||
interface StreamBuffer {
|
||||
/** ID of the assistant message currently receiving deltas (cleared on ``stream_end``). */
|
||||
/** ID of the assistant message currently receiving deltas (cleared when its segment closes). */
|
||||
messageId: string;
|
||||
}
|
||||
|
||||
@@ -780,15 +780,20 @@ export function useNanobotStream(
|
||||
?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current, turn);
|
||||
if (targetIndex !== null) {
|
||||
const target = next[targetIndex];
|
||||
next = replaceMessageAt(next, targetIndex, {
|
||||
const merged = {
|
||||
...target,
|
||||
content: finalAnswerText,
|
||||
isStreaming: true,
|
||||
...turn,
|
||||
});
|
||||
};
|
||||
next = replaceMessageAt(next, targetIndex, merged);
|
||||
if (!options?.closeAnswerSegment) {
|
||||
closedAssistantStreamIdsRef.current.delete(merged.id);
|
||||
activeAssistantRef.current = { id: merged.id, index: targetIndex };
|
||||
buffer.current = { messageId: merged.id };
|
||||
}
|
||||
} else {
|
||||
const id = crypto.randomUUID();
|
||||
closedAssistantStreamIdsRef.current.add(id);
|
||||
next = [
|
||||
...next,
|
||||
{
|
||||
@@ -800,6 +805,12 @@ export function useNanobotStream(
|
||||
createdAt: Date.now(),
|
||||
},
|
||||
];
|
||||
if (options?.closeAnswerSegment) {
|
||||
closedAssistantStreamIdsRef.current.add(id);
|
||||
} else {
|
||||
activeAssistantRef.current = { id, index: next.length - 1 };
|
||||
buffer.current = { messageId: id };
|
||||
}
|
||||
}
|
||||
}
|
||||
if (options?.closeAnswerSegment) closeActiveAssistantStream();
|
||||
@@ -911,8 +922,9 @@ export function useNanobotStream(
|
||||
|
||||
if (ev.event === "stream_end") {
|
||||
const turn = turnFieldsFromEvent(ev, "answer");
|
||||
const mergeNext = ev.resuming === true && ev.merge_next === true;
|
||||
flushPendingStreamEvents({
|
||||
closeAnswerSegment: true,
|
||||
closeAnswerSegment: !mergeNext,
|
||||
...(typeof ev.text === "string" ? { finalAnswerText: ev.text } : {}),
|
||||
turn,
|
||||
});
|
||||
@@ -920,7 +932,9 @@ export function useNanobotStream(
|
||||
if (ev.resuming) {
|
||||
cancelStreamEndTimer();
|
||||
setIsStreaming(true);
|
||||
if (!mergeNext) {
|
||||
setMessages((prev) => finalizeStreamedTurn(prev, turn));
|
||||
}
|
||||
return;
|
||||
}
|
||||
scheduleStreamEndTimer(turn);
|
||||
|
||||
@@ -1118,6 +1118,8 @@ export type InboundEvent =
|
||||
text?: string;
|
||||
/** This answer segment ended, but the active agent turn will continue. */
|
||||
resuming?: boolean;
|
||||
/** The next answer segment continues this same assistant message. */
|
||||
merge_next?: boolean;
|
||||
} & InboundTurnMetadata)
|
||||
| ({
|
||||
event: "reasoning_delta";
|
||||
|
||||
@@ -2009,6 +2009,79 @@ describe("useNanobotStream", () => {
|
||||
expect(result.current.messages.every((message) => !message.isStreaming)).toBe(true);
|
||||
});
|
||||
|
||||
it("keeps length-recovery segments in one assistant message", async () => {
|
||||
const fake = fakeClient();
|
||||
const { result } = renderHook(() => useNanobotStream("chat-length", EMPTY_MESSAGES), {
|
||||
wrapper: wrap(fake.client),
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.send("give a long answer");
|
||||
});
|
||||
const activeTurnId = fake.client.sendMessage.mock.calls.at(-1)![3]?.turnId;
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "delta",
|
||||
chat_id: "chat-length",
|
||||
text: "first ",
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
await flushStreamFrame();
|
||||
const assistantId = result.current.messages[1].id;
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "stream_end",
|
||||
chat_id: "chat-length",
|
||||
text: "first ",
|
||||
resuming: true,
|
||||
merge_next: true,
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
|
||||
expect(result.current.messages).toHaveLength(2);
|
||||
expect(result.current.messages[1]).toMatchObject({
|
||||
id: assistantId,
|
||||
content: "first ",
|
||||
isStreaming: true,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "delta",
|
||||
chat_id: "chat-length",
|
||||
text: "second",
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
await flushStreamFrame();
|
||||
|
||||
expect(result.current.messages).toHaveLength(2);
|
||||
expect(result.current.messages[1]).toMatchObject({
|
||||
id: assistantId,
|
||||
content: "first second",
|
||||
isStreaming: true,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "turn_end",
|
||||
chat_id: "chat-length",
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
|
||||
expect(result.current.messages).toHaveLength(2);
|
||||
expect(result.current.messages[1]).toMatchObject({
|
||||
id: assistantId,
|
||||
content: "first second",
|
||||
isStreaming: false,
|
||||
});
|
||||
});
|
||||
|
||||
it("keeps streaming alive across stream_end when tool activity follows", async () => {
|
||||
const fake = fakeClient();
|
||||
const onTurnEnd = vi.fn();
|
||||
|
||||
Reference in New Issue
Block a user