fix(webui): preserve automation source on streamed replies

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:

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()

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 = []

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"

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<UIMessage, "turnId" | "turnPhase" | "turnSeq">;
@@ -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) {

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. */

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 = [