mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 00:03:01 +03:00
refactor(webui): isolate websocket application orchestration (#5548)
* refactor(webui): extract session attach projection * refactor(webui): isolate websocket application orchestration * refactor(webui): tighten websocket application boundary * test(webui): assert module logger for fork failures
This commit is contained in:
+105
-1123
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
|
||||
|
||||
@@ -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",
|
||||
),
|
||||
):
|
||||
|
||||
@@ -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 "
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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)
|
||||
|
||||
+26
-12
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
@@ -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",
|
||||
) == ()
|
||||
Reference in New Issue
Block a user