diff --git a/.gitignore b/.gitignore index 2820a379..151e6947 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,12 @@ .web .orion +# webui (monorepo frontend) +webui/node_modules/ +webui/dist/ +webui/coverage/ +webui/.vite/ + # Python bytecode & caches *.pyc *.pyo diff --git a/docs/WEBSOCKET.md b/docs/WEBSOCKET.md index 5cc867b4..1e0ddbbb 100644 --- a/docs/WEBSOCKET.md +++ b/docs/WEBSOCKET.md @@ -7,7 +7,7 @@ Nanobot can act as a WebSocket server, allowing external clients (web apps, CLIs - Bidirectional real-time communication over WebSocket - Streaming support — receive agent responses token by token - Token-based authentication (static tokens and short-lived issued tokens) -- Per-connection sessions — each connection gets a unique `chat_id` +- Multi-chat multiplexing — one connection can run many concurrent `chat_id`s - TLS/SSL support (WSS) with enforced TLSv1.2 minimum - Client allow-list via `allowFrom` - Auto-cleanup of dead connections @@ -98,6 +98,7 @@ All frames are JSON text. Each message has an `event` field. ```json { "event": "message", + "chat_id": "uuid-v4", "text": "Hello! How can I help?", "media": ["/tmp/image.png"], "reply_to": "msg-id" @@ -111,6 +112,7 @@ All frames are JSON text. Each message has an `event` field. ```json { "event": "delta", + "chat_id": "uuid-v4", "text": "Hello", "stream_id": "s1" } @@ -121,25 +123,46 @@ All frames are JSON text. Each message has an `event` field. ```json { "event": "stream_end", + "chat_id": "uuid-v4", "stream_id": "s1" } ``` +**`attached`** — confirmation for `new_chat` / `attach` inbound envelopes (see [Multi-chat multiplexing](#multi-chat-multiplexing)): + +```json +{"event": "attached", "chat_id": "uuid-v4"} +``` + +**`error`** — soft error for malformed inbound envelopes. The connection stays open: + +```json +{"event": "error", "detail": "invalid chat_id"} +``` + ### Client → Server -Send plain text: +**Legacy (default chat):** send a plain string, or a JSON object with a recognized text field: ```json "Hello nanobot!" ``` -Or send a JSON object with a recognized text field: - ```json {"content": "Hello nanobot!"} ``` -Recognized fields: `content`, `text`, `message` (checked in that order). Invalid JSON is treated as plain text. +Recognized fields: `content`, `text`, `message` (checked in that order). Invalid JSON is treated as plain text. These frames route to the connection's default `chat_id` (the one announced in `ready`). + +**Typed envelopes (multi-chat):** any JSON object with a string `type` field is a typed envelope: + +| `type` | Fields | Effect | +|--------|--------|--------| +| `new_chat` | — | Server mints a new `chat_id`, subscribes this connection, replies with `attached`. | +| `attach` | `chat_id` | Subscribe to an existing `chat_id` (e.g. after a page reload). Replies with `attached`. | +| `message` | `chat_id`, `content` | Send `content` on `chat_id`. First use auto-attaches; no explicit `attach` needed. | + +See [Multi-chat multiplexing](#multi-chat-multiplexing) for the full flow. ## Configuration Reference @@ -243,11 +266,53 @@ websocat "ws://127.0.0.1:8765/ws?client_id=alice&token=nbwt_aBcDeFg..." - Outstanding tokens are capped at 10,000. Requests beyond this return HTTP 429. - Expired tokens are purged lazily on each issue or validation request. +## Multi-chat multiplexing + +A single WebSocket can carry many concurrent chats. The server tracks `chat_id -> {connections}` as a fan-out set, so the same chat can also be mirrored across multiple connections (e.g. two browser tabs). + +### Typical flow (web UI with a sidebar) + +```text +client server + | --- connect --------------------> | + | <-- {"event":"ready", | + | "chat_id":"d3..."} (default)| + | | + | --- {"type":"new_chat"} ---------> | + | <-- {"event":"attached", | + | "chat_id":"a1..."} | + | | + | --- {"type":"message", | + | "chat_id":"a1...", | + | "content":"hi"} ------------> | + | <-- {"event":"delta", ...} | + | <-- {"event":"stream_end", ...} | + | | + | --- {"type":"attach", | # after page reload + | "chat_id":"a1..."} ---------> | + | <-- {"event":"attached", ...} | +``` + +### Rules + +- Every outbound event carries `chat_id`. Clients must dispatch by that field. +- `chat_id` format: `^[A-Za-z0-9_:-]{1,64}$`. Non-matching values return `error`. +- `message` auto-attaches on first use — no separate `attach` is required for chats the server minted (`new_chat`) on the same connection. +- Errors (invalid envelope, unknown `type`, bad `chat_id`) are soft: the server replies with `{"event":"error","detail":"..."}` and keeps the connection open. + +### Backward compatibility + +Legacy clients that only send plain text or `{"content": ...}` keep working unchanged: those frames route to the connection's default `chat_id` (the one from `ready`). No config flag is needed. + +### Security boundary + +`chat_id` is a *capability*: anyone holding a valid WebSocket auth credential and the chat_id can attach to that conversation and see its output. This is safe for nanobot's local, single-user model. Multi-tenant deployments should namespace chat_ids per user (or introduce a per-tenant auth gate) — nanobot does not do this today. + ## Security Notes - **Timing-safe comparison**: Static token validation uses `hmac.compare_digest` to prevent timing attacks. - **Defense in depth**: `allowFrom` is checked at both the HTTP handshake level and the message level. -- **Token isolation**: Each WebSocket connection gets a unique `chat_id`. Clients cannot access other sessions. +- **chat_id as capability**: see [Multi-chat multiplexing](#multi-chat-multiplexing). Auth on the WebSocket handshake is the single line of defense; callers who pass it can attach to any chat_id they know. - **TLS enforcement**: When SSL is enabled, TLSv1.2 is the minimum allowed version. - **Default-secure**: `websocketRequiresToken` defaults to `true`. Explicitly set it to `false` only on trusted networks. diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index 2af61d6f..882793ae 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -7,6 +7,7 @@ import email.utils import hmac import http import json +import re import secrets import ssl import time @@ -19,7 +20,8 @@ from pydantic import Field, field_validator, model_validator from websockets.asyncio.server import ServerConnection, serve from websockets.datastructures import Headers from websockets.exceptions import ConnectionClosed -from websockets.http11 import Request as WsRequest, Response +from websockets.http11 import Request as WsRequest +from websockets.http11 import Response from nanobot.bus.events import OutboundMessage from nanobot.bus.queue import MessageBus @@ -156,6 +158,37 @@ def _parse_inbound_payload(raw: str) -> str | None: return text +# Accept UUIDs and short scoped keys like "unified:default". Keeps the capability +# namespace small enough to rule out path traversal / quote injection tricks. +_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") + + +def _is_valid_chat_id(value: Any) -> bool: + return isinstance(value, str) and _CHAT_ID_RE.match(value) is not None + + +def _parse_envelope(raw: str) -> dict[str, Any] | None: + """Return a typed envelope dict if the frame is a new-style JSON envelope, else None. + + A frame qualifies when it parses as a JSON object with a string ``type`` field. + Legacy frames (plain text, or ``{"content": ...}`` without ``type``) return None; + callers should fall back to :func:`_parse_inbound_payload` for those. + """ + text = raw.strip() + if not text.startswith("{"): + return None + try: + data = json.loads(text) + except json.JSONDecodeError: + return None + if not isinstance(data, dict): + return None + t = data.get("type") + if not isinstance(t, str): + return None + return data + + def _issue_route_secret_matches(headers: Any, configured_secret: str) -> bool: """Return True if the token-issue HTTP request carries credentials matching ``token_issue_secret``.""" if not configured_secret: @@ -181,11 +214,47 @@ class WebSocketChannel(BaseChannel): config = WebSocketConfig.model_validate(config) super().__init__(config, bus) self.config: WebSocketConfig = config - self._connections: dict[str, Any] = {} + # chat_id -> connections subscribed to it (fan-out target). + self._subs: dict[str, set[Any]] = {} + # connection -> chat_ids it is subscribed to (O(1) cleanup on disconnect). + self._conn_chats: dict[Any, set[str]] = {} + # connection -> default chat_id for legacy frames that omit routing. + self._conn_default: dict[Any, str] = {} self._issued_tokens: dict[str, float] = {} self._stop_event: asyncio.Event | None = None self._server_task: asyncio.Task[None] | None = None + # -- Subscription bookkeeping ------------------------------------------- + + def _attach(self, connection: Any, 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: + """Remove *connection* from every subscription set; safe to call multiple times.""" + chat_ids = self._conn_chats.pop(connection, set()) + for cid in chat_ids: + subs = self._subs.get(cid) + if subs is None: + continue + subs.discard(connection) + if not subs: + self._subs.pop(cid, None) + self._conn_default.pop(connection, None) + + async def _send_event(self, connection: Any, event: str, **fields: Any) -> None: + """Send a control event (attached, error, ...) to a single connection.""" + payload: dict[str, Any] = {"event": event} + payload.update(fields) + raw = json.dumps(payload, ensure_ascii=False) + try: + await connection.send(raw) + except ConnectionClosed: + self._cleanup_connection(connection) + except Exception as e: + logger.warning("websocket: failed to send {} event: {}", event, e) + @classmethod def default_config(cls) -> dict[str, Any]: return WebSocketConfig().model_dump(by_alias=True) @@ -353,21 +422,22 @@ class WebSocketChannel(BaseChannel): logger.warning("websocket: client_id too long ({} chars), truncating", len(client_id)) client_id = client_id[:128] - chat_id = str(uuid.uuid4()) + default_chat_id = str(uuid.uuid4()) try: await connection.send( json.dumps( { "event": "ready", - "chat_id": chat_id, + "chat_id": default_chat_id, "client_id": client_id, }, ensure_ascii=False, ) ) # Register only after ready is successfully sent to avoid out-of-order sends - self._connections[chat_id] = connection + self._conn_default[connection] = default_chat_id + self._attach(connection, default_chat_id) async for raw in connection: if isinstance(raw, bytes): @@ -376,19 +446,66 @@ class WebSocketChannel(BaseChannel): except UnicodeDecodeError: logger.warning("websocket: ignoring non-utf8 binary frame") continue + + envelope = _parse_envelope(raw) + if envelope is not None: + await self._dispatch_envelope(connection, client_id, envelope) + continue + content = _parse_inbound_payload(raw) if content is None: continue await self._handle_message( sender_id=client_id, - chat_id=chat_id, + chat_id=default_chat_id, content=content, metadata={"remote": getattr(connection, "remote_address", None)}, ) except Exception as e: logger.debug("websocket connection ended: {}", e) finally: - self._connections.pop(chat_id, None) + self._cleanup_connection(connection) + + async def _dispatch_envelope( + self, + connection: Any, + client_id: str, + envelope: dict[str, Any], + ) -> None: + """Route one typed inbound envelope (``new_chat`` / ``attach`` / ``message``).""" + t = envelope.get("type") + if t == "new_chat": + new_id = str(uuid.uuid4()) + self._attach(connection, new_id) + await self._send_event(connection, "attached", chat_id=new_id) + return + if t == "attach": + cid = envelope.get("chat_id") + if not _is_valid_chat_id(cid): + await self._send_event(connection, "error", detail="invalid chat_id") + return + self._attach(connection, cid) + await self._send_event(connection, "attached", chat_id=cid) + return + if t == "message": + cid = envelope.get("chat_id") + content = envelope.get("content") + if not _is_valid_chat_id(cid): + await self._send_event(connection, "error", detail="invalid chat_id") + return + if not isinstance(content, str) or not content.strip(): + await self._send_event(connection, "error", detail="missing content") + return + # Auto-attach on first use so clients can one-shot without a separate attach. + self._attach(connection, cid) + await self._handle_message( + sender_id=client_id, + chat_id=cid, + content=content, + metadata={"remote": getattr(connection, "remote_address", None)}, + ) + return + await self._send_event(connection, "error", detail=f"unknown type: {t!r}") async def stop(self) -> None: if not self._running: @@ -402,30 +519,31 @@ class WebSocketChannel(BaseChannel): except Exception as e: logger.warning("websocket: server task error during shutdown: {}", e) self._server_task = None - self._connections.clear() + self._subs.clear() + self._conn_chats.clear() + self._conn_default.clear() self._issued_tokens.clear() - async def _safe_send(self, chat_id: str, raw: str, *, label: str = "") -> None: - """Send a raw frame, cleaning up dead connections on ConnectionClosed.""" - connection = self._connections.get(chat_id) - if connection is None: - return + async def _safe_send_to(self, connection: Any, raw: str, *, label: str = "") -> None: + """Send a raw frame to one connection, cleaning up on ConnectionClosed.""" try: await connection.send(raw) except ConnectionClosed: - self._connections.pop(chat_id, None) - logger.warning("websocket{}connection gone for chat_id={}", label, chat_id) + self._cleanup_connection(connection) + logger.warning("websocket{}connection gone", label) except Exception as e: logger.error("websocket{}send failed: {}", label, e) raise async def send(self, msg: OutboundMessage) -> None: - connection = self._connections.get(msg.chat_id) - if connection is None: - logger.warning("websocket: no active connection for chat_id={}", msg.chat_id) + # Snapshot the subscriber set so ConnectionClosed cleanups mid-iteration are safe. + conns = list(self._subs.get(msg.chat_id, ())) + if not conns: + logger.warning("websocket: no active subscribers for chat_id={}", msg.chat_id) return payload: dict[str, Any] = { "event": "message", + "chat_id": msg.chat_id, "text": msg.content, } if msg.media: @@ -433,7 +551,8 @@ class WebSocketChannel(BaseChannel): if msg.reply_to: payload["reply_to"] = msg.reply_to raw = json.dumps(payload, ensure_ascii=False) - await self._safe_send(msg.chat_id, raw, label=" ") + for connection in conns: + await self._safe_send_to(connection, raw, label=" ") async def send_delta( self, @@ -441,17 +560,20 @@ class WebSocketChannel(BaseChannel): delta: str, metadata: dict[str, Any] | None = None, ) -> None: - if self._connections.get(chat_id) is None: + conns = list(self._subs.get(chat_id, ())) + if not conns: return meta = metadata or {} if meta.get("_stream_end"): - body: dict[str, Any] = {"event": "stream_end"} + body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id} else: body = { "event": "delta", + "chat_id": chat_id, "text": delta, } if meta.get("_stream_id") is not None: body["stream_id"] = meta["_stream_id"] raw = json.dumps(body, ensure_ascii=False) - await self._safe_send(chat_id, raw, label=" stream ") + for connection in conns: + await self._safe_send_to(connection, raw, label=" stream ") diff --git a/tests/channels/test_websocket_channel.py b/tests/channels/test_websocket_channel.py index 89a330a1..c7d4923f 100644 --- a/tests/channels/test_websocket_channel.py +++ b/tests/channels/test_websocket_channel.py @@ -17,9 +17,11 @@ from nanobot.bus.events import OutboundMessage from nanobot.channels.websocket import ( WebSocketChannel, WebSocketConfig, + _is_valid_chat_id, _issue_route_secret_matches, _normalize_config_path, _normalize_http_path, + _parse_envelope, _parse_inbound_payload, _parse_query, _parse_request_path, @@ -168,7 +170,7 @@ async def test_send_delivers_json_message_with_media_and_reply() -> None: bus = MagicMock() channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus) mock_ws = AsyncMock() - channel._connections["chat-1"] = mock_ws + channel._attach(mock_ws, "chat-1") msg = OutboundMessage( channel="websocket", @@ -182,6 +184,7 @@ async def test_send_delivers_json_message_with_media_and_reply() -> None: mock_ws.send.assert_awaited_once() payload = json.loads(mock_ws.send.call_args[0][0]) assert payload["event"] == "message" + assert payload["chat_id"] == "chat-1" assert payload["text"] == "hello" assert payload["reply_to"] == "m1" assert payload["media"] == ["/tmp/a.png"] @@ -201,12 +204,13 @@ async def test_send_removes_connection_on_connection_closed() -> None: channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus) mock_ws = AsyncMock() mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True) - channel._connections["chat-1"] = mock_ws + channel._attach(mock_ws, "chat-1") msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello") await channel.send(msg) - assert "chat-1" not in channel._connections + assert "chat-1" not in channel._subs + assert mock_ws not in channel._conn_chats @pytest.mark.asyncio @@ -215,11 +219,12 @@ async def test_send_delta_removes_connection_on_connection_closed() -> None: channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus) mock_ws = AsyncMock() mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True) - channel._connections["chat-1"] = mock_ws + channel._attach(mock_ws, "chat-1") await channel.send_delta("chat-1", "chunk", {"_stream_delta": True, "_stream_id": "s1"}) - assert "chat-1" not in channel._connections + assert "chat-1" not in channel._subs + assert mock_ws not in channel._conn_chats @pytest.mark.asyncio @@ -227,7 +232,7 @@ async def test_send_delta_emits_delta_and_stream_end() -> None: bus = MagicMock() channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus) mock_ws = AsyncMock() - channel._connections["chat-1"] = mock_ws + channel._attach(mock_ws, "chat-1") await channel.send_delta("chat-1", "part", {"_stream_delta": True, "_stream_id": "sid"}) await channel.send_delta("chat-1", "", {"_stream_end": True, "_stream_id": "sid"}) @@ -236,9 +241,11 @@ async def test_send_delta_emits_delta_and_stream_end() -> None: first = json.loads(mock_ws.send.call_args_list[0][0][0]) second = json.loads(mock_ws.send.call_args_list[1][0][0]) assert first["event"] == "delta" + assert first["chat_id"] == "chat-1" assert first["text"] == "part" assert first["stream_id"] == "sid" assert second["event"] == "stream_end" + assert second["chat_id"] == "chat-1" assert second["stream_id"] == "sid" @@ -248,7 +255,7 @@ async def test_send_non_connection_closed_exception_is_raised() -> None: channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus) mock_ws = AsyncMock() mock_ws.send.side_effect = RuntimeError("unexpected") - channel._connections["chat-1"] = mock_ws + channel._attach(mock_ws, "chat-1") msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello") with pytest.raises(RuntimeError, match="unexpected"): @@ -596,3 +603,223 @@ async def test_websocket_requires_token_without_issue_path(bus: MagicMock) -> No finally: await channel.stop() await server_task + + +# -- Multi-chat multiplexing ------------------------------------------------- +# +# The multiplex protocol lets one WS connection route N logical chats over +# typed envelopes (`new_chat` / `attach` / `message`). Legacy frames must keep +# working on the connection's default chat_id. + + +@pytest.mark.asyncio +async def test_multiplex_legacy_still_works(bus: MagicMock) -> None: + port = 29930 + channel = _ch(bus, port=port) + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + + try: + async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=legacy") as client: + ready = json.loads(await client.recv()) + default_chat = ready["chat_id"] + + # Plain text frame routes to default chat_id + await client.send("hello from legacy") + await asyncio.sleep(0.1) + inbound = bus.publish_inbound.call_args[0][0] + assert inbound.chat_id == default_chat + assert inbound.content == "hello from legacy" + + # {"content": ...} frame routes to default chat_id + await client.send(json.dumps({"content": "structured legacy"})) + await asyncio.sleep(0.1) + assert bus.publish_inbound.call_args[0][0].chat_id == default_chat + assert bus.publish_inbound.call_args[0][0].content == "structured legacy" + + # Outbound still reaches the legacy client, with chat_id annotated + await channel.send( + OutboundMessage(channel="websocket", chat_id=default_chat, content="reply") + ) + reply = json.loads(await client.recv()) + assert reply["event"] == "message" + assert reply["chat_id"] == default_chat + assert reply["text"] == "reply" + finally: + await channel.stop() + await server_task + + +@pytest.mark.asyncio +async def test_multiplex_new_chat_roundtrip(bus: MagicMock) -> None: + port = 29931 + channel = _ch(bus, port=port) + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + + try: + async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=mp") as client: + ready = json.loads(await client.recv()) + default_chat = ready["chat_id"] + + await client.send(json.dumps({"type": "new_chat"})) + attached = json.loads(await client.recv()) + assert attached["event"] == "attached" + new_chat = attached["chat_id"] + assert new_chat and new_chat != default_chat + + # Send on the new chat via typed envelope + await client.send( + json.dumps({"type": "message", "chat_id": new_chat, "content": "hi on new"}) + ) + await asyncio.sleep(0.1) + inbound = bus.publish_inbound.call_args[0][0] + assert inbound.chat_id == new_chat + assert inbound.content == "hi on new" + + # Server pushes a message back; chat_id must match + await channel.send( + OutboundMessage(channel="websocket", chat_id=new_chat, content="ok") + ) + reply = json.loads(await client.recv()) + assert reply["event"] == "message" + assert reply["chat_id"] == new_chat + assert reply["text"] == "ok" + finally: + await channel.stop() + await server_task + + +@pytest.mark.asyncio +async def test_multiplex_two_chats_isolated(bus: MagicMock) -> None: + port = 29932 + channel = _ch(bus, port=port) + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + + try: + async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=two") as client: + await client.recv() # ready + + await client.send(json.dumps({"type": "new_chat"})) + chat_a = json.loads(await client.recv())["chat_id"] + await client.send(json.dumps({"type": "new_chat"})) + chat_b = json.loads(await client.recv())["chat_id"] + assert chat_a != chat_b + + # Push A → client sees A only (FIFO over the single WS). + await channel.send( + OutboundMessage(channel="websocket", chat_id=chat_a, content="for-A") + ) + msg_a = json.loads(await client.recv()) + assert msg_a["chat_id"] == chat_a + assert msg_a["text"] == "for-A" + + # Push B → client sees B only. + await channel.send( + OutboundMessage(channel="websocket", chat_id=chat_b, content="for-B") + ) + msg_b = json.loads(await client.recv()) + assert msg_b["chat_id"] == chat_b + assert msg_b["text"] == "for-B" + finally: + await channel.stop() + await server_task + + +@pytest.mark.asyncio +async def test_multiplex_invalid_frames_return_error(bus: MagicMock) -> None: + port = 29933 + channel = _ch(bus, port=port) + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + + try: + async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=bad") as client: + await client.recv() # ready + + # attach with bad chat_id + await client.send(json.dumps({"type": "attach", "chat_id": "has space"})) + err1 = json.loads(await client.recv()) + assert err1["event"] == "error" + + # message with missing content + await client.send(json.dumps({"type": "message", "chat_id": "abc", "content": ""})) + err2 = json.loads(await client.recv()) + assert err2["event"] == "error" + + # unknown type + await client.send(json.dumps({"type": "nope"})) + err3 = json.loads(await client.recv()) + assert err3["event"] == "error" + + # Connection survives: legacy frame still works. + await client.send("still-alive") + await asyncio.sleep(0.1) + bus.publish_inbound.assert_awaited() + assert bus.publish_inbound.call_args[0][0].content == "still-alive" + finally: + await channel.stop() + await server_task + + +@pytest.mark.asyncio +async def test_multiplex_cleanup_on_disconnect(bus: MagicMock) -> None: + port = 29934 + channel = _ch(bus, port=port) + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + + try: + async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=dc") as client: + ready = json.loads(await client.recv()) + default_chat = ready["chat_id"] + await client.send(json.dumps({"type": "new_chat"})) + extra_chat = json.loads(await client.recv())["chat_id"] + assert default_chat in channel._subs + assert extra_chat in channel._subs + # Client gone. Server-side tracking must be empty. + await asyncio.sleep(0.2) + assert default_chat not in channel._subs + assert extra_chat not in channel._subs + assert not channel._conn_chats + assert not channel._conn_default + finally: + await channel.stop() + await server_task + + +def test_parse_envelope_detects_typed_frames() -> None: + assert _parse_envelope('{"type":"new_chat"}') == {"type": "new_chat"} + env = _parse_envelope('{"type":"message","chat_id":"abc","content":"hi"}') + assert env == {"type": "message", "chat_id": "abc", "content": "hi"} + + +def test_parse_envelope_rejects_legacy_and_garbage() -> None: + # No `type` field → legacy, caller falls back to _parse_inbound_payload. + assert _parse_envelope('{"content":"hi"}') is None + assert _parse_envelope("plain text") is None + assert _parse_envelope("{broken") is None + assert _parse_envelope("[1,2,3]") is None + # Non-string `type` is not a valid envelope. + assert _parse_envelope('{"type":123}') is None + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("abc", True), + ("a1b2_c:d-e", True), + ("x" * 64, True), + ("unified:default", True), + ("", False), + ("x" * 65, False), + ("has space", False), + ("a/b", False), + ("a.b", False), + (None, False), + (123, False), + ], +) +def test_is_valid_chat_id(value: Any, expected: bool) -> None: + assert _is_valid_chat_id(value) is expected diff --git a/tests/channels/test_websocket_integration.py b/tests/channels/test_websocket_integration.py index 8ff15866..ff0d5085 100644 --- a/tests/channels/test_websocket_integration.py +++ b/tests/channels/test_websocket_integration.py @@ -290,10 +290,11 @@ async def test_disconnected_client_cleanup(bus: MagicMock) -> None: async with WsTestClient("ws://127.0.0.1:29914/", client_id="tmp") as c: chat_id = (await c.recv_ready()).chat_id # disconnected + await asyncio.sleep(0.1) await ch.send(OutboundMessage( channel="websocket", chat_id=chat_id, content="orphan", )) - assert chat_id not in ch._connections + assert chat_id not in ch._subs finally: await ch.stop(); await t