From bb2f6cf3241324cad2c63505bb845ea85d5323e8 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Tue, 7 Jul 2026 10:17:23 +0800 Subject: [PATCH] fix(webui): preserve automation source on streamed replies --- nanobot/channels/websocket/runtime.py | 1 + .../websocket/tests/test_websocket_channel.py | 30 ++++++++++ nanobot/webui/transcript.py | 24 ++++++-- tests/utils/test_webui_transcript.py | 36 ++++++++++++ webui/src/hooks/useNanobotStream.ts | 32 +++++++++-- webui/src/lib/types.ts | 4 ++ webui/src/tests/useNanobotStream.test.tsx | 57 +++++++++++++++++++ 7 files changed, 176 insertions(+), 8 deletions(-) diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index 7d5f03e3f..1e0f39904 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -1216,6 +1216,7 @@ class WebSocketChannel(BaseChannel): body, metadata=meta, phase="answer", + include_source=True, ) raw = json.dumps(body, ensure_ascii=False) if not conns: diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index c23ebd088..1e44b1e65 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -51,6 +51,7 @@ from nanobot.webui.http_utils import ( ) from nanobot.webui.metadata import ( WEBSOCKET_TURN_OWNER_METADATA_KEY, + WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_SYSTEM_COMMAND_TURN_PREFIX, WEBUI_TURN_METADATA_KEY, ) @@ -1350,6 +1351,35 @@ async def test_send_delta_emits_delta_and_stream_end() -> None: assert "text" not in second +@pytest.mark.asyncio +async def test_send_delta_preserves_webui_source_metadata() -> None: + bus = MagicMock() + channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus, gateway=_basic_handler(bus)) + mock_ws = AsyncMock() + channel._attach(mock_ws, "chat-source-stream") + source = {"kind": "cron", "label": "Repo check"} + metadata = {WEBUI_MESSAGE_SOURCE_METADATA_KEY: source} + + await channel.send_delta("chat-source-stream", "done", metadata=metadata, stream_id="sid") + await channel.send_delta( + "chat-source-stream", + "", + metadata=metadata, + stream_id="sid", + stream_end=True, + ) + + first = json.loads(mock_ws.send.call_args_list[0][0][0]) + second = json.loads(mock_ws.send.call_args_list[1][0][0]) + assert first["event"] == "delta" + assert first["source"] == source + assert second["event"] == "stream_end" + assert second["source"] == source + lines = read_transcript_lines("websocket:chat-source-stream") + assert lines[-2]["source"] == source + assert lines[-1]["source"] == source + + @pytest.mark.asyncio async def test_send_delta_marks_resuming_stream_end() -> None: bus = MagicMock() diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index b594f1b68..f9c43d415 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -2026,6 +2026,7 @@ def replay_transcript_to_ui_messages( continue close_activity_for_answer() turn_fields = _turn_fields(rec, "answer") + source_fields = _source_fields(rec) adopted = find_active_placeholder(messages, turn_fields) if buffer_message_id is None else None if buffer_message_id is None: if adopted: @@ -2038,7 +2039,8 @@ def replay_transcript_to_ui_messages( "role": "assistant", "content": "", "isStreaming": True, - **_turn_fields(rec, "answer"), + **turn_fields, + **source_fields, "createdAt": _created_at_ms(rec, idx), }, ) @@ -2050,7 +2052,8 @@ def replay_transcript_to_ui_messages( **m, "content": combined, "isStreaming": True, - **_turn_fields(rec, "answer"), + **turn_fields, + **source_fields, } break continue @@ -2062,6 +2065,8 @@ def replay_transcript_to_ui_messages( continue merge_next = rec.get("resuming") is True and rec.get("merge_next") is True final_text = rec.get("text") + turn_fields = _turn_fields(rec, "answer") + source_fields = _source_fields(rec) if isinstance(final_text, str): if buffer_message_id is None: buffer_message_id = _new_id("buf", idx) @@ -2071,7 +2076,8 @@ def replay_transcript_to_ui_messages( "role": "assistant", "content": final_text, "isStreaming": True, - **_turn_fields(rec, "answer"), + **turn_fields, + **source_fields, "createdAt": _created_at_ms(rec, idx), }, ) @@ -2082,11 +2088,21 @@ def replay_transcript_to_ui_messages( **m, "content": final_text, "isStreaming": True, - **_turn_fields(rec, "answer"), + **turn_fields, + **source_fields, } break if merge_next: buffer_parts = [final_text] + elif source_fields and buffer_message_id is not None: + for i, m in enumerate(messages): + if m.get("id") == buffer_message_id: + messages[i] = { + **m, + **turn_fields, + **source_fields, + } + break if not merge_next: buffer_message_id = None buffer_parts = [] diff --git a/tests/utils/test_webui_transcript.py b/tests/utils/test_webui_transcript.py index c31c9705a..1fe3fa6e1 100644 --- a/tests/utils/test_webui_transcript.py +++ b/tests/utils/test_webui_transcript.py @@ -696,6 +696,42 @@ def test_replay_preserves_local_trigger_source_metadata(tmp_path, monkeypatch) - assert msgs[0]["source"] == {"kind": "local_trigger", "label": "PR review"} +def test_replay_preserves_automation_source_metadata_on_streamed_reply( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path) + key = "websocket:t-streamed-cron-source" + source = {"kind": "cron", "label": "Repo check"} + + for record in ( + { + "event": "delta", + "chat_id": "t-streamed-cron-source", + "text": "Repo ", + "source": source, + }, + { + "event": "delta", + "chat_id": "t-streamed-cron-source", + "text": "clean.", + "source": source, + }, + { + "event": "stream_end", + "chat_id": "t-streamed-cron-source", + "source": source, + }, + {"event": "turn_end", "chat_id": "t-streamed-cron-source"}, + ): + append_transcript_object(key, record) + + msgs = replay_transcript_to_ui_messages(read_transcript_lines(key)) + + assert msgs[0]["content"] == "Repo clean." + assert msgs[0]["source"] == source + + def test_replay_preserves_legacy_trigger_source_metadata(tmp_path, monkeypatch) -> None: monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path) key = "websocket:t-trigger-source" diff --git a/webui/src/hooks/useNanobotStream.ts b/webui/src/hooks/useNanobotStream.ts index a81b3a221..c8196e95c 100644 --- a/webui/src/hooks/useNanobotStream.ts +++ b/webui/src/hooks/useNanobotStream.ts @@ -37,7 +37,7 @@ interface ActiveAssistantCursor { } type PendingStreamEvent = - | { kind: "delta"; text: string; turn: UIMessageTurnFields } + | { kind: "delta"; text: string; turn: UIMessageTurnFields; source?: UIMessage["source"] } | { kind: "reasoning"; text: string; turn: UIMessageTurnFields }; type UIMessageTurnFields = Pick; @@ -778,7 +778,12 @@ export function useNanobotStream( }, []); const appendAnswerChunk = useCallback( - (prev: UIMessage[], chunk: string, turn: UIMessageTurnFields = {}): UIMessage[] => { + ( + prev: UIMessage[], + chunk: string, + turn: UIMessageTurnFields = {}, + source?: UIMessage["source"], + ): UIMessage[] => { let next = prev; let targetIndex = resolveActiveAssistantIndex(next, turn); @@ -809,6 +814,7 @@ export function useNanobotStream( content: target.content + chunk, isStreaming: true, ...turn, + ...(source ? { source } : {}), }; closedAssistantStreamIdsRef.current.delete(merged.id); activeAssistantRef.current = { id: merged.id, index: targetIndex }; @@ -823,7 +829,7 @@ export function useNanobotStream( let next = prev; for (const event of events) { if (event.kind === "delta") { - next = appendAnswerChunk(next, event.text, event.turn); + next = appendAnswerChunk(next, event.text, event.turn, event.source); } else { if (closeActiveAssistantStream()) clearActivitySegment(); next = attachReasoningChunk( @@ -843,6 +849,7 @@ export function useNanobotStream( closeAnswerSegment?: boolean; finalAnswerText?: string; turn?: UIMessageTurnFields; + source?: UIMessage["source"]; }) => { if (streamFrameRef.current !== null) { window.cancelAnimationFrame(streamFrameRef.current); @@ -855,7 +862,8 @@ export function useNanobotStream( const events = pendingStreamEventsRef.current; const finalAnswerText = options?.finalAnswerText; const turn = options?.turn ?? {}; - if (events.length === 0 && finalAnswerText === undefined) { + const source = options?.source; + if (events.length === 0 && finalAnswerText === undefined && source === undefined) { if (options?.closeAnswerSegment) closeActiveAssistantStream(); return; } @@ -873,6 +881,7 @@ export function useNanobotStream( content: finalAnswerText, isStreaming: true, ...turn, + ...(source ? { source } : {}), }; next = replaceMessageAt(next, targetIndex, merged); if (!options?.closeAnswerSegment) { @@ -890,6 +899,7 @@ export function useNanobotStream( content: finalAnswerText, isStreaming: true, ...turn, + ...(source ? { source } : {}), createdAt: Date.now(), }, ]; @@ -900,6 +910,18 @@ export function useNanobotStream( buffer.current = { messageId: id }; } } + } else if (source) { + const targetIndex = + resolveActiveAssistantIndex(next, turn) + ?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current, turn); + if (targetIndex !== null) { + const target = next[targetIndex]; + next = replaceMessageAt(next, targetIndex, { + ...target, + ...turn, + source, + }); + } } if (options?.closeAnswerSegment) closeActiveAssistantStream(); return next; @@ -1027,6 +1049,7 @@ export function useNanobotStream( kind: "delta", text: chunk, turn: turnFieldsFromEvent(ev, "answer"), + source: ev.source, }); schedulePendingStreamFlush(); return; @@ -1054,6 +1077,7 @@ export function useNanobotStream( closeAnswerSegment: !mergeNext, ...(typeof ev.text === "string" ? { finalAnswerText: ev.text } : {}), turn, + source: ev.source, }); if (suppressStreamUntilTurnEndRef.current) return; if (ev.resuming) { diff --git a/webui/src/lib/types.ts b/webui/src/lib/types.ts index 4c6b88979..68d623d97 100644 --- a/webui/src/lib/types.ts +++ b/webui/src/lib/types.ts @@ -1185,12 +1185,16 @@ export type InboundEvent = chat_id: string; text: string; stream_id?: string; + /** Lightweight provenance for proactive streamed assistant messages. */ + source?: UIMessageSource; } & InboundTurnMetadata) | ({ event: "stream_end"; chat_id: string; stream_id?: string; text?: string; + /** Lightweight provenance for proactive streamed assistant messages. */ + source?: UIMessageSource; /** This answer segment ended, but the active agent turn will continue. */ resuming?: boolean; /** The next answer segment continues this same assistant message. */ diff --git a/webui/src/tests/useNanobotStream.test.tsx b/webui/src/tests/useNanobotStream.test.tsx index b320318d8..1206461da 100644 --- a/webui/src/tests/useNanobotStream.test.tsx +++ b/webui/src/tests/useNanobotStream.test.tsx @@ -314,6 +314,63 @@ describe("useNanobotStream", () => { }); }); + it("preserves proactive automation source metadata on streamed assistant messages", () => { + const fake = fakeClient(); + const { result } = renderHook(() => useNanobotStream("chat-cron-stream", EMPTY_MESSAGES), { + wrapper: wrap(fake.client), + }); + const source = { kind: "cron", label: "Repo check" }; + + act(() => { + fake.emit("chat-cron-stream", { + event: "delta", + chat_id: "chat-cron-stream", + text: "Repo ", + source, + }); + fake.emit("chat-cron-stream", { + event: "stream_end", + chat_id: "chat-cron-stream", + source, + }); + fake.emit("chat-cron-stream", { + event: "turn_end", + chat_id: "chat-cron-stream", + }); + }); + + expect(result.current.messages[0]).toMatchObject({ + role: "assistant", + content: "Repo ", + isStreaming: false, + source, + }); + }); + + it("preserves proactive automation source metadata on stream_end final text", () => { + const fake = fakeClient(); + const { result } = renderHook(() => useNanobotStream("chat-cron-stream-end", EMPTY_MESSAGES), { + wrapper: wrap(fake.client), + }); + const source = { kind: "cron", label: "Repo check" }; + + act(() => { + fake.emit("chat-cron-stream-end", { + event: "stream_end", + chat_id: "chat-cron-stream-end", + text: "Repo clean.", + source, + }); + }); + + expect(result.current.messages[0]).toMatchObject({ + role: "assistant", + content: "Repo clean.", + isStreaming: true, + source, + }); + }); + it("does not start streaming from completed trailing activity after an answer", () => { const fake = fakeClient(); const initialMessages = [