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

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