Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: enforce BasedPyright strict type checking (#5158)

Browse Source
This commit is contained in:
chengyongru
2026-07-29 21:37:11 +08:00
committed by GitHub
parent e703481755
commit 757ad9c764
166 changed files with 4728 additions and 2621 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/channels/websocket/runtime.py
+73 -26
View File
@@ -11,7 +11,7 @@ import uuid
from collections.abc import Callable
from contextlib import suppress
from pathlib import Path
from typing import Any, Self
from typing import Any, Self, TypeGuard, cast
from pydantic import Field, field_validator, model_validator
from websockets.asyncio.server import ServerConnection, serve, unix_serve
@@ -191,12 +191,13 @@ def _parse_inbound_payload(raw: str) -> str | None:
return None
if text.startswith("{"):
try:
data = json.loads(text)
data = cast(object, json.loads(text))
except json.JSONDecodeError:
return text
if isinstance(data, dict):
payload = cast(dict[str, Any], data)
for key in ("content", "text", "message"):
value = data.get(key)
value = payload.get(key)
if isinstance(value, str) and value.strip():
return value
return None
@@ -209,7 +210,7 @@ def _parse_inbound_payload(raw: str) -> str | None:
_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$")
def _is_valid_chat_id(value: Any) -> bool:
def _is_valid_chat_id(value: Any) -> TypeGuard[str]:
return isinstance(value, str) and _CHAT_ID_RE.match(value) is not None
@@ -224,15 +225,16 @@ def _parse_envelope(raw: str) -> dict[str, Any] | None:
if not text.startswith("{"):
return None
try:
data = json.loads(text)
data = cast(object, json.loads(text))
except json.JSONDecodeError:
return None
if not isinstance(data, dict):
return None
t = data.get("type")
envelope = cast(dict[str, Any], data)
t = envelope.get("type")
if not isinstance(t, str):
return None
return data
return envelope
def _is_websocket_upgrade(request: WsRequest) -> bool:
@@ -264,13 +266,13 @@ class WebSocketChannel(BaseChannel):
super().__init__(config, bus)
self.config: WebSocketConfig = config
# chat_id -> connections subscribed to it (fan-out target).
self._subs: dict[str, set[Any]] = {}
self._subs: dict[str, set[ServerConnection]] = {}
# connection -> chat_ids it is subscribed to (O(1) cleanup on disconnect).
self._conn_chats: dict[Any, set[str]] = {}
self._conn_chats: dict[ServerConnection, set[str]] = {}
# connection -> default chat_id for legacy frames that omit routing.
self._conn_default: dict[Any, str] = {}
self._conn_default: dict[ServerConnection, str] = {}
# Connections authenticated with a one-time token from /webui/bootstrap.
self._webui_connections: set[Any] = set()
self._webui_connections: set[ServerConnection] = set()
self._stop_event: asyncio.Event | None = None
self._server_task: asyncio.Task[None] | None = None
@@ -286,15 +288,43 @@ class WebSocketChannel(BaseChannel):
# -- Subscription bookkeeping -------------------------------------------
def _workspace_controls_available(self, connection: Any) -> bool:
def _workspace_controls_available(self, connection: ServerConnection) -> bool:
return self._http_router.workspace_controls_available(connection)
def _attach(self, connection: Any, chat_id: str) -> None:
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 _cleanup_connection(self, connection: Any) -> None:
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)
await self._send_event(
connection,
"session_updated",
chat_id=fork_id,
scope="metadata",
workspace_scope=scope.payload(),
)
await self._hydrate_after_subscribe(fork_id)
def _cleanup_connection(self, connection: ServerConnection) -> None:
"""Remove *connection* from every subscription set; safe to call multiple times."""
chat_ids = self._conn_chats.pop(connection, set())
for cid in chat_ids:
@@ -317,10 +347,11 @@ class WebSocketChannel(BaseChannel):
if self.gateway.session_manager is None:
return
row = self.gateway.session_manager.read_session_file(f"websocket:{chat_id}")
meta = row.get("metadata", {}) if isinstance(row, dict) else {}
row_data = row if isinstance(row, dict) else {}
meta = row_data.get("metadata", {})
if not isinstance(meta, dict):
meta = {}
blob = goal_state_ws_blob(meta)
blob = goal_state_ws_blob(cast(dict[str, Any], meta))
if not blob.get("active"):
return
await self.send_goal_state(chat_id, blob)
@@ -342,7 +373,12 @@ class WebSocketChannel(BaseChannel):
await self._maybe_push_active_goal_state(chat_id)
await self._maybe_push_turn_run_wall_clock(chat_id)
async def _send_event(self, connection: Any, event: str, **fields: Any) -> None:
async def _send_event(
self,
connection: ServerConnection,
event: str,
**fields: Any,
) -> None:
"""Send a control event (attached, error, ...) to a single connection."""
payload: dict[str, Any] = {"event": event}
payload.update(fields)
@@ -377,7 +413,7 @@ class WebSocketChannel(BaseChannel):
# -- HTTP dispatch ------------------------------------------------------
async def _dispatch_http(self, connection: Any, request: WsRequest) -> Any:
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)
@@ -394,7 +430,11 @@ class WebSocketChannel(BaseChannel):
# Everything else goes to the HTTP handler
return await self._http_router.dispatch(connection, request)
def _authorize_websocket_handshake(self, connection: Any, query: dict[str, list[str]]) -> Any:
def _authorize_websocket_handshake(
self,
connection: ServerConnection,
query: dict[str, list[str]],
) -> Any:
supplied = _query_first(query, "token")
static_token = self.config.token.strip()
@@ -414,7 +454,7 @@ class WebSocketChannel(BaseChannel):
self._consume_issued_token(connection, supplied)
return None
def _consume_issued_token(self, connection: Any, token: str) -> bool:
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)
@@ -509,7 +549,7 @@ class WebSocketChannel(BaseChannel):
self._server_task = asyncio.create_task(runner())
await self._server_task
async def _connection_loop(self, connection: Any) -> None:
async def _connection_loop(self, connection: ServerConnection) -> None:
request = connection.request
path_part = request.path if request else "/"
_, query = _parse_request_path(path_part)
@@ -574,7 +614,7 @@ class WebSocketChannel(BaseChannel):
async def _dispatch_envelope(
self,
connection: Any,
connection: ServerConnection,
client_id: str,
envelope: dict[str, Any],
) -> None:
@@ -700,7 +740,7 @@ class WebSocketChannel(BaseChannel):
**rejection_fields,
)
return
media_paths, reason = self._media.store_inbound_attachments(raw_media)
media_paths, reason = self._media.store_inbound_attachments(cast(list[Any], raw_media))
if reason is not None:
await self._send_event(
connection,
@@ -810,7 +850,7 @@ class WebSocketChannel(BaseChannel):
async def _workspace_scope_or_error(
self,
connection: Any,
connection: ServerConnection,
resolver: Callable[[], Any],
*,
chat_id: str | None = None,
@@ -841,7 +881,8 @@ class WebSocketChannel(BaseChannel):
try:
await self._server_task
except asyncio.CancelledError:
if asyncio.current_task() and asyncio.current_task().cancelling():
current_task = asyncio.current_task()
if current_task is not None and current_task.cancelling():
raise
self.logger.debug("server task was already cancelled during shutdown")
except Exception as e:
@@ -853,7 +894,13 @@ class WebSocketChannel(BaseChannel):
self._webui_connections.clear()
self._tokens.clear()
async def _safe_send_to(self, connection: Any, raw: str, *, label: str = "") -> None:
async def _safe_send_to(
self,
connection: ServerConnection,
raw: str,
*,
label: str = "",
) -> None:
"""Send a raw frame to one connection, cleaning up on ConnectionClosed."""
try:
await connection.send(raw)