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

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

Browse Source
This commit is contained in:
Xubin Ren
2026-04-18 16:49:12 +08:00
committed by Xubin Ren
parent 70a1279b86
commit 6bfb75ed03
5 changed files with 457 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+6
View File
@@ -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
+71 -6
View File
@@ -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
+144 -22
View File
@@ -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 ")
tests/channels/test_websocket_channel.py
+234 -7
View File
@@ -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
+2 -1
View File
@@ -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