diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index fc8e8bf20..d3913b343 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -3,19 +3,13 @@ from __future__ import annotations import asyncio -import hashlib -import hmac import ipaddress import json -import re import ssl -import time import uuid -from collections.abc import Callable from contextlib import suppress -from dataclasses import dataclass from pathlib import Path -from typing import Any, Self, TypeGuard, cast +from typing import TYPE_CHECKING, Any, Self, TypeGuard, cast from urllib.parse import urlsplit, urlunsplit from pydantic import Field, PrivateAttr, field_validator, model_validator @@ -24,56 +18,18 @@ from websockets.exceptions import ConnectionClosed from websockets.http11 import Request as WsRequest from nanobot.bus.events import ( - INBOUND_META_USER_SHELL, OUTBOUND_META_AGENT_UI, OutboundMessage, ) -from nanobot.bus.outbound_events import ( - GoalStateSyncEvent, - GoalStatusEvent, - ProgressEvent, - RecoveryStateEvent, - RuntimeModelUpdatedEvent, - SessionUpdatedEvent, - TurnEndEvent, - TurnModelUpdatedEvent, - UserInputEvent, - outbound_event_from_message, -) from nanobot.bus.queue import MessageBus from nanobot.channels.base import BaseChannel -from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn from nanobot.config.schema import Base -from nanobot.providers.base import LLMUsage -from nanobot.runtime_context import ( - RUNTIME_CONTEXT_INPUT_META, - WEBUI_QUOTE_METADATA, - RuntimeContextBlock, - webui_quote_runtime_context, -) -from nanobot.security.workspace_access import ( - WORKSPACE_SCOPE_METADATA_KEY, - WorkspaceScopeError, -) -from nanobot.session.goal_state import goal_state_ws_blob -from nanobot.session.model_selection import model_preset_from_metadata -from nanobot.session.recovery import recovery_state_from_metadata from nanobot.session.webui_turns import ( clear_websocket_turn_if_current, - clear_websocket_turns, mark_websocket_turn_transcript_persistence_failed, - register_queued_websocket_turn_if_idle, - websocket_turn_id, 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 -from nanobot.webui.http_utils import ( - is_trusted_proxy_authenticated_request as _is_trusted_proxy_authenticated_request, -) from nanobot.webui.http_utils import ( normalize_config_path as _normalize_config_path, ) @@ -83,27 +39,22 @@ from nanobot.webui.http_utils import ( 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.inbound_commands import WebUICommandRouter from nanobot.webui.metadata import ( WEBSOCKET_TURN_OWNER_METADATA_KEY, - WEBUI_SYSTEM_COMMAND_TURN_PREFIX, WEBUI_TURN_METADATA_KEY, ) -from nanobot.webui.session_access import ( - SessionMention, - WebuiSessionAccess, - session_mentions_runtime_context, -) -from nanobot.webui.sidebar_state import write_webui_sidebar_state -from nanobot.webui.temporary_chats import TemporaryChatError +from nanobot.webui.outbound_projection import WebUIOutboundProjector +from nanobot.webui.session_identity import is_valid_webui_chat_id from nanobot.webui.transcript import WEBUI_TRANSCRIPT_INCOMPLETE_KEY -from nanobot.webui.transcription_ws import webui_transcription_event from nanobot.webui.websocket_logging import websockets_server_logger +if TYPE_CHECKING: + from nanobot.bus.outbound_events import ProgressEvent, RecoveryStateEvent + from nanobot.providers.base import LLMUsage + # Plain HTTP WebUI routes also run through websockets.process_request. _WEBUI_HTTP_OPEN_TIMEOUT_S = 360.0 -_WEBUI_REQUEST_CACHE_TTL_S = 5 * 60.0 -_WEBUI_REQUEST_CACHE_MAX = 256 _ROUTING_ASSERTION_HEADERS = frozenset( @@ -317,13 +268,8 @@ def _parse_inbound_payload(raw: str) -> str | None: return text -# Accept UUIDs and short scoped keys like "unified:default". Keeps the capability -# namespace small enough to rule out path traversal / quote injection tricks. -_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") - - -def _is_valid_chat_id(value: Any) -> TypeGuard[str]: - return isinstance(value, str) and _CHAT_ID_RE.match(value) is not None +def _is_valid_chat_id(value: Any) -> TypeGuard[str]: # pyright: ignore[reportUnusedFunction] + return is_valid_webui_chat_id(value) def _parse_envelope(raw: str) -> dict[str, Any] | None: @@ -349,32 +295,6 @@ def _parse_envelope(raw: str) -> dict[str, Any] | None: return envelope -def _is_websocket_upgrade(request: WsRequest) -> bool: - """Detect an actual WS upgrade; plain HTTP GETs to the same path should fall through.""" - upgrade = request.headers.get("Upgrade") or request.headers.get("upgrade") - connection = request.headers.get("Connection") or request.headers.get("connection") - if not upgrade or "websocket" not in upgrade.lower(): - return False - if not connection or "upgrade" not in connection.lower(): - return False - return True - - -@dataclass(frozen=True) -class _WebUIRequestResult: - result: Any = None - status: int | None = None - message: str | None = None - - -@dataclass -class _WebUIRequestOperation: - action: str - payload_digest: bytes - task: asyncio.Task[_WebUIRequestResult] - completed_at: float | None = None - - class WebSocketChannel(BaseChannel): """Run a local WebSocket server; forward text/JSON messages to the message bus.""" @@ -399,71 +319,93 @@ class WebSocketChannel(BaseChannel): # connection -> default chat_id for legacy frames that omit routing. self._conn_default: dict[ServerConnection, str] = {} # Connections authenticated with a one-time token from /webui/bootstrap. - self._webui_connections: set[ServerConnection] = set() - # Delivery tasks are connection-bound, while operations are keyed only - # by request_id so reconnect retries join or replay the original work. - self._webui_request_tasks: dict[ - tuple[ServerConnection, str], - asyncio.Task[None], - ] = {} - self._webui_request_operations: dict[str, _WebUIRequestOperation] = {} - # Preserve request/response order for mutations from one - # UI. Without this, an earlier slow settings response can overwrite a - # newer settings snapshot in the client. - self._webui_request_locks: dict[ServerConnection, asyncio.Lock] = {} + self._webui_connections = gateway.endpoint.webui_connections self._stop_event: asyncio.Event | None = None self._server_task: asyncio.Task[None] | None = None self.gateway = gateway - self._http_router = gateway.http - self._tokens = gateway.tokens self._media = gateway.media - self._ingress = gateway.ingress self._transcripts = gateway.transcripts - self._workspaces = gateway.workspaces self._temporary_chats = gateway.temporary_chats - self._session_access = ( - WebuiSessionAccess(gateway.session_manager) - if gateway.session_manager is not None - else None - ) + self._session_projection = gateway.session_projection + self._commands = WebUICommandRouter(self, gateway) + self._webui_request_tasks = self._commands.request_tasks + self._webui_request_operations = self._commands.request_operations + self._webui_request_locks = self._commands.request_locks + self._outbound = WebUIOutboundProjector(self, self._session_projection) self._stream_text_buffers: dict[tuple[str, str], list[str]] = {} self._reasoning_text_buffers: dict[tuple[str, str], list[str]] = {} # -- Subscription bookkeeping ------------------------------------------- - def _workspace_controls_available(self, connection: ServerConnection) -> bool: - return self._http_router.workspace_controls_available(connection) + def webui_subscribers(self, chat_id: str) -> tuple[ServerConnection, ...]: + """Return a stable snapshot of one chat's transport subscribers.""" + return tuple(self._subs.get(chat_id, ())) + + def webui_connection_chats(self, connection: ServerConnection) -> tuple[str, ...]: + return tuple(self._conn_chats.get(connection, ())) + + def webui_attach(self, connection: ServerConnection, chat_id: str) -> None: + self._attach(connection, chat_id) + + def webui_detach(self, connection: ServerConnection, chat_id: str) -> None: + self._detach(connection, chat_id) + + def webui_clear_connection_default(self, connection: ServerConnection) -> None: + self._conn_default.pop(connection, None) + + def webui_clear_stream_buffers(self, chat_id: str) -> None: + self._clear_stream_buffers(chat_id) + + async def webui_hydrate(self, chat_id: str) -> None: + await self._hydrate_after_subscribe(chat_id) + + async def webui_send_event( + self, + connection: ServerConnection, + event: str, + **fields: Any, + ) -> None: + await self._send_event(connection, event, **fields) + + async def webui_send_raw( + self, + connection: ServerConnection, + raw: str, + *, + label: str = "", + ) -> None: + await self._safe_send_to(connection, raw, label=label) + + async def webui_dispatch_message( + self, + *, + sender_id: str, + chat_id: str, + content: str, + media: list[str] | None, + metadata: dict[str, Any], + is_dm: bool, + session_key: str | None, + require_existing_session: bool, + ) -> None: + await self._handle_message( + sender_id=sender_id, + chat_id=chat_id, + content=content, + media=media, + metadata=metadata, + is_dm=is_dm, + session_key=session_key, + require_existing_session=require_existing_session, + ) def _attach(self, connection: ServerConnection, chat_id: str) -> None: """Idempotently subscribe *connection* to *chat_id*.""" self._subs.setdefault(chat_id, set()).add(connection) self._conn_chats.setdefault(connection, set()).add(chat_id) - def _attached_model_fields(self, chat_id: str) -> dict[str, Any]: - """Expose small session runtime facts on the attach handshake.""" - sessions = self.gateway.session_manager - if sessions is None: - return {} - snapshot = sessions.read_session_metadata(f"websocket:{chat_id}") - raw_metadata = snapshot.get("metadata") if snapshot is not None else None - metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None - fields: dict[str, Any] = {} - try: - fields["model_preset"] = model_preset_from_metadata(metadata) - except ValueError: - self.logger.warning("ignoring invalid model preset metadata for chat_id={}", chat_id) - fields["model_preset"] = None - if isinstance(metadata, dict): - recovery_state = recovery_state_from_metadata(metadata) - if recovery_state is not None: - fields["recovery_state"] = recovery_state - usage = LLMUsage.from_dict(metadata.get("_last_usage")) - if usage is not None: - fields["usage"] = usage.to_turn_dict() - return fields - def _detach(self, connection: ServerConnection, chat_id: str) -> None: chats = self._conn_chats.get(connection) if chats is not None: @@ -484,98 +426,13 @@ class WebSocketChannel(BaseChannel): if key[0] == chat_id: self._reasoning_text_buffers.pop(key, None) - async def _discard_connection_owned_chat( - self, - connection: ServerConnection, - chat_id: str, - ) -> None: - await self._temporary_chats.discard(connection, chat_id) - self._detach(connection, chat_id) - clear_websocket_turns(chat_id) - self._clear_stream_buffers(chat_id) - - async def send_webui_protocol_error( - self, - connection: ServerConnection, - detail: str, - ) -> None: - """Send a stable protocol error from a WebUI-owned orchestration helper.""" - await self._send_event(connection, "error", detail=detail) - - async def attach_webui_fork( - self, - connection: ServerConnection, - *, - fork_id: str, - fork_key: str, - ) -> None: - """Attach and hydrate a newly created WebUI chat fork.""" - scope = self._workspaces.scope_for_session_key(fork_key) - self._attach(connection, fork_id) - await self._send_event( - connection, - "attached", - chat_id=fork_id, - **self._attached_model_fields(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) - async def _cleanup_connection(self, connection: ServerConnection) -> None: """Remove *connection* from every subscription set; safe to call multiple times.""" - chat_ids = tuple(self._conn_chats.get(connection, ())) - for cid in chat_ids: - if self._temporary_chats.owns(connection, cid): - await self._discard_connection_owned_chat(connection, cid) - else: - self._detach(connection, cid) - for cid in self._temporary_chats.chat_ids_for_owner(connection): - await self._discard_connection_owned_chat(connection, cid) - self._conn_default.pop(connection, None) - self._webui_connections.discard(connection) - self._discard_webui_request_lock_if_idle(connection) - - async def _maybe_push_persisted_goal_state(self, chat_id: str) -> None: - """Replay actionable goal state after *chat_id* is subscribed. - - Goal metadata lives on the session JSONL and survives gateway restarts, but - connected clients normally see it via ``goal_state`` / ``turn_end`` frames. - Pushing here makes refresh + reconnect restore the strip without a new model turn. - """ - if self.gateway.session_manager is None: - return - row = self.gateway.session_manager.read_session_file(f"websocket:{chat_id}") - row_data = row if isinstance(row, dict) else {} - meta = row_data.get("metadata", {}) - if not isinstance(meta, dict): - meta = {} - blob = goal_state_ws_blob(cast(dict[str, Any], meta)) - if not blob.get("active") and blob.get("status") != "blocked": - return - await self.send_goal_state(chat_id, blob) - - async def _maybe_push_turn_run_wall_clock(self, chat_id: str) -> None: - """Replay ``goal_status: running`` when a turn is still active (same-process refresh).""" - t0 = websocket_turn_wall_started_at(chat_id) - if t0 is None: - return - await self.send_goal_status( - chat_id, - "running", - started_at=t0, - turn_id=websocket_turn_id(chat_id), - ) + await self._commands.cleanup_connection(connection) async def _hydrate_after_subscribe(self, chat_id: str) -> None: """Replay persisted or actively running per-chat state after subscribe.""" - await self._maybe_push_persisted_goal_state(chat_id) - await self._maybe_push_turn_run_wall_clock(chat_id) + await self._outbound.hydrate(chat_id) async def _send_event( self, @@ -594,61 +451,6 @@ class WebSocketChannel(BaseChannel): except Exception as e: self.logger.warning("failed to send {} event: {}", event, e) - async def _broadcast_webui_event(self, event: str, **fields: Any) -> None: - 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) @@ -673,21 +475,12 @@ class WebSocketChannel(BaseChannel): # -- HTTP dispatch ------------------------------------------------------ async def _dispatch_http(self, connection: ServerConnection, request: WsRequest) -> Any: - """Route an inbound HTTP request to the HTTP handler or WS upgrade.""" - got, query = _parse_request_path(request.path) - expected_ws = self._expected_path() - - # WebSocket upgrade — channel handles this itself - if got == expected_ws and _is_websocket_upgrade(request): - client_id = _query_first(query, "client_id") or "" - if len(client_id) > 128: - client_id = client_id[:128] - if not self.is_allowed(client_id): - return connection.respond(403, "Forbidden") - return self._authorize_websocket_handshake(connection, query, request.headers) - - # Everything else goes to the HTTP handler - return await self._http_router.dispatch(connection, request) + """Compatibility proxy for the externally composed gateway endpoint.""" + return await self.gateway.endpoint.process_request( + connection, + request, + is_allowed=self.is_allowed, + ) def _authorize_websocket_handshake( self, @@ -695,34 +488,15 @@ class WebSocketChannel(BaseChannel): query: dict[str, list[str]], headers: Any = None, ) -> Any: - if _is_trusted_proxy_authenticated_request(connection, headers or {}, self.config): - self._webui_connections.add(connection) - return None - - supplied = _query_first(query, "token") - static_token = self.config.token.strip() - - if static_token: - if supplied and hmac.compare_digest(supplied, static_token): - return None - if supplied and self._consume_issued_token(connection, supplied): - return None - return connection.respond(401, "Unauthorized") - - if self.config.websocket_requires_token: - if supplied and self._consume_issued_token(connection, supplied): - return None - return connection.respond(401, "Unauthorized") - - if supplied: - self._consume_issued_token(connection, supplied) - return None + """Compatibility proxy for handshake tests and integrations.""" + return self.gateway.endpoint.authorize_websocket_handshake( + connection, + query, + headers, + ) def _consume_issued_token(self, connection: ServerConnection, token: str) -> bool: - audience = self._tokens.take_issued_token_audience(token) - if audience == "webui": - self._webui_connections.add(connection) - return audience is not None + return self.gateway.endpoint.consume_issued_token(connection, token) # -- Server lifecycle and connection ingress --------------------------- @@ -882,688 +656,12 @@ class WebSocketChannel(BaseChannel): client_id: str, envelope: dict[str, Any], ) -> None: - """Route one typed inbound envelope (``new_chat`` / ``attach`` / ``message``).""" - t = envelope.get("type") - if t == "webui_request": - await self._start_webui_request(connection, envelope) - return - if t == "new_chat": - new_id = str(uuid.uuid4()) - scope = await self._workspace_scope_or_error( - connection, - lambda: self._workspaces.scope_for_new_chat( - envelope, - controls_available=self._workspace_controls_available(connection), - ), - ) - if scope is None: - return - self._workspaces.stage_scope(new_id, scope) - self._attach(connection, new_id) - await self._send_event( - connection, - "attached", - chat_id=new_id, - **self._attached_model_fields(new_id), - ) - await self._send_event( - connection, - "session_updated", - chat_id=new_id, - scope="metadata", - workspace_scope=scope.payload(), - ) - await self._hydrate_after_subscribe(new_id) - return - if t == "new_temporary_chat": - try: - new_id = self._temporary_chats.create( - connection, - trusted_webui=connection in self._webui_connections, - ) - except TemporaryChatError as exc: - await self._send_event(connection, "error", detail=exc.detail) - return - self._attach(connection, new_id) - await self._send_event( - connection, - "attached", - chat_id=new_id, - temporary=True, - ) - return - if t == "fork_chat": - await handle_webui_fork_chat(self, connection, envelope) - return - if t == "discard_temporary_chat": - cid = envelope.get("chat_id") - if not _is_valid_chat_id(cid): - await self._send_event(connection, "error", detail="invalid temporary chat_id") - return - try: - await self._discard_connection_owned_chat(connection, cid) - except TemporaryChatError as exc: - await self._send_event(connection, "error", detail=exc.detail, chat_id=cid) - return - if t == "attach": - cid = envelope.get("chat_id") - if not _is_valid_chat_id(cid): - await self._send_event(connection, "error", detail="invalid chat_id") - return - try: - self._temporary_chats.validate_attach(cid) - except TemporaryChatError as exc: - await self._send_event(connection, "error", detail=exc.detail, chat_id=cid) - return - self._attach(connection, cid) - await self._send_event( - connection, - "attached", - chat_id=cid, - **self._attached_model_fields(cid), - ) - await self._hydrate_after_subscribe(cid) - return - if t == "set_sidebar_state": - if connection not in self._webui_connections: - await self._send_event(connection, "error", detail="access_denied") - return - state = envelope.get("state") - if not isinstance(state, dict): - await self._send_event( - connection, - "error", - detail="invalid_sidebar_state", - ) - return - try: - saved_state = await asyncio.to_thread( - write_webui_sidebar_state, - cast(dict[str, Any], state), - ) - except (OSError, ValueError): - await self._send_event( - connection, - "error", - detail="invalid_sidebar_state", - ) - return - await self._broadcast_webui_event( - "sidebar_state_updated", - state=saved_state, - ) - return - if t == "set_workspace_scope": - cid = envelope.get("chat_id") - if not _is_valid_chat_id(cid): - await self._send_event(connection, "error", detail="invalid chat_id") - return - try: - self._temporary_chats.validate_workspace_update(cid) - except TemporaryChatError as exc: - await self._send_event(connection, "error", detail=exc.detail, chat_id=cid) - return - scope = await self._workspace_scope_or_error( - connection, - lambda: self._workspaces.scope_for_set_request( - envelope, - chat_id=cid, - chat_running=websocket_turn_wall_started_at(cid) is not None, - controls_available=self._workspace_controls_available(connection), - ), - chat_id=cid, - ) - if scope is None: - return - self._workspaces.stage_scope(cid, scope) - # Other clients on the same gateway only need an invalidation; they - # can reload the authoritative session row without receiving a - # local project path that belongs to another connection. - await self.send_session_updated(cid, scope="metadata") - await self._send_event( - connection, - "session_updated", - chat_id=cid, - scope="metadata", - workspace_scope=scope.payload(), - ) - return - if t == "transcribe_audio": - event, payload = await webui_transcription_event( - envelope, - config_path=self.gateway.settings.config.path, - ) - await self._send_event(connection, event, **payload) - return - if t == "message": - cid = envelope.get("chat_id") - content = envelope.get("content") - if not _is_valid_chat_id(cid): - await self._send_event(connection, "error", detail="invalid chat_id") - return - raw_turn_id = envelope.get("turn_id") - turn_id = raw_turn_id if isinstance(raw_turn_id, str) and raw_turn_id else None - rejection_fields = { - "chat_id": cid, - **({"turn_id": turn_id} if turn_id else {}), - } - # The allowlist can change while an authenticated websocket stays - # open. Reject the exact application turn before hydration, - # transcript persistence, or an acceptance ACK; BaseChannel's - # silent authorization return must not look like successful ingress. - if not self.is_allowed(client_id): - await self._send_event( - connection, - "error", - detail="access_denied", - **rejection_fields, - ) - return - if not isinstance(content, str): - await self._send_event( - connection, - "error", - detail="missing content", - **rejection_fields, - ) - return - message_rejection = self._ingress.validate_text(content) - if message_rejection is not None: - await self._send_event( - connection, - "error", - detail="message_rejected", - reason=message_rejection, - **rejection_fields, - ) - return - - try: - temporary_policy = self._temporary_chats.message_policy( - connection, - cid, - content, - ) - except TemporaryChatError as exc: - await self._send_event( - connection, - "error", - detail=exc.detail, - **rejection_fields, - ) - return - - 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( - connection, - "error", - detail="attachment_rejected", - reason="malformed", - **rejection_fields, - ) - return - media_paths, reason = self._media.store_inbound_attachments(cast(list[Any], raw_media)) - if reason is not None: - await self._send_event( - connection, - "error", - detail="attachment_rejected", - reason=reason, - **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) - - # Allow media-only turns (content may be empty when attachments are present). - if not content.strip() and not media_paths: - await self._send_event( - connection, - "error", - detail="missing content", - **rejection_fields, - ) - return - # Auto-attach on first use so clients can one-shot without a separate attach. - self._attach(connection, cid) - if temporary_policy is None or temporary_policy.hydrate_transcript: - await self._hydrate_after_subscribe(cid) - - # Resolve after hydration so a concurrent downgrade cannot be overwritten. - scope = await self._workspace_scope_or_error( - connection, - lambda: ( - temporary_policy.workspace_scope - if temporary_policy is not None - else self._workspaces.scope_for_message( - envelope, - chat_id=cid, - chat_running=websocket_turn_wall_started_at(cid) is not None, - controls_available=self._workspace_controls_available(connection), - ) - ), - chat_id=cid, - turn_id=turn_id, - ) - if scope is None: - return - - # Hydration and scope resolution can yield. Re-check immediately - # before transcript/bus mutation so a mid-flight revocation cannot - # fall through BaseChannel's silent deny and still receive an ACK. - if not self.is_allowed(client_id): - await self._send_event( - connection, - "error", - detail="access_denied", - **rejection_fields, - ) - return - - metadata: dict[str, Any] = {"remote": getattr(connection, "remote_address", None)} - if envelope.get("webui") is True: - metadata["webui"] = True - metadata.update(self._transcripts.client_turn_metadata(envelope.get("turn_id"))) - trusted_webui = metadata.get("webui") is True and connection in self._webui_connections - is_user_shell = ( - trusted_webui - and envelope.get("user_shell") is True - and content.startswith("!") - ) - if is_user_shell: - metadata[INBOUND_META_USER_SHELL] = True - dispatch_content = ( - f"{USER_SHELL_COMMAND} {content[1:].lstrip()}" - if is_user_shell - else content - ) - cli_apps = normalize_cli_app_mentions(envelope.get("cli_apps")) - if cli_apps: - metadata["cli_apps"] = cli_apps - mcp_presets = normalize_mcp_preset_mentions( - envelope.get("mcp_presets"), - config_path=self.gateway.settings.config.path, - ) - if mcp_presets: - metadata["mcp_presets"] = mcp_presets - session_mentions: list[SessionMention] = [] - if ( - trusted_webui - and self._session_access is not None - ): - session_mentions = await asyncio.to_thread( - self._session_access.normalize_mentions, - envelope.get("session_mentions"), - exclude_session_key=f"{self.name}:{cid}", - ) - if session_mentions: - metadata["session_mentions"] = session_mentions - metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata() - is_webui = metadata.get("webui") is True - queued_owner = None - if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content): - queued_owner = register_queued_websocket_turn_if_idle(cid, turn_id) - if queued_owner is not None: - metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] = queued_owner - accepted = False - try: - if ( - is_webui - and ( - temporary_policy is None - or temporary_policy.persist_transcript - ) - ): - self._transcripts.append_user_message( - cid, - content, - metadata=metadata, - media_paths=media_paths or None, - cli_apps=cli_apps or None, - mcp_presets=mcp_presets or None, - session_mentions=session_mentions or None, - ) - if trusted_webui: - context_blocks: list[RuntimeContextBlock] = [] - quote = webui_quote_runtime_context({ - WEBUI_QUOTE_METADATA: envelope.get("quoted_context"), - }) - if quote is not None: - context_blocks.append(quote) - session_context = session_mentions_runtime_context(session_mentions) - if session_context is not None: - context_blocks.append(session_context) - if context_blocks: - metadata[RUNTIME_CONTEXT_INPUT_META] = context_blocks - await self._handle_message( - sender_id=client_id, - chat_id=cid, - content=dispatch_content, - media=media_paths or None, - metadata=metadata, - is_dm=False, - session_key=( - temporary_policy.session_key - if temporary_policy is not None - else None - ), - require_existing_session=( - temporary_policy.require_existing_session - if temporary_policy is not None - else False - ), - ) - self._workspaces.persist_scope(cid, scope) - accepted = True - 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}") - - async def _start_webui_request( - self, - connection: ServerConnection, - envelope: dict[str, Any], - ) -> None: - request_id = envelope.get("request_id") - if not isinstance(request_id, str) or re.fullmatch( - r"[A-Za-z0-9._:-]{1,128}", - request_id, - ) is None: - await self._send_event( - connection, - "error", - detail="invalid webui request_id", - ) - return - if connection not in self._webui_connections: - await self._send_webui_response( - connection, - request_id, - status=403, - message="access_denied", - ) - return - - action = envelope.get("action") - payload = envelope.get("payload") - if not isinstance(action, str) or re.fullmatch( - r"[a-z][a-z0-9_.]{0,127}", - action, - ) is None: - await self._send_webui_response( - connection, - request_id, - status=400, - message="invalid WebUI mutation action", - ) - return - if not isinstance(payload, dict): - await self._send_webui_response( - connection, - request_id, - status=400, - message="WebUI mutation payload must be an object", - ) - return - - payload_digest = hashlib.sha256( - json.dumps( - payload, - ensure_ascii=False, - sort_keys=True, - separators=(",", ":"), - ).encode("utf-8") - ).digest() - self._prune_webui_request_operations() - operation = self._webui_request_operations.get(request_id) - is_replay = operation is not None - if operation is not None and ( - operation.action != action or operation.payload_digest != payload_digest - ): - await self._send_webui_response( - connection, - request_id, - status=409, - message="request_id was already used for a different WebUI mutation", - ) - return - if operation is None: - operation_task = asyncio.create_task( - self._execute_webui_request( - connection, - action, - cast(dict[str, Any], payload), - ) - ) - new_operation = _WebUIRequestOperation( - action=action, - payload_digest=payload_digest, - task=operation_task, - ) - operation = new_operation - self._webui_request_operations[request_id] = new_operation - - def mark_complete(_task: asyncio.Task[_WebUIRequestResult]) -> None: - current = self._webui_request_operations.get(request_id) - if current is not new_operation: - return - new_operation.completed_at = time.monotonic() - self._prune_webui_request_operations() - - operation_task.add_done_callback(mark_complete) - - key = (connection, request_id) - if key in self._webui_request_tasks: - return - delivery_task = asyncio.create_task( - self._deliver_webui_request( - connection, - request_id, - operation.task, - sequence=is_replay, - ) - ) - self._webui_request_tasks[key] = delivery_task + """Delegate one typed envelope to the WebUI application router.""" + await self._commands.dispatch(connection, client_id, envelope) def _prune_webui_request_operations(self) -> None: - now = time.monotonic() - for request_id, operation in tuple(self._webui_request_operations.items()): - if ( - operation.completed_at is not None - and now - operation.completed_at >= _WEBUI_REQUEST_CACHE_TTL_S - ): - self._webui_request_operations.pop(request_id, None) - - completed = sorted( - ( - (operation.completed_at, request_id) - for request_id, operation in self._webui_request_operations.items() - if operation.completed_at is not None - ), - key=lambda item: item[0], - ) - for _, request_id in completed[:-_WEBUI_REQUEST_CACHE_MAX]: - self._webui_request_operations.pop(request_id, None) - - def _discard_webui_request_lock_if_idle(self, connection: ServerConnection) -> None: - if connection in self._webui_connections: - return - if any(task_connection is connection for task_connection, _ in self._webui_request_tasks): - return - self._webui_request_locks.pop(connection, None) - - async def _deliver_webui_request( - self, - connection: ServerConnection, - request_id: str, - operation_task: asyncio.Task[_WebUIRequestResult], - *, - sequence: bool = False, - ) -> None: - try: - if sequence: - # Make replayed work the predecessor for subsequent mutations on - # this connection without blocking its receive loop. - lock = self._webui_request_locks.setdefault(connection, asyncio.Lock()) - async with lock: - result = await asyncio.shield(operation_task) - await self._send_webui_response( - connection, - request_id, - result=result.result, - status=result.status, - message=result.message, - ) - return - result = await asyncio.shield(operation_task) - await self._send_webui_response( - connection, - request_id, - result=result.result, - status=result.status, - message=result.message, - ) - finally: - self._webui_request_tasks.pop((connection, request_id), None) - self._discard_webui_request_lock_if_idle(connection) - - async def _execute_webui_request( - self, - connection: ServerConnection, - action: str, - payload: dict[str, Any], - ) -> _WebUIRequestResult: - try: - lock = self._webui_request_locks.setdefault(connection, asyncio.Lock()) - async with lock: - response = await self._http_router.dispatch_webui_mutation( - connection, - action, - payload, - ) - status = response.status_code - body = bytes(response.body).decode("utf-8", errors="replace").strip() - if 200 <= status < 300: - try: - result = json.loads(body) - except json.JSONDecodeError: - return _WebUIRequestResult( - status=502, - message="WebUI mutation returned an invalid response", - ) - if action == "sidebar.update" and isinstance(result, dict): - await self._broadcast_webui_event( - "sidebar_state_updated", - state=result, - ) - return _WebUIRequestResult(result=result) - return _WebUIRequestResult( - status=status, - message=body or response.reason_phrase, - ) - except asyncio.CancelledError: - raise - except Exception: - self.logger.exception("WebUI mutation '{}' failed", action) - return _WebUIRequestResult( - status=500, - message="WebUI mutation failed", - ) - - async def _send_webui_response( - self, - connection: ServerConnection, - request_id: str, - *, - result: Any = None, - status: int | None = None, - message: str | None = None, - ) -> None: - if status is None: - await self._send_event( - connection, - "webui_response", - request_id=request_id, - ok=True, - result=result, - ) - return - await self._send_event( - connection, - "webui_response", - request_id=request_id, - ok=False, - error={ - "status": status, - "message": message or "WebUI mutation failed", - }, - ) - - async def _workspace_scope_or_error( - self, - connection: ServerConnection, - resolver: Callable[[], Any], - *, - chat_id: str | None = None, - turn_id: str | None = None, - ) -> Any | None: - try: - return resolver() - except WorkspaceScopeError as exc: - await self._send_event( - connection, - "error", - detail="workspace_scope_rejected", - reason=exc.message, - **({"chat_id": chat_id} if chat_id else {}), - **({"turn_id": turn_id} if turn_id else {}), - ) - return None + """Compatibility hook for request-cache boundary tests.""" + self._commands.prune_request_operations() # -- Outbound WebSocket events ----------------------------------------- @@ -1584,25 +682,10 @@ class WebSocketChannel(BaseChannel): except Exception as e: self.logger.warning("server task error during shutdown: {}", e) self._server_task = None - delivery_tasks = tuple(self._webui_request_tasks.values()) - operation_tasks = tuple( - operation.task for operation in self._webui_request_operations.values() - ) - for task in (*delivery_tasks, *operation_tasks): - task.cancel() - if delivery_tasks: - await asyncio.gather(*delivery_tasks, return_exceptions=True) - if operation_tasks: - await asyncio.gather(*operation_tasks, return_exceptions=True) - self._webui_request_tasks.clear() - self._webui_request_locks.clear() - self._webui_request_operations.clear() + await self._commands.close() self._subs.clear() self._conn_chats.clear() self._conn_default.clear() - self._webui_connections.clear() - self._tokens.clear() - self._temporary_chats.close() async def _safe_send_to( self, @@ -1694,116 +777,15 @@ class WebSocketChannel(BaseChannel): ) async def send(self, msg: OutboundMessage) -> None: - event = outbound_event_from_message(msg) - progress_event = event if isinstance(event, ProgressEvent) else None - if isinstance(event, RuntimeModelUpdatedEvent): - await self.send_runtime_model_updated( - model_name=event.model, - model_preset=event.model_preset, - ) - return + await self._outbound.send(msg) - # Snapshot the subscriber set so ConnectionClosed cleanups mid-iteration are safe. + async def send_projected_message( + self, + msg: OutboundMessage, + progress_event: ProgressEvent | None, + ) -> None: + """Serialize one ordinary outbound message selected by the projector.""" conns = list(self._subs.get(msg.chat_id, ())) - if not conns: - if isinstance( - event, - ProgressEvent - | UserInputEvent - | TurnEndEvent - | SessionUpdatedEvent - | GoalStatusEvent - | GoalStateSyncEvent, - ): - self.logger.debug("no active subscribers for chat_id={}", msg.chat_id) - else: - self.logger.warning("no active subscribers for chat_id={}", msg.chat_id) - if isinstance(event, TurnModelUpdatedEvent): - if conns: - await self.send_turn_model_updated( - msg.chat_id, - model_name=event.model, - model_preset=event.model_preset, - context_window_tokens=event.context_window_tokens, - fallback=event.fallback, - ) - return - if isinstance(event, UserInputEvent): - if conns: - await self.send_user_input( - msg.chat_id, - content=event.content, - created_at_ms=event.created_at_ms, - provenance=event.provenance, - ) - return - if isinstance(event, RecoveryStateEvent): - if conns: - await self.send_recovery_state(msg.chat_id, event) - return - if isinstance(event, GoalStateSyncEvent): - if conns: - await self.send_goal_state(msg.chat_id, event.goal_state or {"active": False}) - return - if isinstance(event, GoalStatusEvent): - turn_id = (msg.metadata or {}).get(WEBUI_TURN_METADATA_KEY) - current_turn_id = turn_id if isinstance(turn_id, str) else None - turn_owner = (msg.metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY) - current_turn_owner = turn_owner if isinstance(turn_owner, str) else None - try: - if conns and event.status in ("running", "idle"): - await self.send_goal_status( - msg.chat_id, - event.status, - started_at=event.started_at, - turn_id=current_turn_id, - ) - finally: - if event.status == "idle": - # Cancellation/direct runs may have no turn_end, so idle is - # still terminal. A failed canonical completion write is - # the one case that must remain pending for safe resume. - clear_websocket_turn_if_current( - msg.chat_id, - current_turn_owner, - preserve_persistence_failure=True, - ) - return - # Signal that the agent has fully finished processing the current turn. - if isinstance(event, TurnEndEvent): - turn_id = (msg.metadata or {}).get(WEBUI_TURN_METADATA_KEY) - session_update_scope = ( - "metadata" - if isinstance(turn_id, str) - and turn_id.startswith(WEBUI_SYSTEM_COMMAND_TURN_PREFIX) - else "thread" - ) - turn_owner = (msg.metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY) - await self.send_turn_end( - msg.chat_id, - latency_ms=event.latency_ms, - goal_state=event.goal_state, - usage=event.usage, - context_window_tokens=event.context_window_tokens, - metadata=msg.metadata, - turn_owner=turn_owner if isinstance(turn_owner, str) else None, - ) - await self.send_session_updated(msg.chat_id, scope=session_update_scope) - return - if isinstance(event, SessionUpdatedEvent): - if conns: - await self.send_session_updated( - msg.chat_id, - scope=event.scope, - ) - return - if progress_event and progress_event.file_edit_events: - await self.send_file_edit_events( - msg.chat_id, - progress_event.file_edit_events, - msg.metadata, - ) - return text = msg.content wire_text = self._media.rewrite_local_markdown_images(text) payload: dict[str, Any] = { diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index d0125a74d..16533020e 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -1240,7 +1240,7 @@ def test_webui_request_cache_prunes_expired_completed_but_keeps_pending( bus: MagicMock, monkeypatch: pytest.MonkeyPatch, ) -> None: - import nanobot.channels.websocket.runtime as websocket_module + import nanobot.webui.inbound_commands as websocket_module channel = _ch(bus) now = 1_000.0 @@ -1263,7 +1263,7 @@ def test_webui_request_cache_prunes_oldest_completed_at_capacity( bus: MagicMock, monkeypatch: pytest.MonkeyPatch, ) -> None: - import nanobot.channels.websocket.runtime as websocket_module + import nanobot.webui.inbound_commands as websocket_module channel = _ch(bus) now = 1_000.0 @@ -1555,7 +1555,7 @@ async def test_new_chat_without_message_does_not_create_session( attached = json.loads(conn.send.await_args_list[0].args[0]) assert attached["event"] == "attached" assert sessions.list_sessions() == [] - assert channel._workspaces.scope_for_session_key( + assert channel.gateway.workspaces.scope_for_session_key( f"websocket:{attached['chat_id']}" ).access_mode == "full" @@ -1813,9 +1813,9 @@ async def test_webui_set_workspace_scope_rejects_running_chat(bus: MagicMock, tm }, }, ) - channel._workspaces.persist_scope( + channel.gateway.workspaces.persist_scope( "chat-running", - channel._workspaces.scope_for_session_key("websocket:chat-running"), + channel.gateway.workspaces.scope_for_session_key("websocket:chat-running"), ) conn.send.reset_mock() @@ -1960,7 +1960,7 @@ async def test_remote_access_reduction_rejects_stale_in_flight_message_scope( await message_task assert sessions.read_session_file(f"websocket:{chat_id}") is None - assert channel._workspaces.scope_for_session_key( + assert channel.gateway.workspaces.scope_for_session_key( f"websocket:{chat_id}" ).access_mode == "restricted" payload = json.loads(message_conn.send.await_args.args[0]) @@ -2051,7 +2051,7 @@ async def test_native_webui_scope_allows_custom_scope_without_loopback( assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve()) assert sessions.read_session_file("websocket:chat-native") is None - assert channel._workspaces.scope_for_session_key( + assert channel.gateway.workspaces.scope_for_session_key( "websocket:chat-native" ).metadata() == { "project_path": str(project.resolve()), @@ -2190,28 +2190,6 @@ async def test_send_scopes_turn_model_updates_to_the_subscribed_chat() -> None: chat_two.send.assert_not_awaited() -def test_attach_fields_restore_the_session_model_and_latest_usage() -> None: - usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175) - manager = MagicMock() - manager.read_session_metadata.return_value = { - "metadata": { - SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research", - "_last_usage": usage.to_dict(), - } - } - bus = MagicMock() - channel = WebSocketChannel( - {"enabled": True, "allowFrom": ["*"]}, - bus, - gateway=_basic_handler(bus, session_manager=manager), - ) - - assert channel._attached_model_fields("chat-1") == { - "model_preset": "Deep Research", - "usage": usage.to_turn_dict(), - } - - @pytest.mark.asyncio async def test_send_stages_external_media_as_signed_url(monkeypatch, tmp_path) -> None: bus = MagicMock() @@ -3459,20 +3437,20 @@ async def test_send_goal_state_emits_blob_per_chat() -> None: @pytest.mark.asyncio -async def test_maybe_push_active_goal_state_noop_without_session_manager() -> None: +async def test_hydrate_noop_without_session_manager() -> None: bus = MagicMock() channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) mock_ws = AsyncMock() channel._attach(mock_ws, "chat-1") - await channel._maybe_push_persisted_goal_state("chat-1") + await channel._outbound.hydrate("chat-1") mock_ws.send.assert_not_called() @pytest.mark.asyncio -async def test_maybe_push_active_goal_state_skips_when_no_goal_on_disk() -> None: +async def test_hydrate_skips_when_no_goal_on_disk() -> None: bus = MagicMock() sm = MagicMock() - sm.read_session_file.return_value = None + sm.read_session_metadata.return_value = None channel = WebSocketChannel( {"enabled": True, "allowFrom": ["*"]}, bus, @@ -3480,15 +3458,15 @@ async def test_maybe_push_active_goal_state_skips_when_no_goal_on_disk() -> None ) mock_ws = AsyncMock() channel._attach(mock_ws, "chat-1") - await channel._maybe_push_persisted_goal_state("chat-1") + await channel._outbound.hydrate("chat-1") mock_ws.send.assert_not_called() @pytest.mark.asyncio -async def test_maybe_push_active_goal_state_notifies_when_goal_active_on_disk() -> None: +async def test_hydrate_notifies_when_goal_active_on_disk() -> None: bus = MagicMock() sm = MagicMock() - sm.read_session_file.return_value = { + sm.read_session_metadata.return_value = { "metadata": { "goal_state": { "status": "active", @@ -3505,7 +3483,7 @@ async def test_maybe_push_active_goal_state_notifies_when_goal_active_on_disk() ) mock_ws = AsyncMock() channel._attach(mock_ws, "chat-1") - await channel._maybe_push_persisted_goal_state("chat-1") + await channel._outbound.hydrate("chat-1") mock_ws.send.assert_awaited_once() body = json.loads(mock_ws.send.await_args.args[0]) assert body["event"] == "goal_state" @@ -3516,10 +3494,10 @@ async def test_maybe_push_active_goal_state_notifies_when_goal_active_on_disk() @pytest.mark.asyncio -async def test_maybe_push_goal_state_restores_blocked_attention_on_disk() -> None: +async def test_hydrate_restores_blocked_attention_on_disk() -> None: bus = MagicMock() sm = MagicMock() - sm.read_session_file.return_value = { + sm.read_session_metadata.return_value = { "metadata": { "goal_state": { "status": "blocked", @@ -3537,7 +3515,7 @@ async def test_maybe_push_goal_state_restores_blocked_attention_on_disk() -> Non mock_ws = AsyncMock() channel._attach(mock_ws, "chat-1") - await channel._maybe_push_persisted_goal_state("chat-1") + await channel._outbound.hydrate("chat-1") body = json.loads(mock_ws.send.await_args.args[0]) assert body["goal_state"] == { @@ -3549,7 +3527,7 @@ async def test_maybe_push_goal_state_restores_blocked_attention_on_disk() -> Non @pytest.mark.asyncio -async def test_maybe_push_turn_run_wall_clock_skips_when_no_active_turn() -> None: +async def test_hydrate_skips_when_no_active_turn() -> None: bus = MagicMock() channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) mock_ws = AsyncMock() @@ -3557,12 +3535,12 @@ async def test_maybe_push_turn_run_wall_clock_skips_when_no_active_turn() -> Non from nanobot.session import webui_turns as wth wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear() - await channel._maybe_push_turn_run_wall_clock("chat-1") + await channel._outbound.hydrate("chat-1") mock_ws.send.assert_not_called() @pytest.mark.asyncio -async def test_maybe_push_turn_run_wall_clock_replays_running() -> None: +async def test_hydrate_replays_running_turn() -> None: bus = MagicMock() channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) mock_ws = AsyncMock() @@ -3572,7 +3550,7 @@ async def test_maybe_push_turn_run_wall_clock_replays_running() -> None: wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear() try: wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-1"] = 1_700_000_000.0 - await channel._maybe_push_turn_run_wall_clock("chat-1") + await channel._outbound.hydrate("chat-1") finally: wth._WEBSOCKET_TURN_WALL_STARTED_AT.pop("chat-1", None) diff --git a/nanobot/channels/websocket/tests/test_websocket_reconnect_idle.py b/nanobot/channels/websocket/tests/test_websocket_reconnect_idle.py index 951b7d7cd..4f69ec7ad 100644 --- a/nanobot/channels/websocket/tests/test_websocket_reconnect_idle.py +++ b/nanobot/channels/websocket/tests/test_websocket_reconnect_idle.py @@ -5,6 +5,8 @@ from unittest.mock import MagicMock, patch import pytest from nanobot.channels.websocket.runtime import WebSocketChannel +from nanobot.webui.outbound_projection import WebUIOutboundProjector +from nanobot.webui.session_projection import WebUISessionProjection @pytest.mark.asyncio @@ -13,7 +15,9 @@ async def test_hydrate_after_subscribe_is_quiet_when_no_turn_active(): channel = WebSocketChannel.__new__(WebSocketChannel) channel.gateway = MagicMock() channel.gateway.session_manager = MagicMock() - channel.gateway.session_manager.read_session_file = MagicMock(return_value={}) + channel.gateway.session_manager.read_session_metadata = MagicMock(return_value={}) + channel._session_projection = WebUISessionProjection(channel.gateway.session_manager) + channel._outbound = WebUIOutboundProjector(channel, channel._session_projection) channel._turn_models = {} sent_events = [] @@ -27,7 +31,7 @@ async def test_hydrate_after_subscribe_is_quiet_when_no_turn_active(): channel.send_goal_state = mock_send_goal_state channel.send_goal_status = mock_send_goal_status - with patch("nanobot.channels.websocket.runtime.websocket_turn_wall_started_at", return_value=None): + with patch("nanobot.webui.session_projection.websocket_turn_wall_started_at", return_value=None): await channel._hydrate_after_subscribe("test-chat") assert sent_events == [] @@ -39,7 +43,9 @@ async def test_hydrate_after_subscribe_pushes_running_when_turn_active(): channel = WebSocketChannel.__new__(WebSocketChannel) channel.gateway = MagicMock() channel.gateway.session_manager = MagicMock() - channel.gateway.session_manager.read_session_file = MagicMock(return_value={}) + channel.gateway.session_manager.read_session_metadata = MagicMock(return_value={}) + channel._session_projection = WebUISessionProjection(channel.gateway.session_manager) + channel._outbound = WebUIOutboundProjector(channel, channel._session_projection) channel._turn_models = {} sent_events = [] @@ -55,11 +61,11 @@ async def test_hydrate_after_subscribe_pushes_running_when_turn_active(): with ( patch( - "nanobot.channels.websocket.runtime.websocket_turn_wall_started_at", + "nanobot.webui.session_projection.websocket_turn_wall_started_at", return_value=1234567890.0, ), patch( - "nanobot.channels.websocket.runtime.websocket_turn_id", + "nanobot.webui.session_projection.websocket_turn_id", return_value="turn-active", ), ): diff --git a/nanobot/cli/tui_launcher.py b/nanobot/cli/tui_launcher.py index 934cde5db..a76d13e00 100644 --- a/nanobot/cli/tui_launcher.py +++ b/nanobot/cli/tui_launcher.py @@ -27,6 +27,7 @@ from nanobot.cli.webui_support import ( ) from nanobot.config.paths import get_data_dir from nanobot.config.schema import Config +from nanobot.webui.session_identity import is_webui_session_key, webui_chat_id if TYPE_CHECKING: from nanobot.gateway import GatewayClientLease @@ -486,8 +487,8 @@ def _tui_gateway_connection(config: Config) -> tuple[str, str]: 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 is_webui_session_key(session_id): + return webui_chat_id(session_id) if ":" in session_id: raise TuiSessionError( "the native TUI can open only WebSocket sessions; use --classic to resume " diff --git a/nanobot/session/recovery.py b/nanobot/session/recovery.py index b2016f0f6..b77a5b60c 100644 --- a/nanobot/session/recovery.py +++ b/nanobot/session/recovery.py @@ -28,6 +28,7 @@ from nanobot.session import turn_continuation from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata from nanobot.session.manager import Session, SessionManager from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY +from nanobot.webui.session_identity import webui_chat_id, webui_session_key RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint" PENDING_USER_TURN_KEY = "pending_user_turn" @@ -877,7 +878,7 @@ class RecoveryCoordinator: return state def _session_key(self, chat_id: str) -> str: - return UNIFIED_SESSION_KEY if self.unified_session else f"websocket:{chat_id}" + return UNIFIED_SESSION_KEY if self.unified_session else webui_session_key(chat_id) @staticmethod def _has_unfinished_webui_transcript(session_key: str) -> bool: @@ -929,9 +930,9 @@ class RecoveryCoordinator: session_key: str, metadata: Mapping[str, Any], ) -> tuple[str, str] | None: - if session_key.startswith("websocket:"): - chat_id = session_key.split(":", 1)[1] - return ("websocket", chat_id) if chat_id else None + chat_id = webui_chat_id(session_key) + if chat_id is not None: + return ("websocket", chat_id) if session_key == UNIFIED_SESSION_KEY: route = last_channel_from_metadata(metadata) if route and route[0] == "websocket": diff --git a/nanobot/session/webui_turns.py b/nanobot/session/webui_turns.py index c04ff52fc..ff7182995 100644 --- a/nanobot/session/webui_turns.py +++ b/nanobot/session/webui_turns.py @@ -56,6 +56,7 @@ from nanobot.webui.metadata import ( WEBSOCKET_TURN_OWNER_METADATA_KEY, WEBUI_TURN_METADATA_KEY, ) +from nanobot.webui.session_identity import is_webui_session_key from nanobot.webui.transcript import append_session_message_input WEBUI_SESSION_METADATA_KEY = "webui" @@ -628,7 +629,7 @@ class WebuiTurnCoordinator: event.context.channel != "system" or envelope is None or envelope["target_session_key"] != session_key - or not session_key.startswith("websocket:") + or not is_webui_session_key(session_key) ): return persisted = self.sessions.read_session_metadata(session_key) diff --git a/nanobot/webui/forking.py b/nanobot/webui/forking.py index 169777556..11e1f38f2 100644 --- a/nanobot/webui/forking.py +++ b/nanobot/webui/forking.py @@ -2,13 +2,15 @@ from __future__ import annotations -import re import uuid from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, TypeGuard +from typing import TYPE_CHECKING, Any, Protocol + +from loguru import logger from nanobot.session.manager import SessionManager from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title +from nanobot.webui.session_identity import is_valid_webui_chat_id, webui_session_key from nanobot.webui.transcript import ( append_fork_marker, delete_webui_transcript, @@ -19,13 +21,25 @@ from nanobot.webui.transcript import ( if TYPE_CHECKING: from websockets.asyncio.server import ServerConnection - from nanobot.channels.websocket.runtime import WebSocketChannel - -_WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") + from nanobot.webui.gateway_services import GatewayServices -def _valid_webui_chat_id(value: Any) -> TypeGuard[str]: - return isinstance(value, str) and _WEBUI_CHAT_ID_RE.match(value) is not None +class WebUIForkHost(Protocol): + gateway: GatewayServices + + async def send_webui_protocol_error( + self, + connection: ServerConnection, + detail: str, + ) -> None: ... + + async def attach_webui_fork( + self, + connection: ServerConnection, + *, + fork_id: str, + fork_key: str, + ) -> None: ... def create_webui_chat_fork( @@ -37,8 +51,8 @@ def create_webui_chat_fork( ) -> 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}" + source_key = webui_session_key(source_chat_id) + target_key = webui_session_key(new_id) try: forked = session_manager.fork_session_before_user_index( source_key, @@ -69,7 +83,7 @@ def create_webui_chat_fork( async def handle_webui_fork_chat( - channel: WebSocketChannel, + channel: WebUIForkHost, connection: ServerConnection, envelope: Mapping[str, Any], ) -> None: @@ -81,7 +95,7 @@ async def handle_webui_fork_chat( """ source_chat_id = envelope.get("source_chat_id") raw_index = envelope.get("before_user_index") - if not _valid_webui_chat_id(source_chat_id): + if not is_valid_webui_chat_id(source_chat_id): await channel.send_webui_protocol_error(connection, "invalid source_chat_id") return if isinstance(raw_index, bool) or not isinstance(raw_index, int) or raw_index < 0: @@ -105,7 +119,7 @@ async def handle_webui_fork_chat( return fork_id, fork_key = forked except Exception as exc: - channel.logger.warning("fork_chat failed: {}", exc) + logger.warning("fork_chat failed: {}", exc) await channel.send_webui_protocol_error(connection, "fork_chat_failed") return diff --git a/nanobot/webui/gateway_endpoint.py b/nanobot/webui/gateway_endpoint.py new file mode 100644 index 000000000..016e6c0f7 --- /dev/null +++ b/nanobot/webui/gateway_endpoint.py @@ -0,0 +1,114 @@ +"""HTTP and handshake composition for the WebUI gateway listener.""" + +from __future__ import annotations + +import hmac +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from websockets.asyncio.server import ServerConnection +from websockets.http11 import Request as WsRequest + +from nanobot.webui.gateway_tokens import GatewayTokenStore +from nanobot.webui.http_utils import ( + is_trusted_proxy_authenticated_request, + normalize_config_path, + parse_request_path, + query_first, +) +from nanobot.webui.ws_http import GatewayHTTPHandler + +if TYPE_CHECKING: + from nanobot.channels.websocket.runtime import WebSocketConfig + + +def is_websocket_upgrade(request: WsRequest) -> bool: + """Return whether a request contains a complete WebSocket upgrade handshake.""" + upgrade = request.headers.get("Upgrade") or request.headers.get("upgrade") + connection = request.headers.get("Connection") or request.headers.get("connection") + return bool( + upgrade + and "websocket" in upgrade.lower() + and connection + and "upgrade" in connection.lower() + ) + + +class WebUIGatewayEndpoint: + """Compose HTTP routing and WebSocket authentication on one listener.""" + + def __init__( + self, + *, + config: WebSocketConfig, + http: GatewayHTTPHandler, + tokens: GatewayTokenStore, + ) -> None: + self._config = config + self._http = http + self._tokens = tokens + self.webui_connections: set[ServerConnection] = set() + + async def process_request( + self, + connection: ServerConnection, + request: WsRequest, + *, + is_allowed: Callable[[str], bool], + ) -> Any: + """Route one listener request to a WS handshake or the HTTP application.""" + got, query = parse_request_path(request.path) + expected_ws = normalize_config_path(self._config.path) + if got == expected_ws and is_websocket_upgrade(request): + client_id = query_first(query, "client_id") or "" + if len(client_id) > 128: + client_id = client_id[:128] + if not is_allowed(client_id): + return connection.respond(403, "Forbidden") + return self.authorize_websocket_handshake(connection, query, request.headers) + return await self._http.dispatch(connection, request) + + def authorize_websocket_handshake( + self, + connection: ServerConnection, + query: dict[str, list[str]], + headers: Any = None, + ) -> Any: + """Authorize a WebSocket upgrade and remember trusted WebUI connections.""" + if is_trusted_proxy_authenticated_request(connection, headers or {}, self._config): + self.webui_connections.add(connection) + return None + + supplied = query_first(query, "token") + static_token = self._config.token.strip() + if static_token: + if supplied and hmac.compare_digest(supplied, static_token): + return None + if supplied and self.consume_issued_token(connection, supplied): + return None + return connection.respond(401, "Unauthorized") + + if self._config.websocket_requires_token: + if supplied and self.consume_issued_token(connection, supplied): + return None + return connection.respond(401, "Unauthorized") + + if supplied: + self.consume_issued_token(connection, supplied) + return None + + def consume_issued_token(self, connection: ServerConnection, token: str) -> bool: + """Consume one issued token and record its WebUI audience when present.""" + audience = self._tokens.take_issued_token_audience(token) + if audience == "webui": + self.webui_connections.add(connection) + return audience is not None + + def is_webui_connection(self, connection: ServerConnection) -> bool: + return connection in self.webui_connections + + def discard_connection(self, connection: ServerConnection) -> None: + self.webui_connections.discard(connection) + + def clear(self) -> None: + self.webui_connections.clear() diff --git a/nanobot/webui/gateway_services.py b/nanobot/webui/gateway_services.py index d9bf6b62d..33c83c9cc 100644 --- a/nanobot/webui/gateway_services.py +++ b/nanobot/webui/gateway_services.py @@ -10,9 +10,11 @@ from typing import TYPE_CHECKING, Any, Callable from loguru import logger as default_logger from nanobot.config.loader import get_config_path +from nanobot.webui.gateway_endpoint import WebUIGatewayEndpoint from nanobot.webui.gateway_tokens import GatewayTokenStore from nanobot.webui.ingress_policy import DEFAULT_WEBUI_INGRESS_POLICY, WebUIIngressPolicy from nanobot.webui.media_gateway import WebUIMediaGateway +from nanobot.webui.session_projection import WebUISessionProjection from nanobot.webui.settings_services import WebUISettingsServices from nanobot.webui.temporary_chats import WebUITemporaryChats from nanobot.webui.transcript import WebUITranscriptRecorder @@ -32,6 +34,7 @@ class GatewayServices: """Explicit dependencies shared by WebSocket transport and HTTP routes.""" http: GatewayHTTPHandler + endpoint: WebUIGatewayEndpoint settings: WebUISettingsServices tokens: GatewayTokenStore media: WebUIMediaGateway @@ -39,6 +42,7 @@ class GatewayServices: transcripts: WebUITranscriptRecorder workspaces: WebUIWorkspaceController temporary_chats: WebUITemporaryChats + session_projection: WebUISessionProjection session_manager: SessionManager | None cron_service: CronService | None local_trigger_store: LocalTriggerStore | None @@ -108,6 +112,7 @@ def build_gateway_services( workspaces=workspaces, logger=logger, ) + session_projection = WebUISessionProjection(session_manager, log=logger) http = GatewayHTTPHandler( config=config, session_manager=session_manager, @@ -135,8 +140,10 @@ def build_gateway_services( recovery_action=recovery_action, log=logger, ) + endpoint = WebUIGatewayEndpoint(config=config, http=http, tokens=tokens) return GatewayServices( http=http, + endpoint=endpoint, settings=settings, tokens=tokens, media=media, @@ -144,6 +151,7 @@ def build_gateway_services( transcripts=transcripts, workspaces=workspaces, temporary_chats=temporary_chats, + session_projection=session_projection, session_manager=session_manager, cron_service=cron_service, local_trigger_store=local_trigger_store, diff --git a/nanobot/webui/inbound_commands.py b/nanobot/webui/inbound_commands.py new file mode 100644 index 000000000..8416f3c4f --- /dev/null +++ b/nanobot/webui/inbound_commands.py @@ -0,0 +1,978 @@ +"""Application orchestration for typed WebUI WebSocket commands.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import re +import time +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Protocol, cast + +from loguru import logger +from websockets.asyncio.server import ServerConnection + +from nanobot.bus.events import INBOUND_META_USER_SHELL +from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn +from nanobot.runtime_context import ( + RUNTIME_CONTEXT_INPUT_META, + WEBUI_QUOTE_METADATA, + RuntimeContextBlock, + webui_quote_runtime_context, +) +from nanobot.security.workspace_access import ( + WORKSPACE_SCOPE_METADATA_KEY, + WorkspaceScopeError, +) +from nanobot.session.webui_turns import ( + clear_websocket_turn_if_current, + clear_websocket_turns, + register_queued_websocket_turn_if_idle, + websocket_turn_id, + 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 +from nanobot.webui.mcp_presets_api import normalize_mcp_preset_mentions +from nanobot.webui.metadata import WEBSOCKET_TURN_OWNER_METADATA_KEY +from nanobot.webui.session_access import ( + SessionMention, + WebuiSessionAccess, + session_mentions_runtime_context, +) +from nanobot.webui.session_identity import is_valid_webui_chat_id, webui_session_key +from nanobot.webui.sidebar_state import write_webui_sidebar_state +from nanobot.webui.temporary_chats import TemporaryChatError +from nanobot.webui.transcription_ws import webui_transcription_event + +_WEBUI_REQUEST_CACHE_TTL_S = 5 * 60.0 +_WEBUI_REQUEST_CACHE_MAX = 256 + + +@dataclass(frozen=True) +class WebUIRequestResult: + result: Any = None + status: int | None = None + message: str | None = None + + +@dataclass +class WebUIRequestOperation: + action: str + payload_digest: bytes + task: asyncio.Task[WebUIRequestResult] + completed_at: float | None = None + + +class WebUICommandTransport(Protocol): + """Typed transport capabilities consumed by WebUI command orchestration.""" + + def is_allowed(self, sender_id: str) -> bool: ... + + def webui_subscribers(self, chat_id: str) -> tuple[ServerConnection, ...]: ... + + def webui_connection_chats(self, connection: ServerConnection) -> tuple[str, ...]: ... + + def webui_attach(self, connection: ServerConnection, chat_id: str) -> None: ... + + def webui_detach(self, connection: ServerConnection, chat_id: str) -> None: ... + + def webui_clear_connection_default(self, connection: ServerConnection) -> None: ... + + def webui_clear_stream_buffers(self, chat_id: str) -> None: ... + + async def webui_hydrate(self, chat_id: str) -> None: ... + + async def webui_send_event( + self, + connection: ServerConnection, + event: str, + **fields: Any, + ) -> None: ... + + async def webui_send_raw( + self, + connection: ServerConnection, + raw: str, + *, + label: str = "", + ) -> None: ... + + async def webui_dispatch_message( + self, + *, + sender_id: str, + chat_id: str, + content: str, + media: list[str] | None, + metadata: dict[str, Any], + is_dm: bool, + session_key: str | None, + require_existing_session: bool, + ) -> None: ... + + async def send_session_updated( + self, + chat_id: str, + *, + scope: str | None = None, + ) -> None: ... + + +class WebUICommandRouter: + """Own WebUI command semantics while a transport host owns raw connections.""" + + def __init__(self, transport: WebUICommandTransport, gateway: GatewayServices) -> None: + self._transport = transport + self.gateway = gateway + self._http_router = gateway.http + self._media = gateway.media + self._ingress = gateway.ingress + self._transcripts = gateway.transcripts + self._workspaces = gateway.workspaces + self._temporary_chats = gateway.temporary_chats + self._session_projection = gateway.session_projection + self._webui_connections = gateway.endpoint.webui_connections + self._session_access = ( + WebuiSessionAccess(gateway.session_manager) + if gateway.session_manager is not None + else None + ) + self.request_tasks: dict[ + tuple[ServerConnection, str], + asyncio.Task[None], + ] = {} + self.request_operations: dict[str, WebUIRequestOperation] = {} + self.request_locks: dict[ServerConnection, asyncio.Lock] = {} + + def workspace_controls_available(self, connection: ServerConnection) -> bool: + return self._http_router.workspace_controls_available(connection) + + async def send_webui_protocol_error( + self, + connection: ServerConnection, + detail: str, + ) -> None: + await self._transport.webui_send_event(connection, "error", detail=detail) + + async def attach_webui_fork( + self, + connection: ServerConnection, + *, + fork_id: str, + fork_key: str, + ) -> None: + scope = self._workspaces.scope_for_session_key(fork_key) + self._transport.webui_attach(connection, fork_id) + await self._transport.webui_send_event( + connection, + "attached", + chat_id=fork_id, + **self._session_projection.attach_fields(fork_key), + ) + await self._transport.webui_send_event( + connection, + "session_updated", + chat_id=fork_id, + scope="metadata", + workspace_scope=scope.payload(), + ) + await self._transport.webui_hydrate(fork_id) + + async def discard_owned_chat( + self, + connection: ServerConnection, + chat_id: str, + ) -> None: + await self._temporary_chats.discard(connection, chat_id) + self._transport.webui_detach(connection, chat_id) + clear_websocket_turns(chat_id) + self._transport.webui_clear_stream_buffers(chat_id) + + async def cleanup_connection(self, connection: ServerConnection) -> None: + """Release command-owned state associated with one transport connection.""" + chat_ids = self._transport.webui_connection_chats(connection) + for chat_id in chat_ids: + if self._temporary_chats.owns(connection, chat_id): + await self.discard_owned_chat(connection, chat_id) + else: + self._transport.webui_detach(connection, chat_id) + for chat_id in self._temporary_chats.chat_ids_for_owner(connection): + await self.discard_owned_chat(connection, chat_id) + self._transport.webui_clear_connection_default(connection) + self.gateway.endpoint.discard_connection(connection) + self.discard_request_lock_if_idle(connection) + + async def broadcast_webui_event(self, event: str, **fields: Any) -> None: + for connection in tuple(self._webui_connections): + await self._transport.webui_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: + 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 self._transport.webui_subscribers(chat_id): + if connection is not origin: + await self._transport.webui_send_raw(connection, raw, label=" user_message ") + + async def workspace_scope_or_error( + self, + connection: ServerConnection, + resolver: Callable[[], Any], + *, + chat_id: str | None = None, + turn_id: str | None = None, + ) -> Any | None: + try: + return resolver() + except WorkspaceScopeError as exc: + await self._transport.webui_send_event( + connection, + "error", + detail="workspace_scope_rejected", + reason=exc.message, + **({"chat_id": chat_id} if chat_id else {}), + **({"turn_id": turn_id} if turn_id else {}), + ) + return None + + async def dispatch( + self, + connection: ServerConnection, + client_id: str, + envelope: dict[str, Any], + ) -> None: + """Execute one typed WebUI command.""" + command_type = envelope.get("type") + if command_type == "webui_request": + await self.start_webui_request(connection, envelope) + return + if command_type == "new_chat": + new_id = str(uuid.uuid4()) + scope = await self.workspace_scope_or_error( + connection, + lambda: self._workspaces.scope_for_new_chat( + envelope, + controls_available=self.workspace_controls_available(connection), + ), + ) + if scope is None: + return + self._workspaces.stage_scope(new_id, scope) + self._transport.webui_attach(connection, new_id) + await self._transport.webui_send_event( + connection, + "attached", + chat_id=new_id, + **self._session_projection.attach_fields(webui_session_key(new_id)), + ) + await self._transport.webui_send_event( + connection, + "session_updated", + chat_id=new_id, + scope="metadata", + workspace_scope=scope.payload(), + ) + await self._transport.webui_hydrate(new_id) + return + if command_type == "new_temporary_chat": + try: + new_id = self._temporary_chats.create( + connection, + trusted_webui=connection in self._webui_connections, + ) + except TemporaryChatError as exc: + await self._transport.webui_send_event(connection, "error", detail=exc.detail) + return + self._transport.webui_attach(connection, new_id) + await self._transport.webui_send_event( + connection, + "attached", + chat_id=new_id, + temporary=True, + ) + return + if command_type == "fork_chat": + await handle_webui_fork_chat(self, connection, envelope) + return + if command_type == "discard_temporary_chat": + chat_id = envelope.get("chat_id") + if not is_valid_webui_chat_id(chat_id): + await self._transport.webui_send_event( + connection, + "error", + detail="invalid temporary chat_id", + ) + return + try: + await self.discard_owned_chat(connection, chat_id) + except TemporaryChatError as exc: + await self._transport.webui_send_event( + connection, + "error", + detail=exc.detail, + chat_id=chat_id, + ) + return + if command_type == "attach": + chat_id = envelope.get("chat_id") + if not is_valid_webui_chat_id(chat_id): + await self._transport.webui_send_event( + connection, + "error", + detail="invalid chat_id", + ) + return + try: + self._temporary_chats.validate_attach(chat_id) + except TemporaryChatError as exc: + await self._transport.webui_send_event( + connection, + "error", + detail=exc.detail, + chat_id=chat_id, + ) + return + self._transport.webui_attach(connection, chat_id) + await self._transport.webui_send_event( + connection, + "attached", + chat_id=chat_id, + **self._session_projection.attach_fields(webui_session_key(chat_id)), + ) + await self._transport.webui_hydrate(chat_id) + return + if command_type == "set_sidebar_state": + if connection not in self._webui_connections: + await self._transport.webui_send_event(connection, "error", detail="access_denied") + return + state = envelope.get("state") + if not isinstance(state, dict): + await self._transport.webui_send_event( + connection, + "error", + detail="invalid_sidebar_state", + ) + return + try: + saved_state = await asyncio.to_thread( + write_webui_sidebar_state, + cast(dict[str, Any], state), + ) + except (OSError, ValueError): + await self._transport.webui_send_event( + connection, + "error", + detail="invalid_sidebar_state", + ) + return + await self.broadcast_webui_event("sidebar_state_updated", state=saved_state) + return + if command_type == "set_workspace_scope": + chat_id = envelope.get("chat_id") + if not is_valid_webui_chat_id(chat_id): + await self._transport.webui_send_event( + connection, + "error", + detail="invalid chat_id", + ) + return + try: + self._temporary_chats.validate_workspace_update(chat_id) + except TemporaryChatError as exc: + await self._transport.webui_send_event( + connection, + "error", + detail=exc.detail, + chat_id=chat_id, + ) + return + scope = await self.workspace_scope_or_error( + connection, + lambda: self._workspaces.scope_for_set_request( + envelope, + chat_id=chat_id, + chat_running=websocket_turn_wall_started_at(chat_id) is not None, + controls_available=self.workspace_controls_available(connection), + ), + chat_id=chat_id, + ) + if scope is None: + return + self._workspaces.stage_scope(chat_id, scope) + await self._transport.send_session_updated(chat_id, scope="metadata") + await self._transport.webui_send_event( + connection, + "session_updated", + chat_id=chat_id, + scope="metadata", + workspace_scope=scope.payload(), + ) + return + if command_type == "transcribe_audio": + event, payload = await webui_transcription_event( + envelope, + config_path=self.gateway.settings.config.path, + ) + await self._transport.webui_send_event(connection, event, **payload) + return + if command_type == "message": + await self._dispatch_message(connection, client_id, envelope) + return + await self._transport.webui_send_event( + connection, + "error", + detail=f"unknown type: {command_type!r}", + ) + + async def _dispatch_message( + self, + connection: ServerConnection, + client_id: str, + envelope: dict[str, Any], + ) -> None: + chat_id = envelope.get("chat_id") + content = envelope.get("content") + if not is_valid_webui_chat_id(chat_id): + await self._transport.webui_send_event(connection, "error", detail="invalid chat_id") + return + raw_turn_id = envelope.get("turn_id") + turn_id = raw_turn_id if isinstance(raw_turn_id, str) and raw_turn_id else None + rejection_fields = { + "chat_id": chat_id, + **({"turn_id": turn_id} if turn_id else {}), + } + if not self._transport.is_allowed(client_id): + await self._transport.webui_send_event( + connection, + "error", + detail="access_denied", + **rejection_fields, + ) + return + if not isinstance(content, str): + await self._transport.webui_send_event( + connection, + "error", + detail="missing content", + **rejection_fields, + ) + return + message_rejection = self._ingress.validate_text(content) + if message_rejection is not None: + await self._transport.webui_send_event( + connection, + "error", + detail="message_rejected", + reason=message_rejection, + **rejection_fields, + ) + return + + try: + temporary_policy = self._temporary_chats.message_policy( + connection, + chat_id, + content, + ) + except TemporaryChatError as exc: + await self._transport.webui_send_event( + connection, + "error", + detail=exc.detail, + **rejection_fields, + ) + return + + 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._transport.webui_send_event( + connection, + "error", + detail="attachment_rejected", + reason="malformed", + **rejection_fields, + ) + return + media_paths, reason = self._media.store_inbound_attachments( + cast(list[Any], raw_media) + ) + if reason is not None: + await self._transport.webui_send_event( + connection, + "error", + detail="attachment_rejected", + reason=reason, + **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, chat_id, media_paths) + + if not content.strip() and not media_paths: + await self._transport.webui_send_event( + connection, + "error", + detail="missing content", + **rejection_fields, + ) + return + self._transport.webui_attach(connection, chat_id) + if temporary_policy is None or temporary_policy.hydrate_transcript: + await self._transport.webui_hydrate(chat_id) + + scope = await self.workspace_scope_or_error( + connection, + lambda: ( + temporary_policy.workspace_scope + if temporary_policy is not None + else self._workspaces.scope_for_message( + envelope, + chat_id=chat_id, + chat_running=websocket_turn_wall_started_at(chat_id) is not None, + controls_available=self.workspace_controls_available(connection), + ) + ), + chat_id=chat_id, + turn_id=turn_id, + ) + if scope is None: + return + + if not self._transport.is_allowed(client_id): + await self._transport.webui_send_event( + connection, + "error", + detail="access_denied", + **rejection_fields, + ) + return + + metadata: dict[str, Any] = { + "remote": getattr(connection, "remote_address", None) + } + if envelope.get("webui") is True: + metadata["webui"] = True + metadata.update(self._transcripts.client_turn_metadata(envelope.get("turn_id"))) + trusted_webui = metadata.get("webui") is True and connection in self._webui_connections + is_user_shell = ( + trusted_webui + and envelope.get("user_shell") is True + and content.startswith("!") + ) + if is_user_shell: + metadata[INBOUND_META_USER_SHELL] = True + dispatch_content = ( + f"{USER_SHELL_COMMAND} {content[1:].lstrip()}" if is_user_shell else content + ) + cli_apps = normalize_cli_app_mentions(envelope.get("cli_apps")) + if cli_apps: + metadata["cli_apps"] = cli_apps + mcp_presets = normalize_mcp_preset_mentions( + envelope.get("mcp_presets"), + config_path=self.gateway.settings.config.path, + ) + if mcp_presets: + metadata["mcp_presets"] = mcp_presets + session_mentions: list[SessionMention] = [] + if trusted_webui and self._session_access is not None: + session_mentions = await asyncio.to_thread( + self._session_access.normalize_mentions, + envelope.get("session_mentions"), + exclude_session_key=webui_session_key(chat_id), + ) + if session_mentions: + metadata["session_mentions"] = session_mentions + metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata() + is_webui = metadata.get("webui") is True + queued_owner = None + if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content): + queued_owner = register_queued_websocket_turn_if_idle(chat_id, turn_id) + if queued_owner is not None: + metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] = queued_owner + + accepted = False + try: + if is_webui and ( + temporary_policy is None or temporary_policy.persist_transcript + ): + self._transcripts.append_user_message( + chat_id, + content, + metadata=metadata, + media_paths=media_paths or None, + cli_apps=cli_apps or None, + mcp_presets=mcp_presets or None, + session_mentions=session_mentions or None, + ) + if trusted_webui: + context_blocks: list[RuntimeContextBlock] = [] + quote = webui_quote_runtime_context( + {WEBUI_QUOTE_METADATA: envelope.get("quoted_context")} + ) + if quote is not None: + context_blocks.append(quote) + session_context = session_mentions_runtime_context(session_mentions) + if session_context is not None: + context_blocks.append(session_context) + if context_blocks: + metadata[RUNTIME_CONTEXT_INPUT_META] = context_blocks + await self._transport.webui_dispatch_message( + sender_id=client_id, + chat_id=chat_id, + content=dispatch_content, + media=media_paths or None, + metadata=metadata, + is_dm=False, + session_key=( + temporary_policy.session_key if temporary_policy is not None else None + ), + require_existing_session=( + temporary_policy.require_existing_session + if temporary_policy is not None + else False + ), + ) + self._workspaces.persist_scope(chat_id, scope) + accepted = True + finally: + if not accepted and queued_owner is not None: + clear_websocket_turn_if_current(chat_id, queued_owner) + + if is_webui: + await self.broadcast_user_message( + connection, + chat_id, + 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(chat_id) + started_at = websocket_turn_wall_started_at(chat_id) + await self._transport.webui_send_event( + connection, + "message_accepted", + chat_id=chat_id, + 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 {} + ), + ) + + async def start_webui_request( + self, + connection: ServerConnection, + envelope: dict[str, Any], + ) -> None: + request_id = envelope.get("request_id") + if not isinstance(request_id, str) or re.fullmatch( + r"[A-Za-z0-9._:-]{1,128}", + request_id, + ) is None: + await self._transport.webui_send_event( + connection, + "error", + detail="invalid webui request_id", + ) + return + if connection not in self._webui_connections: + await self.send_webui_response( + connection, + request_id, + status=403, + message="access_denied", + ) + return + + action = envelope.get("action") + payload = envelope.get("payload") + if not isinstance(action, str) or re.fullmatch( + r"[a-z][a-z0-9_.]{0,127}", + action, + ) is None: + await self.send_webui_response( + connection, + request_id, + status=400, + message="invalid WebUI mutation action", + ) + return + if not isinstance(payload, dict): + await self.send_webui_response( + connection, + request_id, + status=400, + message="WebUI mutation payload must be an object", + ) + return + + payload_digest = hashlib.sha256( + json.dumps( + payload, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).digest() + self.prune_request_operations() + operation = self.request_operations.get(request_id) + is_replay = operation is not None + if operation is not None and ( + operation.action != action or operation.payload_digest != payload_digest + ): + await self.send_webui_response( + connection, + request_id, + status=409, + message="request_id was already used for a different WebUI mutation", + ) + return + if operation is None: + operation_task = asyncio.create_task( + self.execute_webui_request( + connection, + action, + cast(dict[str, Any], payload), + ) + ) + new_operation = WebUIRequestOperation( + action=action, + payload_digest=payload_digest, + task=operation_task, + ) + operation = new_operation + self.request_operations[request_id] = new_operation + + def mark_complete(_task: asyncio.Task[WebUIRequestResult]) -> None: + current = self.request_operations.get(request_id) + if current is not new_operation: + return + new_operation.completed_at = time.monotonic() + self.prune_request_operations() + + operation_task.add_done_callback(mark_complete) + + key = (connection, request_id) + if key in self.request_tasks: + return + delivery_task = asyncio.create_task( + self.deliver_webui_request( + connection, + request_id, + operation.task, + sequence=is_replay, + ) + ) + self.request_tasks[key] = delivery_task + + def prune_request_operations(self) -> None: + now = time.monotonic() + for request_id, operation in tuple(self.request_operations.items()): + if ( + operation.completed_at is not None + and now - operation.completed_at >= _WEBUI_REQUEST_CACHE_TTL_S + ): + self.request_operations.pop(request_id, None) + + completed = sorted( + ( + (operation.completed_at, request_id) + for request_id, operation in self.request_operations.items() + if operation.completed_at is not None + ), + key=lambda item: item[0], + ) + for _, request_id in completed[:-_WEBUI_REQUEST_CACHE_MAX]: + self.request_operations.pop(request_id, None) + + def discard_request_lock_if_idle(self, connection: ServerConnection) -> None: + if connection in self._webui_connections: + return + if any(task_connection is connection for task_connection, _ in self.request_tasks): + return + self.request_locks.pop(connection, None) + + async def deliver_webui_request( + self, + connection: ServerConnection, + request_id: str, + operation_task: asyncio.Task[WebUIRequestResult], + *, + sequence: bool = False, + ) -> None: + try: + if sequence: + lock = self.request_locks.setdefault(connection, asyncio.Lock()) + async with lock: + result = await asyncio.shield(operation_task) + await self.send_webui_response( + connection, + request_id, + result=result.result, + status=result.status, + message=result.message, + ) + return + result = await asyncio.shield(operation_task) + await self.send_webui_response( + connection, + request_id, + result=result.result, + status=result.status, + message=result.message, + ) + finally: + self.request_tasks.pop((connection, request_id), None) + self.discard_request_lock_if_idle(connection) + + async def execute_webui_request( + self, + connection: ServerConnection, + action: str, + payload: dict[str, Any], + ) -> WebUIRequestResult: + try: + lock = self.request_locks.setdefault(connection, asyncio.Lock()) + async with lock: + response = await self._http_router.dispatch_webui_mutation( + connection, + action, + payload, + ) + status = response.status_code + body = bytes(response.body).decode("utf-8", errors="replace").strip() + if 200 <= status < 300: + try: + result = json.loads(body) + except json.JSONDecodeError: + return WebUIRequestResult( + status=502, + message="WebUI mutation returned an invalid response", + ) + if action == "sidebar.update" and isinstance(result, dict): + await self.broadcast_webui_event( + "sidebar_state_updated", + state=result, + ) + return WebUIRequestResult(result=result) + return WebUIRequestResult( + status=status, + message=body or response.reason_phrase, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("WebUI mutation '{}' failed", action) + return WebUIRequestResult( + status=500, + message="WebUI mutation failed", + ) + + async def send_webui_response( + self, + connection: ServerConnection, + request_id: str, + *, + result: Any = None, + status: int | None = None, + message: str | None = None, + ) -> None: + if status is None: + await self._transport.webui_send_event( + connection, + "webui_response", + request_id=request_id, + ok=True, + result=result, + ) + return + await self._transport.webui_send_event( + connection, + "webui_response", + request_id=request_id, + ok=False, + error={ + "status": status, + "message": message or "WebUI mutation failed", + }, + ) + + async def close(self) -> None: + """Cancel command work and release application-owned gateway state.""" + delivery_tasks = tuple(self.request_tasks.values()) + operation_tasks = tuple(operation.task for operation in self.request_operations.values()) + for task in (*delivery_tasks, *operation_tasks): + task.cancel() + if delivery_tasks: + await asyncio.gather(*delivery_tasks, return_exceptions=True) + if operation_tasks: + await asyncio.gather(*operation_tasks, return_exceptions=True) + self.request_tasks.clear() + self.request_locks.clear() + self.request_operations.clear() + self.gateway.tokens.clear() + self.gateway.endpoint.clear() + self._temporary_chats.close() diff --git a/nanobot/webui/outbound_projection.py b/nanobot/webui/outbound_projection.py new file mode 100644 index 000000000..89ba293c3 --- /dev/null +++ b/nanobot/webui/outbound_projection.py @@ -0,0 +1,245 @@ +"""Project agent runtime events onto the WebUI wire protocol.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol + +from loguru import logger + +from nanobot.bus.events import OutboundMessage +from nanobot.bus.outbound_events import ( + GoalStateSyncEvent, + GoalStatusEvent, + ProgressEvent, + RecoveryStateEvent, + RuntimeModelUpdatedEvent, + SessionUpdatedEvent, + TurnEndEvent, + TurnModelUpdatedEvent, + UserInputEvent, + outbound_event_from_message, +) +from nanobot.session.webui_turns import clear_websocket_turn_if_current +from nanobot.webui.metadata import ( + WEBSOCKET_TURN_OWNER_METADATA_KEY, + WEBUI_SYSTEM_COMMAND_TURN_PREFIX, + WEBUI_TURN_METADATA_KEY, +) +from nanobot.webui.session_identity import webui_session_key +from nanobot.webui.session_projection import WebUISessionProjection + +if TYPE_CHECKING: + from websockets.asyncio.server import ServerConnection + + from nanobot.providers.base import LLMUsage + + +class WebUIOutboundTransport(Protocol): + """Wire operations required by the outbound application projector.""" + + def webui_subscribers(self, chat_id: str) -> tuple[ServerConnection, ...]: ... + + async def send_runtime_model_updated( + self, + *, + model_name: str | None, + model_preset: str | None = None, + ) -> None: ... + + async def send_turn_model_updated( + self, + chat_id: str, + *, + model_name: str, + model_preset: str | None = None, + context_window_tokens: int | None = None, + fallback: bool = False, + ) -> None: ... + + async def send_user_input( + self, + chat_id: str, + *, + content: str, + created_at_ms: int, + provenance: dict[str, Any], + ) -> None: ... + + async def send_recovery_state(self, chat_id: str, event: RecoveryStateEvent) -> None: ... + + async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None: ... + + async def send_goal_status( + self, + chat_id: str, + status: str, + *, + started_at: float | None = None, + turn_id: str | None = None, + ) -> None: ... + + async def send_turn_end( + self, + chat_id: str, + latency_ms: int | None = None, + *, + goal_state: dict[str, Any] | None = None, + usage: LLMUsage | None = None, + context_window_tokens: int | None = None, + metadata: dict[str, Any] | None = None, + turn_owner: str | None = None, + ) -> None: ... + + async def send_session_updated(self, chat_id: str, *, scope: str | None = None) -> None: ... + + async def send_file_edit_events( + self, + chat_id: str, + edits: list[dict[str, Any]], + metadata: dict[str, Any] | None = None, + ) -> None: ... + + async def send_projected_message( + self, + msg: OutboundMessage, + progress_event: ProgressEvent | None, + ) -> None: ... + + +class WebUIOutboundProjector: + """Interpret runtime events without coupling that state machine to the channel.""" + + def __init__( + self, + transport: WebUIOutboundTransport, + session_projection: WebUISessionProjection, + ) -> None: + self._transport = transport + self._session_projection = session_projection + + async def hydrate(self, chat_id: str) -> None: + """Replay reconnect state through the existing stable wire operations.""" + for event in self._session_projection.hydration_events( + webui_session_key(chat_id), + chat_id, + ): + if event["event"] == "goal_state": + await self._transport.send_goal_state(chat_id, event["goal_state"]) + continue + await self._transport.send_goal_status( + chat_id, + "running", + started_at=event["started_at"], + turn_id=event.get("turn_id"), + ) + + async def send(self, msg: OutboundMessage) -> None: + event = outbound_event_from_message(msg) + progress_event = event if isinstance(event, ProgressEvent) else None + if isinstance(event, RuntimeModelUpdatedEvent): + await self._transport.send_runtime_model_updated( + model_name=event.model, + model_preset=event.model_preset, + ) + return + + conns = list(self._transport.webui_subscribers(msg.chat_id)) + if not conns: + quiet_events = ( + ProgressEvent, + UserInputEvent, + TurnEndEvent, + SessionUpdatedEvent, + GoalStatusEvent, + GoalStateSyncEvent, + ) + log = ( + logger.debug + if isinstance(event, quiet_events) + else logger.warning + ) + log("no active subscribers for chat_id={}", msg.chat_id) + + if isinstance(event, TurnModelUpdatedEvent): + if conns: + await self._transport.send_turn_model_updated( + msg.chat_id, + model_name=event.model, + model_preset=event.model_preset, + context_window_tokens=event.context_window_tokens, + fallback=event.fallback, + ) + return + if isinstance(event, UserInputEvent): + if conns: + await self._transport.send_user_input( + msg.chat_id, + content=event.content, + created_at_ms=event.created_at_ms, + provenance=event.provenance, + ) + return + if isinstance(event, RecoveryStateEvent): + if conns: + await self._transport.send_recovery_state(msg.chat_id, event) + return + if isinstance(event, GoalStateSyncEvent): + if conns: + await self._transport.send_goal_state( + msg.chat_id, + event.goal_state or {"active": False}, + ) + return + if isinstance(event, GoalStatusEvent): + turn_id = (msg.metadata or {}).get(WEBUI_TURN_METADATA_KEY) + current_turn_id = turn_id if isinstance(turn_id, str) else None + turn_owner = (msg.metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY) + current_turn_owner = turn_owner if isinstance(turn_owner, str) else None + try: + if conns and event.status in ("running", "idle"): + await self._transport.send_goal_status( + msg.chat_id, + event.status, + started_at=event.started_at, + turn_id=current_turn_id, + ) + finally: + if event.status == "idle": + clear_websocket_turn_if_current( + msg.chat_id, + current_turn_owner, + preserve_persistence_failure=True, + ) + return + if isinstance(event, TurnEndEvent): + turn_id = (msg.metadata or {}).get(WEBUI_TURN_METADATA_KEY) + session_update_scope = ( + "metadata" + if isinstance(turn_id, str) + and turn_id.startswith(WEBUI_SYSTEM_COMMAND_TURN_PREFIX) + else "thread" + ) + turn_owner = (msg.metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY) + await self._transport.send_turn_end( + msg.chat_id, + latency_ms=event.latency_ms, + goal_state=event.goal_state, + usage=event.usage, + context_window_tokens=event.context_window_tokens, + metadata=msg.metadata, + turn_owner=turn_owner if isinstance(turn_owner, str) else None, + ) + await self._transport.send_session_updated(msg.chat_id, scope=session_update_scope) + return + if isinstance(event, SessionUpdatedEvent): + if conns: + await self._transport.send_session_updated(msg.chat_id, scope=event.scope) + return + if progress_event and progress_event.file_edit_events: + await self._transport.send_file_edit_events( + msg.chat_id, + progress_event.file_edit_events, + msg.metadata, + ) + return + await self._transport.send_projected_message(msg, progress_event) diff --git a/nanobot/webui/session_identity.py b/nanobot/webui/session_identity.py new file mode 100644 index 000000000..5f9bea6ef --- /dev/null +++ b/nanobot/webui/session_identity.py @@ -0,0 +1,32 @@ +"""Stable mapping between public WebUI chat IDs and persisted session keys.""" + +from __future__ import annotations + +import re +from typing import Any, TypeGuard + +WEBUI_SESSION_STORAGE_PREFIX = "websocket:" +_WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") + + +def is_valid_webui_chat_id(value: Any) -> TypeGuard[str]: + """Validate the compact chat IDs accepted by the WebUI protocol.""" + return isinstance(value, str) and _WEBUI_CHAT_ID_RE.fullmatch(value) is not None + + +def webui_session_key(chat_id: str) -> str: + """Return the backward-compatible persisted key for a WebUI chat.""" + return f"{WEBUI_SESSION_STORAGE_PREFIX}{chat_id}" + + +def is_webui_session_key(session_key: str) -> bool: + """Return whether *session_key* belongs to the WebUI session namespace.""" + return session_key.startswith(WEBUI_SESSION_STORAGE_PREFIX) + + +def webui_chat_id(session_key: str) -> str | None: + """Extract a non-empty WebUI chat ID from a persisted session key.""" + if not is_webui_session_key(session_key): + return None + chat_id = session_key.removeprefix(WEBUI_SESSION_STORAGE_PREFIX) + return chat_id or None diff --git a/nanobot/webui/session_list_index.py b/nanobot/webui/session_list_index.py index 50612cfde..e59a3aef3 100644 --- a/nanobot/webui/session_list_index.py +++ b/nanobot/webui/session_list_index.py @@ -32,6 +32,12 @@ from nanobot.session.manager import ( ) from nanobot.session.model_selection import model_preset_from_metadata from nanobot.session.recovery import recovery_state_from_metadata +from nanobot.webui.session_identity import ( + WEBUI_SESSION_STORAGE_PREFIX, + is_webui_session_key, + webui_chat_id, + webui_session_key, +) _INDEX_VERSION = 8 _INDEX_FILENAME = ".webui_session_index.json" @@ -50,7 +56,7 @@ _WEBUI_ACTIVITY_MTIME_NS = "webui_activity_mtime_ns" _WEBUI_ACTIVITY_SIZE = "webui_activity_size" _WEBUI_ACTIVITY_FILES = "webui_activity_files" _VISIBLE_TRANSCRIPT_ROLES = {"user", "assistant"} -_WEBUI_SESSION_STEM_PREFIX = SessionManager.safe_key("websocket:") +_WEBUI_SESSION_STEM_PREFIX = SessionManager.safe_key(WEBUI_SESSION_STORAGE_PREFIX) _WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") _TRANSCRIPT_SEGMENTS_SUFFIX = ".segments" _TRANSCRIPT_NON_ANSWER_KINDS = {"progress", "reasoning", "tool_hint"} @@ -90,7 +96,7 @@ def _reconcile_index(session_manager: SessionManager) -> tuple[list[dict[str, An session_keys_by_stem = { SessionManager.safe_key(key): key for key in session_paths - if key.startswith("websocket:") + if is_webui_session_key(key) } rows: list[dict[str, Any]] = [] changed = existing_rows is None @@ -375,9 +381,9 @@ def _transcript_record(line: str) -> dict[str, Any] | None: def _valid_transcript_session_key(key: str, stem: str) -> bool: - if not key.startswith("websocket:"): + chat_id = webui_chat_id(key) + if chat_id is None: return False - chat_id = key.split(":", 1)[1] return _WEBUI_CHAT_ID_RE.fullmatch(chat_id) is not None and SessionManager.safe_key(key) == stem @@ -535,7 +541,9 @@ def _scan_transcript_row( paths: tuple[Path, ...], webui_dir: Path, ) -> dict[str, Any] | None: - path_key = session_key or f"websocket:{stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)}" + path_key = session_key or webui_session_key( + stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX) + ) signature = _webui_activity_signature(path_key, webui_dir) activity_updated_at = _webui_activity_updated_at(signature) if activity_updated_at is None: @@ -560,7 +568,7 @@ def _scan_transcript_row( saw_record = True chat_id = record.get("chat_id") if isinstance(chat_id, str) and chat_id.strip(): - candidate = f"websocket:{chat_id.strip()}" + candidate = webui_session_key(chat_id.strip()) if _valid_transcript_session_key(candidate, stem): session_key = candidate if created_at is None: @@ -586,7 +594,7 @@ def _scan_transcript_row( if not saw_record: return None if session_key is None: - fallback = f"websocket:{stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)}" + fallback = webui_session_key(stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)) if not _valid_transcript_session_key(fallback, stem): return None session_key = fallback diff --git a/nanobot/webui/session_projection.py b/nanobot/webui/session_projection.py new file mode 100644 index 000000000..d132d4394 --- /dev/null +++ b/nanobot/webui/session_projection.py @@ -0,0 +1,103 @@ +"""WebUI session read models exposed to interactive clients.""" + +from __future__ import annotations + +from typing import Any, Protocol, cast + +from loguru import logger as default_logger + +from nanobot.providers.base import LLMUsage +from nanobot.session.goal_state import goal_state_ws_blob +from nanobot.session.model_selection import model_preset_from_metadata +from nanobot.session.recovery import recovery_state_from_metadata +from nanobot.session.webui_turns import websocket_turn_id, websocket_turn_wall_started_at + + +class SessionMetadataReader(Protocol): + """Narrow persisted-session dependency used by WebUI projections.""" + + def read_session_metadata(self, key: str) -> dict[str, Any] | None: ... + + +class WebUISessionProjection: + """Project persisted session metadata into stable WebUI protocol fields.""" + + def __init__( + self, + sessions: SessionMetadataReader | None, + *, + log: Any = default_logger, + ) -> None: + self._sessions = sessions + self._log = log + + def attach_fields(self, session_key: str) -> dict[str, Any]: + """Return the session runtime facts sent with an attach handshake.""" + if self._sessions is None: + return {} + snapshot = self._sessions.read_session_metadata(session_key) + raw_metadata = snapshot.get("metadata") if snapshot is not None else None + metadata = cast(dict[str, object], raw_metadata) if isinstance(raw_metadata, dict) else None + + fields: dict[str, Any] = {} + try: + fields["model_preset"] = model_preset_from_metadata(metadata) + except ValueError: + self._log.warning("ignoring invalid model preset metadata for session_key={}", session_key) + fields["model_preset"] = None + if metadata is None: + return fields + + recovery_state = recovery_state_from_metadata(metadata) + if recovery_state is not None: + fields["recovery_state"] = recovery_state + usage = LLMUsage.from_dict(metadata.get("_last_usage")) + if usage is not None: + fields["usage"] = usage.to_turn_dict() + return fields + + def hydration_events(self, session_key: str, chat_id: str) -> tuple[dict[str, Any], ...]: + """Return reconnect events for durable and same-process session state.""" + events: list[dict[str, Any]] = [] + goal_state = self.persisted_goal_state(session_key) + if goal_state is not None: + events.append( + { + "event": "goal_state", + "chat_id": chat_id, + "goal_state": goal_state, + } + ) + active_turn = self.active_turn_status(chat_id) + if active_turn is not None: + started_at, turn_id = active_turn + event: dict[str, Any] = { + "event": "goal_status", + "chat_id": chat_id, + "status": "running", + "started_at": started_at, + } + if turn_id is not None: + event["turn_id"] = turn_id + events.append(event) + return tuple(events) + + def persisted_goal_state(self, session_key: str) -> dict[str, Any] | None: + """Return an actionable persisted goal state for reconnect hydration.""" + if self._sessions is None: + return None + snapshot = self._sessions.read_session_metadata(session_key) + raw_metadata = snapshot.get("metadata") if snapshot is not None else None + metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {} + goal_state = goal_state_ws_blob(metadata) + if not goal_state.get("active") and goal_state.get("status") != "blocked": + return None + return goal_state + + @staticmethod + def active_turn_status(chat_id: str) -> tuple[float, str | None] | None: + """Return same-process running-turn state for reconnect hydration.""" + started_at = websocket_turn_wall_started_at(chat_id) + if started_at is None: + return None + return started_at, websocket_turn_id(chat_id) diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index 43906be90..22f6c7684 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -23,6 +23,7 @@ from nanobot.session.automation_turns import is_automation_kind from nanobot.session.history_visibility import is_hidden_history_message from nanobot.session.manager import SessionManager from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY +from nanobot.webui.session_identity import webui_chat_id, webui_session_key WEBUI_TRANSCRIPT_SCHEMA_VERSION = 3 WEBUI_FORK_MARKER_EVENT = "fork_marker" @@ -828,7 +829,7 @@ class WebUITranscriptRecorder: def append(self, chat_id: str, event: dict[str, Any]) -> bool: try: dup = json.loads(json.dumps(event, ensure_ascii=False)) - append_transcript_object(f"websocket:{chat_id}", dup) + append_transcript_object(webui_session_key(chat_id), dup) except (OSError, ValueError, TypeError) as e: self._log.warning("webui transcript append failed: {}", e) return False @@ -860,10 +861,10 @@ class WebUITranscriptRecorder: def _chat_id_from_session_key(session_key: str) -> str | None: - if not session_key.startswith("websocket:"): + chat_id = webui_chat_id(session_key) + if chat_id is None: return None - chat_id = session_key.split(":", 1)[1].strip() - return chat_id or None + return chat_id.strip() or None def _is_user_transcript_row(row: dict[str, Any]) -> bool: diff --git a/nanobot/webui/workspaces.py b/nanobot/webui/workspaces.py index 4d0a853b1..6c1d393c8 100644 --- a/nanobot/webui/workspaces.py +++ b/nanobot/webui/workspaces.py @@ -20,6 +20,7 @@ from nanobot.security.workspace_access import ( default_workspace_scope, validate_workspace_scope_payload, ) +from nanobot.webui.session_identity import webui_session_key if TYPE_CHECKING: from nanobot.session.manager import SessionManager @@ -309,7 +310,7 @@ class WebUIWorkspaceController: raise WorkspaceScopeError("chat_running", status=409) return self.scope_from_envelope( envelope, - session_key=f"websocket:{chat_id}", + session_key=webui_session_key(chat_id), controls_available=controls_available, ) @@ -323,19 +324,19 @@ class WebUIWorkspaceController: ) -> WorkspaceScope: scope = self.scope_from_envelope( envelope, - session_key=f"websocket:{chat_id}", + session_key=webui_session_key(chat_id), controls_available=controls_available, ) if ( WORKSPACE_SCOPE_METADATA_KEY in envelope and chat_running - and scope.metadata() != self.scope_for_session_key(f"websocket:{chat_id}").metadata() + and scope.metadata() != self.scope_for_session_key(webui_session_key(chat_id)).metadata() ): raise WorkspaceScopeError("chat_running", status=409) return scope def persist_scope(self, chat_id: str, scope: WorkspaceScope) -> None: - session_key = f"websocket:{chat_id}" + session_key = webui_session_key(chat_id) if self._sessions is not None: session = self._sessions.get_or_create(session_key) session.metadata["webui"] = True @@ -345,7 +346,7 @@ class WebUIWorkspaceController: def stage_scope(self, chat_id: str, scope: WorkspaceScope) -> None: """Keep a new chat's scope transient until its first accepted message.""" - session_key = f"websocket:{chat_id}" + session_key = webui_session_key(chat_id) if ( self._sessions is not None and self._sessions.read_session_metadata(session_key) is not None diff --git a/nanobot/webui/ws_http.py b/nanobot/webui/ws_http.py index eb64e778b..16d2f80d4 100644 --- a/nanobot/webui/ws_http.py +++ b/nanobot/webui/ws_http.py @@ -103,6 +103,7 @@ from nanobot.webui.session_automations import ( session_automations_payload, ) from nanobot.webui.session_context import session_context_payload +from nanobot.webui.session_identity import is_webui_session_key from nanobot.webui.session_list_index import ( WEBUI_SESSION_INDEX_INTERNAL_FIELDS, indexed_workspace_scope, @@ -774,7 +775,7 @@ class GatewayHTTPHandler: default_scope: WorkspaceScope | None = None for s in sessions: key = s.get("key") - if not (isinstance(key, str) and key.startswith("websocket:")): + if not (isinstance(key, str) and is_webui_session_key(key)): continue row = { k: v @@ -1619,4 +1620,4 @@ def _positive_int(value: Any) -> int | None: def _is_websocket_channel_session_key(key: str) -> bool: - return key.startswith("websocket:") + return is_webui_session_key(key) diff --git a/tests/channels/test_websocket_application_boundary.py b/tests/channels/test_websocket_application_boundary.py new file mode 100644 index 000000000..014cd09cd --- /dev/null +++ b/tests/channels/test_websocket_application_boundary.py @@ -0,0 +1,96 @@ +"""Executable architecture constraints for the WebSocket transport adapter.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +from nanobot.channels.websocket import runtime + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +_RUNTIME_PATH = _REPOSITORY_ROOT / "nanobot" / "channels" / "websocket" / "runtime.py" +_SESSION_IDENTITY_PATH = _REPOSITORY_ROOT / "nanobot" / "webui" / "session_identity.py" +_FORBIDDEN_RUNTIME_IMPORTS = ( + "nanobot.bus.outbound_events", + "nanobot.command", + "nanobot.runtime_context", + "nanobot.security.workspace_access", + "nanobot.session.goal_state", + "nanobot.webui.cli_apps_api", + "nanobot.webui.forking", + "nanobot.webui.mcp_presets_api", + "nanobot.webui.sidebar_state", + "nanobot.webui.transcription_ws", +) + + +def _channel_method(name: str) -> ast.AsyncFunctionDef: + tree = ast.parse(_RUNTIME_PATH.read_text(encoding="utf-8")) + channel = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "WebSocketChannel" + ) + return next( + node + for node in channel.body + if isinstance(node, ast.AsyncFunctionDef) and node.name == name + ) + + +def _statements_without_docstring(node: ast.AsyncFunctionDef) -> list[ast.stmt]: + body = list(node.body) + if ( + body + and isinstance(body[0], ast.Expr) + and isinstance(body[0].value, ast.Constant) + and isinstance(body[0].value.value, str) + ): + body.pop(0) + return body + + +def test_websocket_runtime_does_not_import_application_command_trees() -> None: + tree = ast.parse(_RUNTIME_PATH.read_text(encoding="utf-8")) + imported = { + node.module + for node in tree.body + if isinstance(node, ast.ImportFrom) and node.module is not None + } + violations = sorted( + module + for module in imported + if module.startswith(_FORBIDDEN_RUNTIME_IMPORTS) + ) + assert violations == [] + + +def test_business_entrypoints_are_thin_transport_delegations() -> None: + for method_name in ("_dispatch_envelope", "_hydrate_after_subscribe", "send"): + statements = _statements_without_docstring(_channel_method(method_name)) + assert len(statements) == 1, method_name + assert isinstance(statements[0], ast.Expr), method_name + assert isinstance(statements[0].value, ast.Await), method_name + + statements = _statements_without_docstring(_channel_method("_dispatch_http")) + assert len(statements) == 1 + assert isinstance(statements[0], ast.Return) + assert isinstance(statements[0].value, ast.Await) + + +def test_persisted_webui_session_prefix_has_one_production_owner() -> None: + owners = [] + for path in (_REPOSITORY_ROOT / "nanobot").rglob("*.py"): + if "tests" in path.parts or path == _SESSION_IDENTITY_PATH: + continue + if "websocket:" in path.read_text(encoding="utf-8"): + owners.append(path.relative_to(_REPOSITORY_ROOT).as_posix()) + assert owners == [] + assert 'WEBUI_SESSION_STORAGE_PREFIX = "websocket:"' in _SESSION_IDENTITY_PATH.read_text( + encoding="utf-8" + ) + + +def test_runtime_exports_compatibility_protocol_helpers() -> None: + assert runtime._is_valid_chat_id("unified:default") # pyright: ignore[reportPrivateUsage] + assert not runtime._is_valid_chat_id("../escape") # pyright: ignore[reportPrivateUsage] diff --git a/tests/webui/test_forking.py b/tests/webui/test_forking.py index b99a4a40e..73a04d19c 100644 --- a/tests/webui/test_forking.py +++ b/tests/webui/test_forking.py @@ -139,8 +139,9 @@ async def test_fork_handler_maps_invalid_source_and_internal_failure_to_stable_e channel = SimpleNamespace( send_webui_protocol_error=AsyncMock(), gateway=SimpleNamespace(session_manager=MagicMock()), - logger=SimpleNamespace(warning=MagicMock()), ) + warning = MagicMock() + monkeypatch.setattr(forking, "logger", SimpleNamespace(warning=warning)) envelope = {"source_chat_id": "source", "before_user_index": 0} monkeypatch.setattr(forking, "create_webui_chat_fork", lambda *_args, **_kwargs: None) @@ -158,5 +159,5 @@ async def test_fork_handler_maps_invalid_source_and_internal_failure_to_stable_e ) await forking.handle_webui_fork_chat(channel, connection, envelope) - channel.logger.warning.assert_called_once_with("fork_chat failed: {}", ANY) + warning.assert_called_once_with("fork_chat failed: {}", ANY) channel.send_webui_protocol_error.assert_awaited_once_with(connection, "fork_chat_failed") diff --git a/tests/webui/test_session_identity.py b/tests/webui/test_session_identity.py new file mode 100644 index 000000000..84646e933 --- /dev/null +++ b/tests/webui/test_session_identity.py @@ -0,0 +1,24 @@ +from nanobot.webui.session_identity import ( + WEBUI_SESSION_STORAGE_PREFIX, + is_valid_webui_chat_id, + is_webui_session_key, + webui_chat_id, + webui_session_key, +) + + +def test_webui_session_identity_preserves_persisted_wire_compatibility() -> None: + assert WEBUI_SESSION_STORAGE_PREFIX == "websocket:" + assert webui_session_key("chat-1") == "websocket:chat-1" + assert is_webui_session_key("websocket:chat-1") + assert webui_chat_id("websocket:chat-1") == "chat-1" + assert webui_chat_id("websocket: chat-1") == " chat-1" + assert webui_chat_id("websocket:") is None + assert webui_chat_id("telegram:chat-1") is None + + +def test_webui_chat_id_validation_is_protocol_scoped() -> None: + assert is_valid_webui_chat_id("unified:default") + assert is_valid_webui_chat_id("x" * 64) + assert not is_valid_webui_chat_id("x" * 65) + assert not is_valid_webui_chat_id("../escape") diff --git a/tests/webui/test_session_projection.py b/tests/webui/test_session_projection.py new file mode 100644 index 000000000..abe110583 --- /dev/null +++ b/tests/webui/test_session_projection.py @@ -0,0 +1,114 @@ +from unittest.mock import MagicMock + +import pytest + +from nanobot.providers.base import LLMUsage +from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY +from nanobot.session.recovery import RECOVERY_METADATA_KEY +from nanobot.webui.session_projection import WebUISessionProjection + + +def test_attach_fields_restore_session_runtime_metadata() -> None: + usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175) + sessions = MagicMock() + sessions.read_session_metadata.return_value = { + "metadata": { + SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research", + RECOVERY_METADATA_KEY: { + "status": "recovered", + "recovery_id": "recovery-1", + "reason": "answer_restored", + }, + "_last_usage": usage.to_dict(), + } + } + + projection = WebUISessionProjection(sessions) + + assert projection.attach_fields("websocket:chat-1") == { + "model_preset": "Deep Research", + "recovery_state": { + "status": "recovered", + "recovery_id": "recovery-1", + "reason": "answer_restored", + }, + "usage": usage.to_turn_dict(), + } + sessions.read_session_metadata.assert_called_once_with("websocket:chat-1") + + +def test_attach_fields_tolerate_missing_or_invalid_session_metadata() -> None: + sessions = MagicMock() + sessions.read_session_metadata.return_value = { + "metadata": {SESSION_MODEL_PRESET_METADATA_KEY: 42} + } + log = MagicMock() + projection = WebUISessionProjection(sessions, log=log) + + assert projection.attach_fields("websocket:invalid") == {"model_preset": None} + log.warning.assert_called_once() + assert WebUISessionProjection(None).attach_fields("websocket:missing") == {} + + +def test_hydration_events_restore_goal_and_running_turn( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sessions = MagicMock() + sessions.read_session_metadata.return_value = { + "metadata": { + "goal_state": { + "status": "active", + "objective": "finish boundary split", + "ui_summary": "Refactoring", + } + } + } + monkeypatch.setattr( + "nanobot.webui.session_projection.websocket_turn_wall_started_at", + lambda _chat_id: 42.5, + ) + monkeypatch.setattr( + "nanobot.webui.session_projection.websocket_turn_id", + lambda _chat_id: "turn-1", + ) + + events = WebUISessionProjection(sessions).hydration_events( + "websocket:chat-1", + "chat-1", + ) + + assert events == ( + { + "event": "goal_state", + "chat_id": "chat-1", + "goal_state": { + "active": True, + "status": "active", + "ui_summary": "Refactoring", + "objective": "finish boundary split", + }, + }, + { + "event": "goal_status", + "chat_id": "chat-1", + "status": "running", + "started_at": 42.5, + "turn_id": "turn-1", + }, + ) + + +def test_hydration_events_are_quiet_without_actionable_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sessions = MagicMock() + sessions.read_session_metadata.return_value = {"metadata": {}} + monkeypatch.setattr( + "nanobot.webui.session_projection.websocket_turn_wall_started_at", + lambda _chat_id: None, + ) + + assert WebUISessionProjection(sessions).hydration_events( + "websocket:chat-1", + "chat-1", + ) == ()