fix(webui): merge length recovery stream segments

This commit is contained in:
chengyongru
2026-07-27 01:39:46 +08:00
committed by Xubin Ren
parent 3cc5a98d9f
commit 27a00c7a4f
26 changed files with 325 additions and 29 deletions
+1
View File
@@ -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
View File
@@ -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
+7 -1
View File
@@ -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()
+1
View File
@@ -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,
+8 -1
View File
@@ -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
+2
View File
@@ -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(
+4
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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.
+21 -3
View File
@@ -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
+1
View File
@@ -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)
+1
View File
@@ -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
+1
View File
@@ -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:
+8 -1
View File
@@ -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()
+1
View File
@@ -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.
+4
View File
@@ -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
+46
View File
@@ -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,
+3
View File
@@ -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()
+9 -2
View File
@@ -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
+11 -5
View File
@@ -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),
]
+20
View File
@@ -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([
{
+19 -5
View File
@@ -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);
+2
View File
@@ -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";
+73
View File
@@ -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();