refactor(webui): shrink fork implementation
This commit is contained in:
@@ -696,22 +696,23 @@ class WebSocketChannel(BaseChannel):
|
||||
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(forked.session_key)
|
||||
self._attach(connection, forked.chat_id)
|
||||
await self._send_event(connection, "attached", chat_id=forked.chat_id)
|
||||
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=forked.chat_id,
|
||||
chat_id=fork_id,
|
||||
scope="metadata",
|
||||
workspace_scope=scope.payload(),
|
||||
)
|
||||
await self._hydrate_after_subscribe(forked.chat_id)
|
||||
await self._hydrate_after_subscribe(fork_id)
|
||||
return
|
||||
if t == "attach":
|
||||
cid = envelope.get("chat_id")
|
||||
|
||||
@@ -1,14 +1,8 @@
|
||||
"""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.
|
||||
"""
|
||||
"""WebUI chat fork orchestration."""
|
||||
|
||||
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
|
||||
@@ -20,25 +14,14 @@ from nanobot.webui.transcript import (
|
||||
)
|
||||
|
||||
|
||||
@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.
|
||||
"""
|
||||
) -> tuple[str, str] | None:
|
||||
"""Return ``(chat_id, session_key)`` for a new fork, or ``None`` for bad input."""
|
||||
new_id = str(uuid.uuid4())
|
||||
source_key = f"websocket:{source_chat_id}"
|
||||
target_key = f"websocket:{new_id}"
|
||||
@@ -68,4 +51,4 @@ def create_webui_chat_fork(
|
||||
delete_webui_transcript(target_key)
|
||||
session_manager.delete_session(target_key)
|
||||
raise
|
||||
return WebuiForkResult(chat_id=new_id, session_key=target_key)
|
||||
return new_id, target_key
|
||||
|
||||
+46
-82
@@ -286,6 +286,25 @@ def _is_user_transcript_row(row: dict[str, Any]) -> bool:
|
||||
return row.get("event") == "user" or row.get("role") == "user"
|
||||
|
||||
|
||||
def _write_transcript_lines(session_key: str, rows: list[dict[str, Any]]) -> None:
|
||||
path = webui_transcript_path(session_key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = path.with_suffix(".jsonl.tmp")
|
||||
try:
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
for row in rows:
|
||||
raw = json.dumps(row, ensure_ascii=False, separators=(",", ":"))
|
||||
if len(raw.encode("utf-8")) > _MAX_TRANSCRIPT_FILE_BYTES:
|
||||
raise ValueError("webui transcript line too large")
|
||||
f.write(raw + "\n")
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
except BaseException:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def fork_transcript_before_user_index(
|
||||
source_key: str,
|
||||
target_key: str,
|
||||
@@ -324,22 +343,7 @@ def fork_transcript_before_user_index(
|
||||
if not found_target:
|
||||
return False
|
||||
|
||||
path = webui_transcript_path(target_key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = path.with_suffix(".jsonl.tmp")
|
||||
try:
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
for row in copied:
|
||||
raw = json.dumps(row, ensure_ascii=False, separators=(",", ":"))
|
||||
if len(raw.encode("utf-8")) > _MAX_TRANSCRIPT_FILE_BYTES:
|
||||
raise ValueError("webui transcript line too large")
|
||||
f.write(raw + "\n")
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
except BaseException:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
_write_transcript_lines(target_key, copied)
|
||||
return True
|
||||
|
||||
|
||||
@@ -360,51 +364,29 @@ def write_session_messages_as_transcript(
|
||||
) -> None:
|
||||
"""Write a minimal WebUI transcript from already-truncated session messages."""
|
||||
target_chat_id = _chat_id_from_session_key(target_key)
|
||||
path = webui_transcript_path(target_key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = path.with_suffix(".jsonl.tmp")
|
||||
try:
|
||||
with open(tmp_path, "w", encoding="utf-8") as f:
|
||||
for msg in messages:
|
||||
role = msg.get("role")
|
||||
content = msg.get("content")
|
||||
text = content if isinstance(content, str) else ""
|
||||
if role == "user":
|
||||
row: dict[str, Any] = {
|
||||
"event": "user",
|
||||
"chat_id": target_chat_id,
|
||||
"text": text,
|
||||
}
|
||||
media = msg.get("media")
|
||||
if isinstance(media, list) and media:
|
||||
row["media_paths"] = [str(p) for p in media if isinstance(p, str) and p]
|
||||
for key in ("cli_apps", "mcp_presets"):
|
||||
value = msg.get(key)
|
||||
if isinstance(value, list) and value:
|
||||
row[key] = json.loads(json.dumps(value, ensure_ascii=False))
|
||||
elif role == "assistant":
|
||||
if not text.strip():
|
||||
continue
|
||||
row = {
|
||||
"event": "message",
|
||||
"chat_id": target_chat_id,
|
||||
"text": text,
|
||||
}
|
||||
media = msg.get("media")
|
||||
if isinstance(media, list) and media:
|
||||
row["media"] = [str(p) for p in media if isinstance(p, str) and p]
|
||||
else:
|
||||
continue
|
||||
raw = json.dumps(row, ensure_ascii=False, separators=(",", ":"))
|
||||
if len(raw.encode("utf-8")) > _MAX_TRANSCRIPT_FILE_BYTES:
|
||||
raise ValueError("webui transcript line too large")
|
||||
f.write(raw + "\n")
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
except BaseException:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
rows: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
role = msg.get("role")
|
||||
content = msg.get("content")
|
||||
text = content if isinstance(content, str) else ""
|
||||
if role == "user":
|
||||
row: dict[str, Any] = {"event": "user", "chat_id": target_chat_id, "text": text}
|
||||
media = msg.get("media")
|
||||
if isinstance(media, list) and media:
|
||||
row["media_paths"] = [str(p) for p in media if isinstance(p, str) and p]
|
||||
for key in ("cli_apps", "mcp_presets"):
|
||||
value = msg.get(key)
|
||||
if isinstance(value, list) and value:
|
||||
row[key] = json.loads(json.dumps(value, ensure_ascii=False))
|
||||
elif role == "assistant" and text.strip():
|
||||
row = {"event": "message", "chat_id": target_chat_id, "text": text}
|
||||
media = msg.get("media")
|
||||
if isinstance(media, list) and media:
|
||||
row["media"] = [str(p) for p in media if isinstance(p, str) and p]
|
||||
else:
|
||||
continue
|
||||
rows.append(row)
|
||||
_write_transcript_lines(target_key, rows)
|
||||
|
||||
|
||||
def delete_webui_transcript(session_key: str) -> bool:
|
||||
@@ -1411,25 +1393,12 @@ def replay_transcript_to_ui_messages(
|
||||
return messages
|
||||
|
||||
|
||||
def fork_boundary_message_count(
|
||||
lines: list[dict[str, Any]],
|
||||
*,
|
||||
augment_user_media: Callable[[list[str]], list[dict[str, Any]]] | None = None,
|
||||
augment_assistant_media: Callable[[list[str]], list[dict[str, Any]]] | None = None,
|
||||
augment_assistant_text: Callable[[str], str] | None = None,
|
||||
) -> int | None:
|
||||
def fork_boundary_message_count(lines: list[dict[str, Any]]) -> int | None:
|
||||
"""Return the replayed UI message count before the first fork marker, if any."""
|
||||
for idx, rec in enumerate(lines):
|
||||
if rec.get("event") != WEBUI_FORK_MARKER_EVENT:
|
||||
continue
|
||||
return len(
|
||||
replay_transcript_to_ui_messages(
|
||||
lines[:idx],
|
||||
augment_user_media=augment_user_media,
|
||||
augment_assistant_media=augment_assistant_media,
|
||||
augment_assistant_text=augment_assistant_text,
|
||||
),
|
||||
)
|
||||
return len(replay_transcript_to_ui_messages(lines[:idx]))
|
||||
return None
|
||||
|
||||
|
||||
@@ -1446,12 +1415,7 @@ def build_webui_thread_response(
|
||||
if not lines:
|
||||
return None
|
||||
lines = inject_missing_user_events_from_session(session_key, lines, session_messages)
|
||||
fork_boundary = fork_boundary_message_count(
|
||||
lines,
|
||||
augment_user_media=augment_user_media,
|
||||
augment_assistant_media=augment_assistant_media,
|
||||
augment_assistant_text=augment_assistant_text,
|
||||
)
|
||||
fork_boundary = fork_boundary_message_count(lines)
|
||||
msgs = replay_transcript_to_ui_messages(
|
||||
lines,
|
||||
augment_user_media=augment_user_media,
|
||||
|
||||
Reference in New Issue
Block a user