refactor(webui): isolate fork websocket handler

This commit is contained in:
Xubin Ren 2026-06-10 04:01:06 +08:00
parent 916525f94a
commit 1432094bb5
2 changed files with 61 additions and 44 deletions

View File

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

View File

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