mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-04 08:28:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(webui): preserve automation source on streamed replies

Browse Source
This commit is contained in:
chengyongru 2026-07-07 10:17:23 +08:00 committed by chengyongru
parent 606ac56e8f
commit bb2f6cf324
7 changed files with 176 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
nanobot/channels/websocket/runtime.py
View File

@ -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:

30
nanobot/channels/websocket/tests/test_websocket_channel.py
View File

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

24
nanobot/webui/transcript.py
View File

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

36
tests/utils/test_webui_transcript.py
View File

@ -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"

32
webui/src/hooks/useNanobotStream.ts
View File

@ -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) {

4
webui/src/lib/types.ts
View File

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

57
webui/src/tests/useNanobotStream.test.tsx
View File

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