Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(webui): unify turn observability

Browse Source
This commit is contained in:
Xubin Ren
2026-08-22 20:51:24 +08:00
parent dbc1801d3c
commit 48eea29313
44 changed files with 1871 additions and 930 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/channels/websocket/runtime.py
+65 -8
View File
@@ -426,6 +426,7 @@ class WebSocketChannel(BaseChannel):
)
self._stream_text_buffers: dict[tuple[str, str], list[str]] = {}
self._reasoning_text_buffers: dict[tuple[str, str], list[str]] = {}
# -- Subscription bookkeeping -------------------------------------------
@@ -482,6 +483,9 @@ class WebSocketChannel(BaseChannel):
for key in tuple(self._stream_text_buffers):
if key[0] == chat_id:
self._stream_text_buffers.pop(key, None)
for key in tuple(self._reasoning_text_buffers):
if key[0] == chat_id:
self._reasoning_text_buffers.pop(key, None)
async def _discard_connection_owned_chat(
self,
@@ -1641,11 +1645,22 @@ class WebSocketChannel(BaseChannel):
include_source=include_source,
transcript_overrides=transcript_overrides,
)
if (
not persisted
and phase in {"answer", "complete"}
and (metadata or {}).get("webui") is True
):
return self._retain_turn_on_transcript_failure(
chat_id,
persisted=persisted,
metadata=metadata,
phase=phase,
)
@staticmethod
def _retain_turn_on_transcript_failure(
chat_id: str,
*,
persisted: bool,
metadata: dict[str, Any] | None,
phase: str,
) -> bool:
if not persisted and phase in {"answer", "complete"} and (metadata or {}).get("webui") is True:
owner = (metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY)
mark_websocket_turn_transcript_persistence_failed(
chat_id,
@@ -1653,6 +1668,34 @@ class WebSocketChannel(BaseChannel):
)
return persisted
def _persist_turn_stream_event(
self,
chat_id: str,
event: dict[str, Any],
*,
completed_text: str | None,
metadata: dict[str, Any] | None,
phase: str,
include_source: bool = False,
) -> bool:
"""Persist the canonical end of a live stream, never its wire chunks."""
if not self._temporary_chats.should_persist_transcript(chat_id):
return True
persisted = self._transcripts.prepare_and_append_stream_event(
chat_id,
event,
completed_text=completed_text,
metadata=metadata,
phase=phase,
include_source=include_source,
)
return self._retain_turn_on_transcript_failure(
chat_id,
persisted=persisted,
metadata=metadata,
phase=phase,
)
async def send(self, msg: OutboundMessage) -> None:
event = outbound_event_from_message(msg)
progress_event = event if isinstance(event, ProgressEvent) else None
@@ -1685,6 +1728,7 @@ class WebSocketChannel(BaseChannel):
model_name=event.model,
model_preset=event.model_preset,
context_window_tokens=event.context_window_tokens,
fallback=event.fallback,
)
return
if isinstance(event, UserInputEvent):
@@ -1834,9 +1878,12 @@ class WebSocketChannel(BaseChannel):
}
if stream_id is not None:
body["stream_id"] = stream_id
self._persist_turn_transcript_event(
stream_key = (chat_id, str(stream_id or ""))
self._reasoning_text_buffers.setdefault(stream_key, []).append(delta)
self._persist_turn_stream_event(
chat_id,
body,
completed_text=None,
metadata=meta,
phase="reasoning",
)
@@ -1862,9 +1909,12 @@ class WebSocketChannel(BaseChannel):
}
if stream_id is not None:
body["stream_id"] = stream_id
self._persist_turn_transcript_event(
stream_key = (chat_id, str(stream_id or ""))
reasoning_text = "".join(self._reasoning_text_buffers.pop(stream_key, []))
self._persist_turn_stream_event(
chat_id,
body,
completed_text=reasoning_text or None,
metadata=meta,
phase="reasoning",
)
@@ -1912,6 +1962,7 @@ class WebSocketChannel(BaseChannel):
conns = list(self._subs.get(chat_id, ()))
meta = metadata or {}
stream_key = (chat_id, str(stream_id or ""))
completed_text: str | None = None
if stream_end:
body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id}
buffered = (
@@ -1923,6 +1974,7 @@ class WebSocketChannel(BaseChannel):
buffered.append(delta)
full_text = "".join(buffered)
rewritten = self._media.rewrite_local_markdown_images(full_text)
completed_text = rewritten
if delta or rewritten != full_text:
body["text"] = rewritten
else:
@@ -1938,9 +1990,10 @@ class WebSocketChannel(BaseChannel):
body["resuming"] = True
if stream_end and merge_next:
body["merge_next"] = True
self._persist_turn_transcript_event(
self._persist_turn_stream_event(
chat_id,
body,
completed_text=completed_text,
metadata=meta,
phase="answer",
include_source=True,
@@ -1997,6 +2050,7 @@ class WebSocketChannel(BaseChannel):
# carries a durable incomplete marker. The HTTP replay path can
# recover the latter from session history after a gateway restart.
clear_websocket_turn_if_current(chat_id, turn_owner)
self._clear_stream_buffers(chat_id)
raw = json.dumps(body, ensure_ascii=False)
if not conns:
return
@@ -2102,6 +2156,7 @@ class WebSocketChannel(BaseChannel):
model_name: Any,
model_preset: Any = None,
context_window_tokens: Any = None,
fallback: bool = False,
) -> None:
"""Notify one chat's subscribers which model is handling its current request."""
conns = list(self._subs.get(chat_id, ()))
@@ -2120,6 +2175,8 @@ class WebSocketChannel(BaseChannel):
body["model_preset"] = model_preset.strip()
if isinstance(context_window_tokens, int) and context_window_tokens > 0:
body["context_window_tokens"] = context_window_tokens
if fallback:
body["fallback"] = True
raw = json.dumps(body, ensure_ascii=False)
for connection in conns:
await self._safe_send_to(connection, raw, label=" turn_model_updated ")
nanobot/channels/websocket/tests/test_websocket_channel.py
+98 -2
View File
@@ -2073,6 +2073,21 @@ async def test_send_scopes_turn_model_updates_to_the_subscribed_chat() -> None:
"model_preset": "Deep Research",
"context_window_tokens": 128_000,
}
await channel.send(
OutboundMessage(
channel="websocket",
chat_id="chat-1",
content="",
event=TurnModelUpdatedEvent(
model="deepseek/deepseek-chat",
model_preset="Deep Research",
fallback=True,
),
)
)
fallback_payload = json.loads(chat_one.send.call_args.args[0])
assert fallback_payload["fallback"] is True
chat_two.send.assert_not_awaited()
@@ -2348,8 +2363,9 @@ async def test_send_delta_preserves_webui_source_metadata() -> None:
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
assert lines[-1]["event"] == "stream_end"
assert lines[-1]["text"] == "done"
@pytest.mark.asyncio
@@ -2374,6 +2390,8 @@ async def test_send_delta_marks_resuming_stream_end() -> None:
@pytest.mark.asyncio
async def test_send_delta_keeps_buffer_across_merged_stream_boundary() -> None:
from nanobot.webui.transcript import build_webui_thread_response, read_transcript_lines
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "streaming": True},
@@ -2403,6 +2421,12 @@ async def test_send_delta_keeps_buffer_across_merged_stream_boundary() -> None:
"second",
]
assert ("chat-1", "sid") not in channel._stream_text_buffers
lines = read_transcript_lines("websocket:chat-1")
assert [line["event"] for line in lines] == ["stream_end", "stream_end"]
assert [line["text"] for line in lines] == ["first ", "first second"]
body = build_webui_thread_response("websocket:chat-1")
assert body is not None
assert body["messages"][-1]["content"] == "first second"
@pytest.mark.asyncio
@@ -2596,7 +2620,8 @@ async def test_stream_transcript_persists_without_subscribers() -> None:
assert channel._subs == {}
lines = read_transcript_lines("websocket:chat-1")
assert [line["event"] for line in lines] == ["delta", "delta", "stream_end", "turn_end"]
assert [line["event"] for line in lines] == ["stream_end", "turn_end"]
assert lines[0]["text"] == "hello world"
body = build_webui_thread_response("websocket:chat-1")
assert body is not None
assert body["messages"][-1]["role"] == "assistant"
@@ -2604,6 +2629,77 @@ async def test_stream_transcript_persists_without_subscribers() -> None:
assert body["messages"][-1]["latencyMs"] == 42
@pytest.mark.asyncio
async def test_stream_transcript_writes_once_per_completed_segment(monkeypatch) -> None:
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "streaming": True},
bus,
gateway=_basic_handler(bus),
)
append = MagicMock()
monkeypatch.setattr("nanobot.webui.transcript.append_transcript_object", append)
await channel.send_delta("chat-write-rate", "one", stream_id="s1")
await channel.send_delta("chat-write-rate", " two", stream_id="s1")
await channel.send_delta("chat-write-rate", " three", stream_id="s1")
append.assert_not_called()
await channel.send_delta("chat-write-rate", "", stream_id="s1", stream_end=True)
append.assert_called_once()
persisted = append.call_args.args[1]
assert persisted["event"] == "stream_end"
assert persisted["text"] == "one two three"
@pytest.mark.asyncio
async def test_reasoning_transcript_persists_one_canonical_record(monkeypatch) -> None:
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"]},
bus,
gateway=_basic_handler(bus),
)
append = MagicMock()
monkeypatch.setattr("nanobot.webui.transcript.append_transcript_object", append)
await channel.send_reasoning_delta("chat-reasoning-write-rate", "plan ", stream_id="r1")
await channel.send_reasoning_delta("chat-reasoning-write-rate", "then act", stream_id="r1")
append.assert_not_called()
await channel.send_reasoning_end("chat-reasoning-write-rate", stream_id="r1")
append.assert_called_once()
persisted = append.call_args.args[1]
assert persisted["event"] == "reasoning_end"
assert persisted["text"] == "plan then act"
@pytest.mark.asyncio
async def test_turn_end_discards_unclosed_stream_buffers() -> None:
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "streaming": True},
bus,
gateway=_basic_handler(bus),
)
await channel.send_delta("chat-unclosed", "partial", stream_id="s1")
await channel.send_reasoning_delta("chat-unclosed", "thinking", stream_id="r1")
await channel.send(OutboundMessage(
channel="websocket",
chat_id="chat-unclosed",
content="",
event=TurnEndEvent(),
))
assert channel._stream_text_buffers == {}
assert channel._reasoning_text_buffers == {}
@pytest.mark.asyncio
async def test_send_turn_end_emits_turn_end_event() -> None:
bus = MagicMock()