feat(websocket): multiplex multiple chat_ids over a single connection

.gitignore

@@ -6,6 +6,12 @@
 .web
 .orion
 
+# webui (monorepo frontend)
+webui/node_modules/
+webui/dist/
+webui/coverage/
+webui/.vite/
+
 # Python bytecode & caches
 *.pyc
 *.pyo

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.
 

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:
+    async def _safe_send_to(self, connection: Any, raw: str, *, label: str = "") -> None:
-        """Send a raw frame, cleaning up dead connections on ConnectionClosed."""
+        """Send a raw frame to one connection, cleaning up on ConnectionClosed."""
-        connection = self._connections.get(chat_id)
-        if connection is None:
-            return
         try:
             await connection.send(raw)
         except ConnectionClosed:
-            self._connections.pop(chat_id, None)
+            self._cleanup_connection(connection)
-            logger.warning("websocket{}connection gone for chat_id={}", label, chat_id)
+            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)
+        # Snapshot the subscriber set so ConnectionClosed cleanups mid-iteration are safe.
-        if connection is None:
+        conns = list(self._subs.get(msg.chat_id, ()))
-            logger.warning("websocket: no active connection for chat_id={}", 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 ")

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

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