diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index 74c8077f4..9527c0dd7 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -34,7 +34,7 @@ from nanobot.utils.media_decode import ( 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.forking import handle_webui_fork_chat from nanobot.webui.gateway_services import GatewayServices from nanobot.webui.http_utils import ( normalize_config_path as _normalize_config_path, @@ -670,49 +670,7 @@ class WebSocketChannel(BaseChannel): await self._hydrate_after_subscribe(new_id) return if t == "fork_chat": - source_chat_id = envelope.get("source_chat_id") - raw_index = envelope.get("before_user_index") - if not _is_valid_chat_id(source_chat_id): - await self._send_event(connection, "error", detail="invalid source_chat_id") - return - if ( - isinstance(raw_index, bool) - or not isinstance(raw_index, int) - or raw_index < 0 - ): - await self._send_event(connection, "error", detail="invalid before_user_index") - return - if self.gateway.session_manager is None: - await self._send_event(connection, "error", detail="session_manager_unavailable") - return - - try: - 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 - fork_id, fork_key = forked - except Exception as exc: - 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(fork_key) - self._attach(connection, fork_id) - await self._send_event(connection, "attached", chat_id=fork_id) - await self._send_event( - connection, - "session_updated", - chat_id=fork_id, - scope="metadata", - workspace_scope=scope.payload(), - ) - await self._hydrate_after_subscribe(fork_id) + await handle_webui_fork_chat(self, connection, envelope) return if t == "attach": cid = envelope.get("chat_id") diff --git a/nanobot/webui/forking.py b/nanobot/webui/forking.py index c867ffc66..247cb8e6f 100644 --- a/nanobot/webui/forking.py +++ b/nanobot/webui/forking.py @@ -2,7 +2,10 @@ from __future__ import annotations +import re import uuid +from collections.abc import Mapping +from typing import Any from nanobot.session.manager import SessionManager from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title @@ -13,6 +16,12 @@ from nanobot.webui.transcript import ( write_session_messages_as_transcript, ) +_WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") + + +def _valid_webui_chat_id(value: Any) -> bool: + return isinstance(value, str) and _WEBUI_CHAT_ID_RE.match(value) is not None + def create_webui_chat_fork( session_manager: SessionManager, @@ -52,3 +61,53 @@ def create_webui_chat_fork( session_manager.delete_session(target_key) raise return new_id, target_key + + +async def handle_webui_fork_chat(channel: Any, connection: Any, envelope: Mapping[str, Any]) -> None: + """Handle the WebUI/desktop ``fork_chat`` websocket command. + + ``websocket.py`` owns the transport. This module owns WebUI fork semantics: + validate the request, clone session/transcript state, attach the new chat, + and hydrate the client. + """ + source_chat_id = envelope.get("source_chat_id") + raw_index = envelope.get("before_user_index") + if not _valid_webui_chat_id(source_chat_id): + await channel._send_event(connection, "error", detail="invalid source_chat_id") + return + if isinstance(raw_index, bool) or not isinstance(raw_index, int) or raw_index < 0: + await channel._send_event(connection, "error", detail="invalid before_user_index") + return + + session_manager = channel.gateway.session_manager + if session_manager is None: + await channel._send_event(connection, "error", detail="session_manager_unavailable") + return + + try: + forked = create_webui_chat_fork( + 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 channel._send_event(connection, "error", detail="invalid fork source or index") + return + fork_id, fork_key = forked + except Exception as exc: + channel.logger.warning("fork_chat failed: {}", exc) + await channel._send_event(connection, "error", detail="fork_chat_failed") + return + + scope = channel._workspaces.scope_for_session_key(fork_key) + channel._attach(connection, fork_id) + await channel._send_event(connection, "attached", chat_id=fork_id) + await channel._send_event( + connection, + "session_updated", + chat_id=fork_id, + scope="metadata", + workspace_scope=scope.payload(), + ) + await channel._hydrate_after_subscribe(fork_id)