fix(tui): synchronize shared session clients

This commit is contained in:
Xubin Ren
2026-08-17 20:56:10 +08:00
parent 2f78f7fbc5
commit 9e47d8106c
25 changed files with 803 additions and 101 deletions
+85
View File
@@ -62,6 +62,7 @@ from nanobot.session.webui_turns import (
websocket_turn_transcript_persistence_failed,
websocket_turn_wall_started_at,
)
from nanobot.utils.helpers import safe_filename
from nanobot.webui.cli_apps_api import normalize_cli_app_mentions
from nanobot.webui.forking import handle_webui_fork_chat
from nanobot.webui.gateway_services import GatewayServices
@@ -594,6 +595,57 @@ class WebSocketChannel(BaseChannel):
for connection in tuple(self._webui_connections):
await self._send_event(connection, event, **fields)
async def _broadcast_user_message(
self,
origin: ServerConnection,
chat_id: str,
text: str,
*,
turn_id: str | None,
starts_turn: bool,
media_paths: list[str],
media_names: list[str | None],
cli_apps: list[dict[str, Any]],
mcp_presets: list[dict[str, Any]],
session_mentions: list[SessionMention],
) -> None:
"""Project one accepted user message to the other clients on the chat.
The origin already has an optimistic row and receives canonical turn
ownership in ``message_accepted``. Peers need the ingress projection.
"""
body: dict[str, Any] = {
"event": "user_message",
"chat_id": chat_id,
"text": text,
"starts_turn": starts_turn,
}
if turn_id is not None:
body["turn_id"] = turn_id
media = self._media.augment_transcript_user_media(media_paths)
for attachment, name in zip(media, media_names, strict=False):
if name:
attachment["name"] = name
if media:
body["media_urls"] = media
if cli_apps:
body["cli_apps"] = cli_apps
if mcp_presets:
body["mcp_presets"] = mcp_presets
if session_mentions:
body["session_mentions"] = session_mentions
active_turn_id = websocket_turn_id(chat_id)
if active_turn_id is not None:
body["active_turn_id"] = active_turn_id
started_at = websocket_turn_wall_started_at(chat_id)
if active_turn_id is not None and started_at is not None:
body["started_at"] = started_at
raw = json.dumps(body, ensure_ascii=False)
for connection in tuple(self._subs.get(chat_id, ())):
if connection is origin:
continue
await self._safe_send_to(connection, raw, label=" user_message ")
@classmethod
def default_config(cls) -> dict[str, Any]:
return WebSocketConfig().model_dump(by_alias=True)
@@ -1036,6 +1088,7 @@ class WebSocketChannel(BaseChannel):
raw_media = envelope.get("media")
media_paths: list[str] = []
media_names: list[str | None] = []
if raw_media is not None:
if not isinstance(raw_media, list):
await self._send_event(
@@ -1056,6 +1109,12 @@ class WebSocketChannel(BaseChannel):
**rejection_fields,
)
return
for item in cast(list[Any], raw_media):
attachment = cast(dict[str, Any], item) if isinstance(item, dict) else {}
name = attachment.get("name")
media_names.append(
(safe_filename(name) or None) if isinstance(name, str) else None
)
if temporary_policy is not None:
self._temporary_chats.register_media(connection, cid, media_paths)
@@ -1190,12 +1249,38 @@ class WebSocketChannel(BaseChannel):
finally:
if not accepted and queued_owner is not None:
clear_websocket_turn_if_current(cid, queued_owner)
if is_webui:
await self._broadcast_user_message(
connection,
cid,
content,
turn_id=turn_id,
starts_turn=queued_owner is not None,
media_paths=media_paths,
media_names=media_names,
cli_apps=cli_apps,
mcp_presets=mcp_presets,
session_mentions=session_mentions,
)
if is_webui and turn_id:
active_turn_id = websocket_turn_id(cid)
started_at = websocket_turn_wall_started_at(cid)
await self._send_event(
connection,
"message_accepted",
chat_id=cid,
turn_id=turn_id,
starts_turn=queued_owner is not None,
**(
{"active_turn_id": active_turn_id}
if active_turn_id is not None
else {}
),
**(
{"started_at": started_at}
if active_turn_id is not None and started_at is not None
else {}
),
)
return
await self._send_event(connection, "error", detail=f"unknown type: {t!r}")
@@ -4263,9 +4263,89 @@ async def test_authorized_webui_turn_is_acked_after_bus_acceptance(
"event": "message_accepted",
"chat_id": "chat-accepted",
"turn_id": "turn-accepted",
"starts_turn": True,
"active_turn_id": "turn-accepted",
"started_at": wth.websocket_turn_wall_started_at("chat-accepted"),
}
@pytest.mark.asyncio
async def test_user_messages_fan_out_to_other_clients_on_the_same_chat(
bus: MagicMock,
) -> None:
channel = _ch(bus)
origin = AsyncMock()
origin.remote_address = ("127.0.0.1", 50123)
peer = AsyncMock()
outsider = AsyncMock()
channel._attach(peer, "shared-chat")
channel._attach(outsider, "other-chat")
await channel._dispatch_envelope(
origin,
"terminal-a",
{
"type": "message",
"chat_id": "shared-chat",
"content": "hello from A",
"webui": True,
"turn_id": "turn-a",
},
)
await channel._dispatch_envelope(
origin,
"terminal-a",
{
"type": "message",
"chat_id": "shared-chat",
"content": "one more detail",
"webui": True,
"turn_id": "steer-a",
},
)
peer_payloads = [
payload
for payload in _sent_ws_payloads(peer)
if payload["event"] == "user_message"
]
assert len(peer_payloads) == 2
assert peer_payloads[0] == {
"event": "user_message",
"chat_id": "shared-chat",
"text": "hello from A",
"starts_turn": True,
"turn_id": "turn-a",
"active_turn_id": "turn-a",
"started_at": pytest.approx(wth.websocket_turn_wall_started_at("shared-chat")),
}
assert peer_payloads[1] == {
"event": "user_message",
"chat_id": "shared-chat",
"text": "one more detail",
"starts_turn": False,
"turn_id": "steer-a",
"active_turn_id": "turn-a",
"started_at": pytest.approx(wth.websocket_turn_wall_started_at("shared-chat")),
}
origin_payloads = _sent_ws_payloads(origin)
assert [
(
payload["event"],
payload["turn_id"],
payload["starts_turn"],
payload["active_turn_id"],
)
for payload in origin_payloads
if payload["event"] == "message_accepted"
] == [
("message_accepted", "turn-a", True, "turn-a"),
("message_accepted", "steer-a", False, "turn-a"),
]
outsider.send.assert_not_awaited()
assert bus.publish_inbound.await_count == 2
@pytest.mark.asyncio
async def test_side_channel_command_does_not_register_queued_turn(
bus: MagicMock,
@@ -4294,6 +4374,7 @@ async def test_side_channel_command_does_not_register_queued_turn(
"event": "message_accepted",
"chat_id": "chat-status",
"turn_id": "turn-status",
"starts_turn": False,
}
@@ -129,9 +129,47 @@ async def test_webui_message_acceptance_echoes_turn_id() -> None:
"event": "message_accepted",
"chat_id": "abc123",
"turn_id": "turn-accepted",
"starts_turn": True,
"active_turn_id": "turn-accepted",
"started_at": wth.websocket_turn_wall_started_at("abc123"),
}
@pytest.mark.asyncio
async def test_webui_message_projects_attachments_to_other_clients(tmp_path: Path) -> None:
channel = _make_channel()
origin = AsyncMock()
peer = AsyncMock()
channel._attach(origin, "abc123")
channel._attach(peer, "abc123")
channel._webui_connections.add(origin)
envelope = {
"type": "message",
"chat_id": "abc123",
"content": "please inspect @drawio",
"webui": True,
"turn_id": "turn-shared",
"media": [{"data_url": _tiny_png_data_url(), "name": "shot.png"}],
"cli_apps": [{"name": "DrawIO", "entry_point": "cli-anything-drawio"}],
}
with patch("nanobot.webui.media_gateway.get_media_dir", return_value=tmp_path):
await channel._dispatch_envelope(origin, "client-1", envelope)
event = json.loads(peer.send.await_args.args[0])
assert event["event"] == "user_message"
assert event["turn_id"] == "turn-shared"
assert event["text"] == "please inspect @drawio"
assert event["cli_apps"] == [{
"name": "drawio",
"entry_point": "cli-anything-drawio",
}]
assert event["media_urls"][0]["kind"] == "image"
assert event["media_urls"][0]["name"] == "shot.png"
assert event["media_urls"][0]["url"].startswith("/api/media/")
assert json.loads(origin.send.await_args.args[0])["event"] == "message_accepted"
@pytest.mark.asyncio
async def test_message_text_policy_is_independent_from_transport_limit() -> None:
channel = _make_channel()
@@ -334,6 +334,56 @@ async def test_independent_sessions(bus: MagicMock) -> None:
await t
@pytest.mark.asyncio
async def test_same_session_projects_one_turn_to_both_clients(bus: MagicMock) -> None:
ch = _ch(bus, 29925)
t = asyncio.create_task(ch.start())
try:
async with WsTestClient("ws://127.0.0.1:29925/", client_id="terminal-a") as a:
async with WsTestClient("ws://127.0.0.1:29925/", client_id="terminal-b") as b:
chat_id = (await a.recv_ready()).chat_id
await b.recv_ready()
await b.send_json({"type": "attach", "chat_id": chat_id})
attached = await b.recv()
assert attached.event == "attached"
assert attached.chat_id == chat_id
await a.send_json(
{
"type": "message",
"chat_id": chat_id,
"content": "hello from terminal A",
"webui": True,
"turn_id": "turn-a",
}
)
accepted = await a.recv()
projected = await b.recv()
assert accepted.event == "message_accepted"
assert accepted.raw["turn_id"] == "turn-a"
assert accepted.raw["starts_turn"] is True
assert accepted.raw["active_turn_id"] == "turn-a"
assert projected.raw == {
"event": "user_message",
"chat_id": chat_id,
"text": "hello from terminal A",
"starts_turn": True,
"turn_id": "turn-a",
"active_turn_id": "turn-a",
"started_at": projected.raw["started_at"],
}
await ch.send_delta(chat_id, "shared reply", stream_id="stream-a")
assert (await a.recv_delta()).text == "shared reply"
assert (await b.recv_delta()).text == "shared reply"
assert bus.publish_inbound.await_count == 1
finally:
await ch.stop()
await t
@pytest.mark.asyncio
async def test_disconnected_client_cleanup(bus: MagicMock) -> None:
ch = _ch(bus, 29914)
+45 -51
View File
@@ -33,15 +33,9 @@ class TuiUnavailableError(RuntimeError):
@dataclass(frozen=True)
class _GatewayLease:
runtime: Any
owned: bool
class _GatewayHandle:
base_url: str
def close(self) -> None:
if self.owned:
self.runtime.stop(timeout_s=20)
def launch_tui(
config: Config,
@@ -51,47 +45,44 @@ def launch_tui(
session_id: str | None,
theme: str,
) -> int:
"""Run the native TUI, owning a local gateway only when one is not running."""
"""Run the native TUI against the shared local gateway."""
command = _resolve_tui_command()
lease = _ensure_gateway(
gateway = _ensure_gateway(
config,
config_path=config_path,
workspace_override=workspace_override,
)
bootstrap = _fetch_bootstrap(
gateway.base_url,
secret=webui_bootstrap_secret(config),
)
env = os.environ.copy()
env.update(
{
"NANOBOT_TUI_WS_URL": _authenticated_ws_url(bootstrap),
"NANOBOT_TUI_API_URL": gateway.base_url,
"NANOBOT_TUI_API_TOKEN": str(bootstrap.get("api_token") or ""),
"NANOBOT_TUI_MODEL": _model_display(config)[0],
"NANOBOT_TUI_MODEL_PRESET": config.agents.defaults.model_preset or "default",
"NANOBOT_TUI_WORKSPACE": str(config.workspace_path),
"NANOBOT_TUI_VERSION": __version__,
"NANOBOT_TUI_ACCESS": (
"workspace access" if config.tools.restrict_to_workspace else "full access"
),
"NANOBOT_TUI_THEME": theme,
}
)
state_path = config_path.parent / "tui" / "state.json"
env["NANOBOT_TUI_STATE_PATH"] = str(state_path)
chat_id = _initial_tui_chat_id(session_id, state_path)
if chat_id:
env["NANOBOT_TUI_CHAT_ID"] = chat_id
else:
env.pop("NANOBOT_TUI_CHAT_ID", None)
try:
bootstrap = _fetch_bootstrap(
lease.base_url,
secret=webui_bootstrap_secret(config),
)
env = os.environ.copy()
env.update(
{
"NANOBOT_TUI_WS_URL": _authenticated_ws_url(bootstrap),
"NANOBOT_TUI_API_URL": lease.base_url,
"NANOBOT_TUI_API_TOKEN": str(bootstrap.get("api_token") or ""),
"NANOBOT_TUI_MODEL": _model_display(config)[0],
"NANOBOT_TUI_MODEL_PRESET": config.agents.defaults.model_preset or "default",
"NANOBOT_TUI_WORKSPACE": str(config.workspace_path),
"NANOBOT_TUI_VERSION": __version__,
"NANOBOT_TUI_ACCESS": (
"workspace access" if config.tools.restrict_to_workspace else "full access"
),
"NANOBOT_TUI_THEME": theme,
}
)
state_path = config_path.parent / "tui" / "state.json"
env["NANOBOT_TUI_STATE_PATH"] = str(state_path)
chat_id = _initial_tui_chat_id(session_id, state_path)
if chat_id:
env["NANOBOT_TUI_CHAT_ID"] = chat_id
else:
env.pop("NANOBOT_TUI_CHAT_ID", None)
try:
return subprocess.run(command, env=env, check=False).returncode
except OSError as exc:
raise TuiUnavailableError(f"could not start the native TUI: {exc}") from exc
finally:
lease.close()
return subprocess.run(command, env=env, check=False).returncode
except OSError as exc:
raise TuiUnavailableError(f"could not start the native TUI: {exc}") from exc
def _resolve_tui_command() -> list[str]:
@@ -236,7 +227,7 @@ def _ensure_gateway(
*,
config_path: Path,
workspace_override: str | None,
) -> _GatewayLease:
) -> _GatewayHandle:
from nanobot.gateway import GatewayRuntime, GatewayRuntimePaths, GatewayStartOptions
base_url = _webui_browser_url(config).split("/#/", 1)[0].rstrip("/")
@@ -265,7 +256,7 @@ def _ensure_gateway(
"restart it or use `nanobot agent --classic`"
)
if endpoint_reachable:
return _GatewayLease(runtime=runtime, owned=False, base_url=base_url)
return _GatewayHandle(base_url=base_url)
elif endpoint_reachable:
raise TuiUnavailableError(
"the configured gateway port belongs to a different nanobot instance; "
@@ -279,7 +270,7 @@ def _ensure_gateway(
config_path=str(config_path),
)
)
owned = result.ok
started_here = result.ok
if not result.ok and result.message != "gateway_already_running":
raise TuiUnavailableError(
f"could not start the local gateway ({result.message}); logs: {result.status.log_path}"
@@ -290,7 +281,7 @@ def _ensure_gateway(
if _webui_endpoint_reachable(base_url):
current = runtime.status()
if current.running and current.port in {None, config.gateway.port}:
return _GatewayLease(runtime=runtime, owned=owned, base_url=base_url)
return _GatewayHandle(base_url=base_url)
break
if not runtime.status().running and not _gateway_health_ready(
config.gateway.host,
@@ -299,7 +290,7 @@ def _ensure_gateway(
break
time.sleep(0.1)
if owned:
if started_here:
runtime.stop(timeout_s=5)
raise TuiUnavailableError(
f"local gateway did not become ready; logs: {result.status.log_path}"
@@ -343,16 +334,19 @@ def _websocket_chat_id(session_id: str) -> str | None:
"""Map the CLI selector to the WebSocket namespace used by the native TUI."""
if session_id.startswith("websocket:"):
return session_id.split(":", 1)[1] or None
if session_id == "cli:direct":
return "tui-direct"
return session_id.split(":", 1)[-1] or None
if ":" in session_id:
raise TuiUnavailableError(
"the native TUI can open only WebSocket sessions; use --classic to resume "
f"{session_id!r}"
)
return session_id or None
def _initial_tui_chat_id(session_id: str | None, state_path: Path) -> str | None:
"""Resume the default TUI, while keeping an explicit selector authoritative."""
if session_id is not None:
return _websocket_chat_id(session_id)
return _read_tui_chat_id(state_path) or _websocket_chat_id("cli:direct")
return _read_tui_chat_id(state_path) or "tui-direct"
def _read_tui_chat_id(path: Path) -> str | None: