From 1f926e3769b5a7bf5ed66dd98e62503a322532ea Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Wed, 10 Jun 2026 02:20:31 +0800 Subject: [PATCH] refactor(webui): isolate chat fork creation --- nanobot/channels/websocket.py | 51 ++++++------------------- nanobot/webui/forking.py | 71 +++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 39 deletions(-) create mode 100644 nanobot/webui/forking.py diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index ec26198e6..9ed3a0e76 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -28,16 +28,13 @@ from nanobot.security.workspace_access import ( WorkspaceScopeError, ) from nanobot.session.goal_state import goal_state_ws_blob -from nanobot.session.webui_turns import ( - WEBUI_TITLE_METADATA_KEY, - clean_generated_title, - websocket_turn_wall_started_at, -) +from nanobot.session.webui_turns import websocket_turn_wall_started_at from nanobot.utils.media_decode import ( FileSizeExceeded, save_base64_data_url, ) from nanobot.webui.cli_apps_api import normalize_cli_app_mentions +from nanobot.webui.forking import create_webui_chat_fork from nanobot.webui.gateway_services import GatewayServices from nanobot.webui.http_utils import ( normalize_config_path as _normalize_config_path, @@ -49,12 +46,6 @@ from nanobot.webui.http_utils import ( query_first as _query_first, ) from nanobot.webui.mcp_presets_api import normalize_mcp_preset_mentions -from nanobot.webui.transcript import ( - append_fork_marker, - delete_webui_transcript, - fork_transcript_before_user_index, - write_session_messages_as_transcript, -) from nanobot.webui.transcription_ws import webui_transcription_event from nanobot.webui.websocket_logging import websockets_server_logger @@ -695,50 +686,32 @@ class WebSocketChannel(BaseChannel): await self._send_event(connection, "error", detail="session_manager_unavailable") return - new_id = str(uuid.uuid4()) - source_key = f"websocket:{source_chat_id}" - target_key = f"websocket:{new_id}" try: - forked = self.gateway.session_manager.fork_session_before_user_index( - source_key, - target_key, - raw_index, + forked = create_webui_chat_fork( + self.gateway.session_manager, + source_chat_id=source_chat_id, + before_user_index=raw_index, + title=envelope.get("title") if isinstance(envelope.get("title"), str) else None, ) if forked is None: await self._send_event(connection, "error", detail="invalid fork source or index") return - transcript_ok = fork_transcript_before_user_index( - source_key, - target_key, - raw_index, - ) - if not transcript_ok: - write_session_messages_as_transcript(target_key, forked.messages) - append_fork_marker(target_key) - fork_title = clean_generated_title( - envelope.get("title") if isinstance(envelope.get("title"), str) else None, - ) - if fork_title: - forked.metadata[WEBUI_TITLE_METADATA_KEY] = fork_title - self.gateway.session_manager.save(forked, fsync=True) except Exception as exc: - delete_webui_transcript(target_key) - self.gateway.session_manager.delete_session(target_key) self.logger.warning("fork_chat failed: {}", exc) await self._send_event(connection, "error", detail="fork_chat_failed") return - scope = self._workspaces.scope_for_session_key(target_key) - self._attach(connection, new_id) - await self._send_event(connection, "attached", chat_id=new_id) + scope = self._workspaces.scope_for_session_key(forked.session_key) + self._attach(connection, forked.chat_id) + await self._send_event(connection, "attached", chat_id=forked.chat_id) await self._send_event( connection, "session_updated", - chat_id=new_id, + chat_id=forked.chat_id, scope="metadata", workspace_scope=scope.payload(), ) - await self._hydrate_after_subscribe(new_id) + await self._hydrate_after_subscribe(forked.chat_id) return if t == "attach": cid = envelope.get("chat_id") diff --git a/nanobot/webui/forking.py b/nanobot/webui/forking.py new file mode 100644 index 000000000..69669ab92 --- /dev/null +++ b/nanobot/webui/forking.py @@ -0,0 +1,71 @@ +"""Helpers for WebUI chat forking. + +The WebSocket channel owns transport concerns only. This module owns the +WebUI-specific session/transcript work needed to make a fork look like a normal +chat in both browser WebUI and desktop. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass + +from nanobot.session.manager import SessionManager +from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title +from nanobot.webui.transcript import ( + append_fork_marker, + delete_webui_transcript, + fork_transcript_before_user_index, + write_session_messages_as_transcript, +) + + +@dataclass(frozen=True) +class WebuiForkResult: + chat_id: str + session_key: str + + +def create_webui_chat_fork( + session_manager: SessionManager, + *, + source_chat_id: str, + before_user_index: int, + title: str | None = None, +) -> WebuiForkResult | None: + """Create a WebUI chat fork from a completed assistant-turn boundary. + + Returns ``None`` when the source/index is invalid. Exceptions are reserved + for unexpected I/O or persistence failures and are rolled back before being + re-raised. + """ + new_id = str(uuid.uuid4()) + source_key = f"websocket:{source_chat_id}" + target_key = f"websocket:{new_id}" + try: + forked = session_manager.fork_session_before_user_index( + source_key, + target_key, + before_user_index, + ) + if forked is None: + return None + + transcript_ok = fork_transcript_before_user_index( + source_key, + target_key, + before_user_index, + ) + if not transcript_ok: + write_session_messages_as_transcript(target_key, forked.messages) + append_fork_marker(target_key) + + fork_title = clean_generated_title(title) + if fork_title: + forked.metadata[WEBUI_TITLE_METADATA_KEY] = fork_title + session_manager.save(forked, fsync=True) + except Exception: + delete_webui_transcript(target_key) + session_manager.delete_session(target_key) + raise + return WebuiForkResult(chat_id=new_id, session_key=target_key)