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