diff --git a/docker-example/production-compose/Caddyfile b/docker-example/production-compose/Caddyfile index 32ac014e6da..9298ff0f31e 100644 --- a/docker-example/production-compose/Caddyfile +++ b/docker-example/production-compose/Caddyfile @@ -2,7 +2,7 @@ encode gzip -@backend_routes path /_event/* /ping /_health /_upload /_upload/* +@backend_routes path /_event /_event/* /ping /_health /_upload /_upload/* handle @backend_routes { reverse_proxy app:8000 } diff --git a/docker-example/production/Caddyfile b/docker-example/production/Caddyfile index e0bfb2705ef..ab80a6ed823 100644 --- a/docker-example/production/Caddyfile +++ b/docker-example/production/Caddyfile @@ -2,7 +2,7 @@ encode gzip -@backend_routes path /_event/* /ping /_health /_upload /_upload/* +@backend_routes path /_event /_event/* /ping /_health /_upload /_upload/* handle @backend_routes { reverse_proxy localhost:8000 } diff --git a/docs/api-reference/channels.md b/docs/api-reference/channels.md new file mode 100644 index 00000000000..b40b6b52cd0 --- /dev/null +++ b/docs/api-reference/channels.md @@ -0,0 +1,193 @@ +```python exec +import reflex as rx +``` + +# Channels + +A channel is an application-defined message stream multiplexed onto the +websocket your app already uses for state updates. It is the supported way for +a component or a third-party package to move data that does not belong in +state — streaming chart columns, a live cursor position, an audio buffer — +without opening a second connection. + +Because a channel rides the app's own socket, it inherits the connection's +origin checks, client token, reconnect handling and reverse-proxy setup. + +Channels require state to be enabled (the default) and the `websocket` +transport. They are not available under `transport="socketio"` or +`transport="polling"`: there, registering a channel logs a warning and skips +it, and its browser-side handle reports a `channels_unsupported` error. + +## Defining a channel + +Subclass `rx.channels.Channel`, give it a name, and handle messages: + +```python +import reflex as rx + + +class Ticks(rx.channels.Channel): + name = "ticks" + + async def on_open(self, session): + session.data["symbols"] = set() + + async def on_message(self, session, event, data, buffers): + if event == "subscribe": + session.join(data["symbol"]) + session.data["symbols"].add(data["symbol"]) + elif event == "unsubscribe": + session.leave(data["symbol"]) + + async def on_close(self, session): + del session.data["symbols"] + + +app = rx.App() +app.register_channel(Ticks()) +``` + +A package can register its channel from a plugin's `post_compile` hook instead, +which runs at backend startup with the live app. Registration never stops an +app that cannot serve channels from starting, so a package needs no check of +its own. + +Every field of an inbound message is client-controlled and unvalidated. A +handler that raises is logged and the connection keeps serving, so a failing +handler never drops the app's socket. + +Handlers run inline on their connection's receive loop, which keeps messages +in order and lets a slow channel push back on its client. It also means a +handler that awaits something slow stalls that connection — its state updates +wait, its heartbeat replies stop, and after `REFLEX_SOCKET_INTERVAL + +REFLEX_SOCKET_TIMEOUT` the server closes it as unresponsive. That holds for +anything the handler awaits, `asyncio.to_thread` included, so run long work in +a task of its own and answer when it finishes: + +```python +import asyncio + + +class Rows(rx.channels.Channel): + name = "rows" + + def __init__(self): + super().__init__() + # Held, so the event loop does not drop a running task. + self._tasks = set() + + async def on_message(self, session, event, data, buffers): + task = asyncio.create_task(self._answer(session, data["filter"])) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + + async def _answer(self, session, query_filter): + rows = await asyncio.to_thread(expensive_query, query_filter) + if session.open: + await session.send("rows", {"count": len(rows)}, [rows.tobytes()]) +``` + +## Sending to clients + +`ChannelSession.send` answers one client. `Channel.send_to_room` fans out to +everyone who joined a room, and `Channel.send_to_token` addresses every +session belonging to one client token (browser tab): + +```python +await session.send("tick", {"symbol": "RFX", "price": 42.0}) +await self.send_to_room("RFX", "tick", {"price": 42.0}) +``` + +Work running in a task of its own can outlive its client; `session.open` +reports whether the client is still connected before you commit to expensive +or long-lived work. Inside `on_message` itself it stays true: the connection +cannot close while its receive loop is waiting for the handler. + +Rooms and sessions are local to the worker holding the connection. A client +reconnecting to another worker opens its session there, so anything that must +outlive a connection belongs in Reflex state, not in the channel. + +To push from a background task or a thread, hold the serving event loop and +schedule the coroutine on it with `asyncio.run_coroutine_threadsafe`. + +## Binary payloads + +Messages may carry binary attachments beside their JSON metadata. They travel +as raw bytes — no base64, no JSON numbers — and arrive in the browser as +`Uint8Array` views, each aligned so it can be read as a typed array without +copying: + +```python +class Frames(rx.channels.Channel): + name = "frames" + # Inbound attachments are refused unless the channel opts in. + accepts_binary = True + + async def on_message(self, session, event, data, buffers): + await session.send("frame", {"rows": len(buffers[0]) // 8}, buffers) +``` + +A message may carry up to 64 attachments, in either direction: +`session.send` raises and `channel.emit` throws rather than build a frame the +other end would refuse. + +Size is capped in one direction only. A frame a client sends must fit +`REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE` (1 MB by default) or the backend closes +the connection, so `channel.emit` throws first and a message queued before the +channel opened is dropped with an `error` when the limit turns out to exclude +it. Raise the setting if your clients send larger payloads; Granian, the +default server, receives frames of at most 16 MiB, so Reflex caps the setting +there and you need uvicorn for anything larger. Messages the server sends are +not capped by it — mind the section below on sharing the connection. + +`connect`, `disconnect` and `error` are reserved message names: the client +handle reports its own lifecycle under them. So is every name starting with +`_`, which the channel protocol uses for its own messages. Neither end puts a +reserved name on the wire: `session.send` and `channel.emit` refuse them, and a +message that arrives under one is answered with a `reserved_event` error +instead of reaching `on_message`. + +## Using a channel from the frontend + +`getChannel` returns a handle that survives reconnects and remounts, so a +component can register its handlers once: + +```javascript +import { getChannel } from "$/utils/state"; + +const channel = getChannel("ticks"); + +channel.on("connect", () => channel.emit("subscribe", { symbol: "RFX" })); +channel.on("tick", (data, buffers) => console.log(data.price)); +channel.on("error", (error) => console.error(error.code, error.message)); + +// Attachments are an array of ArrayBuffers or typed arrays. +channel.emit("frame", { seq: 1 }, [new Float64Array([1, 2, 3])]); +``` + +`channel.emit` throws on anything the backend would refuse — a name that is +not a string or is reserved, attachments that are not an array of binary +buffers — rather than cost the app its socket. Read a `Blob` with +`await blob.arrayBuffer()` before attaching it. + +Messages emitted before the channel is open are queued and flushed on +`connect`. `error` reports a channel-level failure: an unknown channel name, a +backend too old to speak channels, or a transport that cannot carry them. + +## Performance + +On uvicorn, websocket messages are compressed with permessage-deflate by +default. That is a good trade for JSON state updates and a poor one for binary +data, which barely shrinks while costing milliseconds of event loop time per +message. Apps that stream binary over a channel should turn it off: + +```bash +REFLEX_SOCKET_PER_MESSAGE_DEFLATE=false reflex run --env prod +``` + +Granian, the default server, does not negotiate permessage-deflate at all, so +the setting only changes uvicorn deployments. + +Channel messages share one connection with state updates, so a very large +message delays the deltas queued behind it. Prefer messages bounded by what the +client actually needs — a screenful of data, not a whole dataset. diff --git a/docs/app/news/+channels-reference.docs.md b/docs/app/news/+channels-reference.docs.md new file mode 100644 index 00000000000..afb19494174 --- /dev/null +++ b/docs/app/news/+channels-reference.docs.md @@ -0,0 +1 @@ +Document channels: multiplexed message streams, with binary payloads, carried on the app's event websocket. diff --git a/docs/app/pyproject.toml b/docs/app/pyproject.toml index d89369e758d..6fe448176c1 100644 --- a/docs/app/pyproject.toml +++ b/docs/app/pyproject.toml @@ -15,7 +15,8 @@ dependencies = [ "psycopg[binary]; python_version < '3.15'", "psycopg; python_version >= '3.15'", "python-frontmatter", - "reflex", + # xy 0.0.7's charts ride a Socket.IO namespace; see rxconfig.py. + "reflex[socketio]", "reflex-docgen", "reflex-enterprise>=0.9.4", "reflex-hosting-cli", diff --git a/docs/app/reflex_docs/pages/docs/apiref.py b/docs/app/reflex_docs/pages/docs/apiref.py index e19310ccb72..7d836fe2945 100644 --- a/docs/app/reflex_docs/pages/docs/apiref.py +++ b/docs/app/reflex_docs/pages/docs/apiref.py @@ -49,6 +49,7 @@ "cli", "browser-storage", "browser-javascript", + "channels", "plugins", "utils", "telemetry", diff --git a/docs/app/rxconfig.py b/docs/app/rxconfig.py index 810f2f6c9f3..8418f2df1ad 100644 --- a/docs/app/rxconfig.py +++ b/docs/app/rxconfig.py @@ -18,6 +18,9 @@ "@fontsource-variable/jetbrains-mono@5.3.0", ], telemetry_enabled=False, + # xy 0.0.7 serves its charts' data plane as a namespace on the app's + # Socket.IO server; back to the default once xy rides a channel instead. + transport="socketio", plugins=[ reflex_xy.XYPlugin(), rx.plugins.TailwindV4Plugin(), diff --git a/docs/app/tests/test_sidebar.py b/docs/app/tests/test_sidebar.py index 5a87878ec4f..5a4e8b1f7a2 100644 --- a/docs/app/tests/test_sidebar.py +++ b/docs/app/tests/test_sidebar.py @@ -141,6 +141,7 @@ def test_api_reference_groups_related_symbols(): "CLI", "Browser Storage", "Browser Javascript", + "Channels", "Plugins", "Utils", "Telemetry", diff --git a/docs/changelog/upgrading/upgrading-to-0-10.md b/docs/changelog/upgrading/upgrading-to-0-10.md index d6f26a7f8d1..34819ff0c16 100644 --- a/docs/changelog/upgrading/upgrading-to-0-10.md +++ b/docs/changelog/upgrading/upgrading-to-0-10.md @@ -1,6 +1,6 @@ --- title: "Upgrading to Reflex 0.10" -meta_description: "Changes to review when upgrading an app from Reflex 0.9 to 0.10, with the fix for each: backend vars read on a state class, class-level defaults, background tasks, and state stores shared with 0.9 instances." +meta_description: "Changes to review when upgrading an app from Reflex 0.9 to 0.10, with the fix for each: backend vars read on a state class, class-level defaults, background tasks, the event transport, and state stores shared with 0.9 instances." --- # Upgrading to Reflex 0.10 @@ -131,6 +131,62 @@ class Child(Parent): self.bump() ``` +## The event transport + +`transport="websocket"`, the default, now means a plain WebSocket speaking +JSON frames rather than Socket.IO over a websocket +([#6932](https://github.com/reflex-dev/reflex/pull/6932)). Socket.IO remains +available as `transport="socketio"` (or `"polling"`) with the optional extra, +`pip install 'reflex[socketio]'`. + +| If your app | Do this | +| --- | --- | +| Passes `rx.App(sio=...)` | Set `transport="socketio"` in `rxconfig.py` and install `reflex[socketio]`. Without a transport setting the app still falls back to Socket.IO, with a deprecation warning. | +| Uses `app.sio` for its own Socket.IO events | The same: under the default transport `app.sio` is `None`. | +| Imports `socketio`, or `AsyncServer` or `EventNamespace` from `reflex.app` | Install `reflex[socketio]`: `python-socketio` is no longer a dependency of `reflex`. | +| Replaces `reflex.app.EventNamespace` with a subclass | That class serves only `transport="socketio"` and `"polling"`; replace `reflex.app.WebsocketEventNamespace` as well, and override `handle_event` rather than `on_event` (see [below](#running-code-for-every-event)). | +| Runs the backend on uvicorn | Install `reflex[uvicorn]`, which brings the `websockets` library uvicorn needs to serve the event socket. Without a websocket library, Reflex runs Granian instead. | +| Routes the backend through a reverse proxy | Send both `/_event` and `/_event/*` to the backend; the client dials `/_event/`. | +| Raised `REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE` | It counts UTF-8 bytes rather than characters, and Granian caps it at 16 MiB; run uvicorn for larger messages. | +| Uses `AppHarness.backend` or `AppHarnessProd.frontend_server` in tests | Both are an embedded Granian server rather than a `uvicorn.Server`: wait for startup with `wait_started()` (`started` is deprecated), and drop uses of `servers`, `config`, `serve()` and the other uvicorn internals. `reflex[testing]` no longer installs uvicorn; add it yourself if your tests import it. | + +Browser tabs left open across the upgrade, and frontends exported with 0.9 and +hosted separately, cannot reconnect: reload the tabs and re-export the +frontend. + +Under the default transport `app.event_namespace` is a +`WebsocketEventNamespace`. Its `emit(event, data, to=sid)` reaches one session, +or every session when `to` is omitted; Socket.IO's rooms, `skip_sid` and +callbacks are not available. An event whose payload is not a JSON object closes +the connection instead of being parsed from a string, and a Socket.IO +connection that carries no client token is refused. + +### Running code for every event + +To run code for every event, such as an authorization check or an audit log, +subclass the event namespace and override `handle_event`. On either transport +it receives every event, including the hydrate event a client sends with its +connect. Replace the class in `reflex.app` before creating the app: + +```python +import reflex as rx +import reflex.app +from reflex.event_namespace import WebsocketEventNamespace + + +class AuditingNamespace(WebsocketEventNamespace): + async def handle_event(self, sid, data, asgi_scope): + audit(sid, data) + await super().handle_event(sid, data, asgi_scope) + + +reflex.app.WebsocketEventNamespace = AuditingNamespace +app = rx.App() +``` + +Under `transport="socketio"` or `"polling"`, subclass `reflex.app.EventNamespace` +the same way. Overriding its `on_event` still works, but is deprecated. + ## Sharing a state store between 0.9 and 0.10 instances Instances of 0.9 and 0.10 cannot share a Redis or disk state store. A 0.10 diff --git a/docs/hosting/self-hosting.md b/docs/hosting/self-hosting.md index 439c2fe4723..7b44df219f9 100644 --- a/docs/hosting/self-hosting.md +++ b/docs/hosting/self-hosting.md @@ -71,8 +71,10 @@ REFLEX_API_URL=http://localhost:8000 reflex run --env prod --frontend-only --fro ``` Alternatively, put a reverse proxy in front of both processes and route the -backend routes (`/_event`, `/ping`, `/_upload`, `/_health`) to the backend -port, as the `Caddyfile`s in the container examples below do. +backend routes (`/_event` and `/_event/*`, `/ping`, `/_upload`, `/_health`) to +the backend port, as the `Caddyfile`s in the container examples below do. The +event websocket is dialed at `/_event/`; it must be proxied as a websocket +upgrade. Instances of Reflex 0.9 and 0.10 cannot share a state store, whether Redis or the disk state manager's directory. A 0.10 instance loads state saved by 0.9, diff --git a/news/+event-namespace-handle-event.deprecation.md b/news/+event-namespace-handle-event.deprecation.md new file mode 100644 index 00000000000..11662aae2a0 --- /dev/null +++ b/news/+event-namespace-handle-event.deprecation.md @@ -0,0 +1,16 @@ +Overriding `EventNamespace.on_event` is deprecated: override `handle_event`, which runs for every event, including the hydrate on connect, on both transports. Replace `reflex.app.WebsocketEventNamespace` the same way to cover the default transport; see [Running code for every event](https://reflex.dev/docs/changelog/upgrading/upgrading-to-0-10/#running-code-for-every-event). + +```python +# Before +class AuditingNamespace(EventNamespace): + async def on_event(self, sid, data): + audit(sid, data) + await super().on_event(sid, data) + + +# After +class AuditingNamespace(EventNamespace): + async def handle_event(self, sid, data, asgi_scope): + audit(sid, data) + await super().handle_event(sid, data, asgi_scope) +``` diff --git a/news/+socket-deflate-flag.performance.md b/news/+socket-deflate-flag.performance.md new file mode 100644 index 00000000000..534e67a4111 --- /dev/null +++ b/news/+socket-deflate-flag.performance.md @@ -0,0 +1 @@ +Add `REFLEX_SOCKET_PER_MESSAGE_DEFLATE` to control websocket permessage-deflate compression (on by default wherever Uvicorn serves, including the Gunicorn production worker; Granian never negotiates it). Turning it off is worthwhile for apps sending binary data over a channel: compressing it costs milliseconds of event loop time per message and barely shrinks it. diff --git a/news/+socketio-handler-failures.bugfix.md b/news/+socketio-handler-failures.bugfix.md new file mode 100644 index 00000000000..c55f0a2ba7d --- /dev/null +++ b/news/+socketio-handler-failures.bugfix.md @@ -0,0 +1 @@ +The Socket.IO transport now escapes unpaired surrogates in outgoing packets, which a non-UTF-8 file name in state could otherwise turn into a failed send, and logs event handler failures within the per-session error budget instead of raising them into python-socketio. diff --git a/news/+websocket-channels.feature.md b/news/+websocket-channels.feature.md new file mode 100644 index 00000000000..08afff09e2a --- /dev/null +++ b/news/+websocket-channels.feature.md @@ -0,0 +1 @@ +Add channels: an application-defined message stream, including binary payloads, multiplexed onto the websocket your app already uses for state updates. A component or package can stream data that does not belong in state without opening a second connection. See the [channels reference](https://reflex.dev/docs/api-reference/channels/). diff --git a/news/6932.breaking.md b/news/6932.breaking.md new file mode 100644 index 00000000000..f00ef48a542 --- /dev/null +++ b/news/6932.breaking.md @@ -0,0 +1 @@ +`transport="websocket"`, the default, now means a plain WebSocket rather than Socket.IO, and `python-socketio` and `uvicorn` moved to the `reflex[socketio]` and `reflex[uvicorn]` extras. Apps that pass `rx.App(sio=...)`, use `app.sio`, run uvicorn, proxy `/_event`, or use `AppHarness.backend` need changes: see [Upgrading to Reflex 0.10](https://reflex.dev/docs/changelog/upgrading/upgrading-to-0-10/#the-event-transport). diff --git a/news/6932.feature.md b/news/6932.feature.md new file mode 100644 index 00000000000..48b92a24827 --- /dev/null +++ b/news/6932.feature.md @@ -0,0 +1 @@ +The client and server now talk over a plain WebSocket with a lightweight JSON event protocol instead of Socket.IO, and the browser loads `socket.io-client` only when `transport="socketio"` or `"polling"` is configured. diff --git a/packages/reflex-base/news/+cors-origins-string.bugfix.md b/packages/reflex-base/news/+cors-origins-string.bugfix.md new file mode 100644 index 00000000000..f7fb4195326 --- /dev/null +++ b/packages/reflex-base/news/+cors-origins-string.bugfix.md @@ -0,0 +1 @@ +A plain-string `cors_allowed_origins` is read as a comma-separated list, as `REFLEX_CORS_ALLOWED_ORIGINS` is, instead of admitting every origin it contains as a substring. diff --git a/packages/reflex-base/news/+socket-deflate-flag.performance.md b/packages/reflex-base/news/+socket-deflate-flag.performance.md new file mode 100644 index 00000000000..670dcaa7922 --- /dev/null +++ b/packages/reflex-base/news/+socket-deflate-flag.performance.md @@ -0,0 +1 @@ +Add the `REFLEX_SOCKET_PER_MESSAGE_DEFLATE` environment variable, controlling whether uvicorn compresses websocket messages with permessage-deflate; Granian never negotiates it, so the variable has no effect there. diff --git a/packages/reflex-base/news/+websocket-channels.feature.md b/packages/reflex-base/news/+websocket-channels.feature.md new file mode 100644 index 00000000000..93c27b9888d --- /dev/null +++ b/packages/reflex-base/news/+websocket-channels.feature.md @@ -0,0 +1 @@ +The frontend websocket client can carry channels: named message streams, with binary attachments, multiplexed onto the app's event socket. Components reach one with `getChannel(name)` from `$/utils/state`. diff --git a/packages/reflex-base/news/6932.breaking.md b/packages/reflex-base/news/6932.breaking.md new file mode 100644 index 00000000000..3dcb16dfe38 --- /dev/null +++ b/packages/reflex-base/news/6932.breaking.md @@ -0,0 +1 @@ +`transport="websocket"`, the default, now means a plain WebSocket speaking JSON `[event_name, payload]` frames rather than Socket.IO over a websocket; set `transport="socketio"` to keep Socket.IO. See [Upgrading to Reflex 0.10](https://reflex.dev/docs/changelog/upgrading/upgrading-to-0-10/#the-event-transport). diff --git a/packages/reflex-base/news/6932.feature.md b/packages/reflex-base/news/6932.feature.md new file mode 100644 index 00000000000..64ce4283ac2 --- /dev/null +++ b/packages/reflex-base/news/6932.feature.md @@ -0,0 +1 @@ +`config.transport` gains a `"socketio"` value for Socket.IO over a websocket, and the browser loads socket.io-client only when `"socketio"` or `"polling"` is configured. diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/json.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/json.js index cd52a484ab6..5e5082cdf57 100644 --- a/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/json.js +++ b/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/json.js @@ -70,3 +70,26 @@ export const parseJson = (str) => { return parseNonFiniteFloats(str); } }; + +/** + * Parse JSON, tolerating bare non-finite float tokens. + * @param text The text to parse. + * @param fallback The value to return if the text is unparsable. + * @returns The parsed value, or the fallback. + */ +export const parseJsonLenient = (text, fallback) => { + try { + return parseJson(text); + } catch { + return fallback; + } +}; + +/** + * JSON.stringify replacer that sends undefined fields as null instead of + * removing them. Also assigned as the socket.io encoder replacer. + * @param _k The key being serialized. + * @param v The value being serialized. + * @returns The value to serialize. + */ +export const undefinedToNull = (_k, v) => (v === undefined ? null : v); diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/websocket.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/websocket.js new file mode 100644 index 00000000000..7eab3b01714 --- /dev/null +++ b/packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/websocket.js @@ -0,0 +1,961 @@ +// Plain WebSocket transport speaking the Reflex JSON event protocol: each +// frame is a JSON array `[event_name, payload]`. Mirrors the socket.io-client +// surface that state.js and upload.js rely on: connected, connect(), +// disconnect(), emit(), on(), auth, io.opts.query, and _callbacks. +import { parseJsonLenient, undefinedToNull } from "$/utils/helpers/json"; + +// Protocol-level message names (must match reflex/event_namespace.py). +const CONNECT_MESSAGE = "_connect"; +const HANDSHAKE_MESSAGE = "_handshake"; +const PING_MESSAGE = "_ping"; +const PONG_MESSAGE = "_pong"; +const OPEN_MESSAGE = "_open"; +const OPENED_MESSAGE = "_opened"; +const CHANNEL_ERROR_MESSAGE = "_error"; + +// Close reason prefix of a connect frame the backend refused for its size; the +// limit follows it (must match reflex/event_namespace.py). +const MAX_MESSAGE_SIZE_REASON = "max_message_size="; + +// The code of the error that reports such a refused hydrate (must match +// ClientErrorType.HYDRATION_PAYLOAD_TOO_LARGE in reflex_base/constants/event.py). +const HYDRATION_PAYLOAD_TOO_LARGE = "hydration_payload_too_large"; + +// Backend protocol version that speaks channels. A backend older than this +// closes the connection on a binary frame, so channels stay shut until the +// handshake proves otherwise. +const CHANNEL_PROTOCOL_VERSION = 2; + +// Binary frames align every attachment to this boundary, so a handler can +// view one as a Float64Array without copying. +const FRAME_ALIGNMENT = 8; + +// Reused across frames: both are stateless, and a channel streaming binary +// would otherwise allocate one per message. +const TEXT_ENCODER = new TextEncoder(); +const TEXT_DECODER = new TextDecoder(); + +// Messages a channel buffers while it is not open, oldest dropped first. +const MAX_QUEUED_CHANNEL_MESSAGES = 64; + +// Attachments one channel message may carry (must match +// reflex.channels.MAX_MESSAGE_BUFFERS). The backend closes the connection over +// a frame that breaks its limits, so they are enforced here, where the mistake +// is, rather than losing the app's socket for it. +const MAX_MESSAGE_BUFFERS = 64; + +// A channel handle reports its own lifecycle under these names, so a message +// may not use them (must match reflex.channels.RESERVED_EVENTS). Names +// starting with an underscore are reserved too: the channel protocol uses them. +const LIFECYCLE_EVENTS = new Set(["connect", "disconnect", "error"]); + +// Names a channel can be registered under (must match reflex.channels). +const CHANNEL_NAME_PATTERN = /^[A-Za-z0-9_./:-]{1,64}$/; + +/** + * Whether a channel message name is reserved for the handle or the protocol. + * @param event The message name. + * @returns Whether an application message may not use it. + */ +const isReservedEvent = (event) => + LIFECYCLE_EVENTS.has(event) || event.startsWith("_"); + +/** + * Report an error without throwing it into the code that found it. + * @param error The error to report. + */ +const report = (error) => (globalThis.reportError ?? console.error)(error); + +/** + * Serialize one channel message: binary when it carries attachments. + * @param event The message name. + * @param data The JSON metadata. + * @param channel The channel name. + * @param buffers Binary attachments (ArrayBuffers or typed arrays). + * @returns The serialized frame, text or binary. + */ +const channelFrame = (event, data, channel, buffers) => + buffers?.length + ? encodeChannelFrame(event, data, channel, buffers) + : stringifyFrame([event, data, channel]); + +/** + * Whether a parsed frame names a channel, which only a channel message does. + * @param message The parsed frame. + * @returns Whether it carries a channel name after the event and its payload. + */ +const isChannelMessage = (message) => + Array.isArray(message) && message.length > 2; + +/** + * Parse a serialized frame back into its message. + * @param frame The serialized text or binary frame. + * @returns The parsed message, or undefined if malformed. + */ +const decodeFrame = (frame) => + typeof frame === "string" + ? parseJsonLenient(frame, undefined) + : decodeChannelFrame(frame); + +/** + * The size a serialized frame occupies on the wire. + * @param frame The serialized text or binary frame. + * @returns The size in bytes. + */ +const frameByteLength = (frame) => + typeof frame === "string" + ? TEXT_ENCODER.encode(frame).byteLength + : frame.byteLength; + +/** + * Whether a serialized frame is over the backend's inbound message limit. + * + * Mirrors the check the backend applies before closing the connection: the + * limit counts UTF-8 bytes, and UTF-8 is 1-4 bytes per character, so a text + * frame is only measured exactly when the cheap bounds cannot settle it. + * @param frame The serialized text or binary frame. + * @param limit The limit in bytes; falsy while it is unknown. + * @returns Whether the frame exceeds it. + */ +const exceedsMessageLimit = (frame, limit) => { + if (!limit) { + return false; + } + if (typeof frame === "string") { + if (frame.length > limit) { + return true; + } + if (frame.length * 4 <= limit) { + return false; + } + } + return frameByteLength(frame) > limit; +}; + +/** + * Describe a frame over the backend's inbound message limit. + * @param subject What the frame carries, starting the sentence. + * @param frame The serialized frame. + * @param limit The limit in bytes. + * @returns The description, for the caller to finish with the consequence. + */ +const oversizedMessage = (subject, frame, limit) => + `${subject} is ${frameByteLength(frame)} bytes, over the ${limit} the ` + + "backend accepts (REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE)"; + +/** + * The error for a connect frame over the backend's inbound message limit. + * + * Every retry would carry the same hydrate, so the code marks it as final. + * @param frame The serialized connect frame. + * @param limit The limit in bytes; falsy when the backend did not name it. + * @returns The error. + */ +const hydrationTooLargeError = (frame, limit) => { + const error = new Error( + `Hydration payload (${frameByteLength(frame)} bytes) exceeds the server limit` + + (limit ? ` (${limit} bytes)` : "") + + ". Reduce browser storage data and reload.", + ); + error.code = HYDRATION_PAYLOAD_TOO_LARGE; + return error; +}; + +/** + * Serialize an outgoing frame. + * @param frame The frame array to serialize. + * @returns The JSON string. + */ +const stringifyFrame = (frame) => JSON.stringify(frame, undefinedToNull); + +/** + * Whether a value is binary data an attachment can carry. + * + * Checked by tag rather than instanceof, so buffers from another realm (an + * iframe, a worker) and shared buffers qualify. + * @param value The candidate attachment. + * @returns Whether it is an ArrayBuffer, SharedArrayBuffer or view of one. + */ +const isBinary = (value) => + ArrayBuffer.isView(value) || + ["[object ArrayBuffer]", "[object SharedArrayBuffer]"].includes( + Object.prototype.toString.call(value), + ); + +/** + * View any binary value as bytes without copying it. + * @param buffer An ArrayBuffer, SharedArrayBuffer, typed array or DataView. + * @returns A Uint8Array over the same memory. + */ +const asBytes = (buffer) => + ArrayBuffer.isView(buffer) + ? new Uint8Array(buffer.buffer, buffer.byteOffset, buffer.byteLength) + : new Uint8Array(buffer); + +/** + * Bytes of padding needed to reach the next attachment boundary. + * @param offset The current offset. + * @returns The padding length. + */ +const padding = (offset) => + (FRAME_ALIGNMENT - (offset % FRAME_ALIGNMENT)) % FRAME_ALIGNMENT; + +/** + * Serialize a channel message carrying binary attachments. + * @param event The message name. + * @param data The JSON metadata. + * @param channel The channel name. + * @param buffers The binary attachments. + * @returns The frame as an ArrayBuffer. + */ +export const encodeChannelFrame = (event, data, channel, buffers) => { + const views = buffers.map(asBytes); + const header = TEXT_ENCODER.encode( + stringifyFrame([event, data, channel, views.map((v) => v.byteLength)]), + ); + let size = 4 + header.byteLength; + for (const view of views) { + size += padding(size) + view.byteLength; + } + const frame = new ArrayBuffer(size); + const bytes = new Uint8Array(frame); + new DataView(frame).setUint32(0, header.byteLength, true); + bytes.set(header, 4); + let offset = 4 + header.byteLength; + for (const view of views) { + offset += padding(offset); + bytes.set(view, offset); + offset += view.byteLength; + } + return frame; +}; + +/** + * Deserialize a binary channel frame. + * @param frame The received ArrayBuffer. + * @returns [event, data, channel, buffers], or undefined if malformed. + */ +export const decodeChannelFrame = (frame) => { + if (frame.byteLength < 4) { + return undefined; + } + const headerSize = new DataView(frame).getUint32(0, true); + if (4 + headerSize > frame.byteLength) { + return undefined; + } + const header = parseJsonLenient( + TEXT_DECODER.decode(new Uint8Array(frame, 4, headerSize)), + undefined, + ); + if (!Array.isArray(header) || !Array.isArray(header[3])) { + return undefined; + } + const [event, data, channel, lengths] = header; + const buffers = []; + let offset = 4 + headerSize; + for (const length of lengths) { + offset += padding(offset); + if (offset + length > frame.byteLength) { + return undefined; + } + buffers.push(new Uint8Array(frame, offset, length)); + offset += length; + } + return [event, data, channel, buffers]; +}; + +/** + * Local handler registry shared by the transport and its channels. + * + * Handlers are keyed with socket.io's "$"-prefixed convention because + * upload.js reads `socket._callbacks.$event` directly. + */ +class LocalEmitter { + /** + * Create an emitter with no handlers registered. + */ + constructor() { + this._callbacks = {}; + } + + /** + * Register a handler for an event. + * @param event The event name. + * @param fn The handler function. + */ + on(event, fn) { + (this._callbacks["$" + event] ??= []).push(fn); + } + + /** + * Remove a handler, every handler for an event, or all of them. + * @param event The event name; omit to remove all handlers. + * @param fn The handler to remove; omit to remove all handlers for event. + */ + off(event, fn) { + if (event === undefined) { + this._callbacks = {}; + return; + } + if (fn === undefined) { + delete this._callbacks["$" + event]; + return; + } + const handlers = this._callbacks["$" + event]; + const ix = handlers ? handlers.indexOf(fn) : -1; + if (ix !== -1) { + handlers.splice(ix, 1); + } + } + + /** + * Invoke the registered handlers for a local event. + * @param event The event name. + * @param args The handler arguments. + */ + _emitLocal(event, ...args) { + const handlers = this._callbacks["$" + event]; + if (handlers === undefined) { + return; + } + // A copy: a handler may remove itself, or add one, while this runs. + for (const fn of [...handlers]) { + try { + fn(...args); + } catch (error) { + // A failing handler must not skip the others, nor abort the + // transport bookkeeping that is reporting the event. + report(error); + } + } + } +} + +// Channel handles by name, and the transport they currently ride. A handle +// outlives every transport: the event loop recreates the socket on remount and +// hot reload, and a channel must survive that without its consumer +// re-registering handlers. +const channels = new Map(); +let activeTransport = null; +let channelsUnsupportedReason = null; + +class ReflexChannel extends LocalEmitter { + /** + * Create a channel handle. Use getChannel() instead of constructing one. + * + * Handlers live on the channel rather than on the transport, whose table + * the event loop clears wholesale with socket.off() on unmount. + * @param name The channel name, matching the backend registration. + */ + constructor(name) { + super(); + this.name = name; + this.connected = false; + this._queue = []; + this._transport = null; + } + + /** + * Send a message to the channel's backend, buffering until it is open. + * @param event The message name. + * @param data The JSON metadata. + * @param buffers Binary attachments (ArrayBuffers or typed arrays). + */ + emit(event, data, buffers = []) { + // The backend closes the app's socket over a frame it cannot read, so a + // malformed message is refused here, where the mistake is. + if (typeof event !== "string" || event === "") { + throw new TypeError( + `Channel message name must be a non-empty string, not ${String(event)}.`, + ); + } + if (isReservedEvent(event)) { + throw new Error( + `Channel message name "${event}" is reserved: this handle reports its ` + + 'own lifecycle under it, names starting with "_" belong to the ' + + "channel protocol, and the backend refuses both.", + ); + } + if (!Array.isArray(buffers) || !buffers.every(isBinary)) { + throw new TypeError( + `Channel message "${event}" attachments must be an array of ` + + "ArrayBuffers or typed arrays; read a Blob with " + + "`await blob.arrayBuffer()` first.", + ); + } + if (buffers.length > MAX_MESSAGE_BUFFERS) { + throw new Error( + `Channel message "${event}" carries ${buffers.length} attachments, ` + + `over the ${MAX_MESSAGE_BUFFERS} a frame may hold.`, + ); + } + // Serialize now, connected or not: a queued frame must carry what was + // emitted, not whatever the caller's payload and typed arrays hold by the + // time the channel opens. + const frame = channelFrame(event, data, this.name, buffers); + const limit = this._transport?._maxMessageSize; + if (exceedsMessageLimit(frame, limit)) { + throw new Error( + oversizedMessage(`Channel message "${event}"`, frame, limit) + ".", + ); + } + if (this.connected && this._transport) { + this._transport._send(frame); + return; + } + if (this._queue.length >= MAX_QUEUED_CHANNEL_MESSAGES) { + // The backend is unreachable and the producer is not waiting for + // "connect"; drop the oldest rather than grow without bound. + this._queue.shift(); + } + this._queue.push(frame); + } + + /** + * Open the channel on a newly connected transport. + * @param transport The connected transport. + */ + _attach(transport) { + this._transport = transport; + transport._send(stringifyFrame([OPEN_MESSAGE, null, this.name])); + } + + /** + * Report the transport going away; queued messages survive for the next one. + * @param reason The disconnect reason. + */ + _detach(reason) { + this._transport = null; + if (this.connected) { + this.connected = false; + this._emitLocal("disconnect", reason); + } + } + + /** + * Report a queued message dropped for being over the backend's limit. + * @param frame The serialized frame. + * @param limit The limit in bytes. + */ + _dropOversized(frame, limit) { + this._emitLocal("error", { + code: "message_too_large", + message: + oversizedMessage("A queued channel message", frame, limit) + + "; it was dropped.", + }); + } + + /** + * Report that this deployment cannot carry channels at all. + * @param reason Why channels are unavailable. + */ + _unsupported(reason) { + this._emitLocal("error", { code: "channels_unsupported", message: reason }); + } + + /** + * Dispatch one message received for this channel. + * @param event The message name. + * @param data The JSON metadata. + * @param buffers Binary attachments, as Uint8Array views. + */ + _receive(event, data, buffers) { + if (event === OPENED_MESSAGE) { + if (this._transport === null || this.connected) { + // A straggler from a transport this channel is no longer on (the + // queue belongs to whichever it attaches to next), or the answer to + // an open sent twice. + return; + } + this.connected = true; + const queued = this._queue; + this._queue = []; + flushFrames(this._transport, queued); + this._emitLocal("connect"); + return; + } + if (event === CHANNEL_ERROR_MESSAGE) { + this._emitLocal("error", data); + return; + } + if (isReservedEvent(event)) { + // The backend rejects these names; a frame carrying one is not from a + // Reflex backend and must not be mistaken for the handle's own events. + console.error(`Ignoring channel message named "${event}" (reserved)`); + return; + } + this._emitLocal(event, data, buffers); + } +} + +/** + * Get the handle for a named channel, creating it on first use. + * @param name The channel name, matching the backend registration. + * @returns The channel handle. + */ +export const getChannel = (name) => { + if (typeof name !== "string" || !CHANNEL_NAME_PATTERN.test(name)) { + throw new TypeError( + `Invalid channel name ${String(name)}: expected 1-64 characters from ` + + "[A-Za-z0-9_./:-].", + ); + } + let channel = channels.get(name); + if (channel === undefined) { + channel = new ReflexChannel(name); + channels.set(name, channel); + if (channelsUnsupportedReason !== null) { + // After this turn: the caller registers its "error" handler on the + // handle we are still returning. + const reason = channelsUnsupportedReason; + queueMicrotask(() => channel._unsupported(reason)); + } else if (activeTransport !== null) { + channel._attach(activeTransport); + } + } + return channel; +}; + +/** + * Declare that channels cannot run against this backend, failing every handle. + * @param reason Why channels are unavailable. + */ +export const disableChannels = (reason) => { + if (channelsUnsupportedReason === reason) { + // Already reported; a remount must not fire "error" at every consumer again. + return; + } + channelsUnsupportedReason = reason; + activeTransport = null; + for (const channel of channels.values()) { + channel._detach(reason); + channel._unsupported(reason); + } +}; + +/** + * The error reported for an event message over the backend's limit. + * @param name The event name. + * @param frame The serialized frame. + * @param limit The limit in bytes. + * @returns The error. + */ +const oversizedEventError = (name, frame, limit) => + new Error( + oversizedMessage(`Message "${name}"`, frame, limit) + "; it was not sent.", + ); + +/** + * Report a queued frame dropped for being over the backend's limit: to its + * channel, or as an error for an event. + * @param frame The serialized text or binary frame. + * @param limit The limit in bytes. + */ +const dropOversizedFrame = (frame, limit) => { + const [event, data, channel] = decodeFrame(frame) ?? []; + if (channel === undefined) { + report(oversizedEventError(data?.name ?? event, frame, limit)); + } else { + channels.get(channel)?._dropOversized(frame, limit); + } +}; + +/** + * Send frames queued before the backend's limit was known, dropping any over + * it, which would cost the app its websocket and every frame flushed after it. + * @param transport The connected transport. + * @param frames The serialized frames, oldest first. + */ +const flushFrames = (transport, frames) => { + const limit = transport._maxMessageSize; + for (const frame of frames) { + if (exceedsMessageLimit(frame, limit)) { + dropOversizedFrame(frame, limit); + } else { + // The transport re-queues it if its socket closed mid-flush. + transport._send(frame); + } + } +}; + +/** + * Open every channel on a transport that just finished its handshake. + * @param transport The connected transport. + */ +const attachChannels = (transport) => { + channelsUnsupportedReason = null; + activeTransport = transport; + for (const channel of channels.values()) { + channel._attach(transport); + } +}; + +/** + * Detach every channel from a transport that went away. + * @param transport The transport reporting the disconnect. + * @param reason The disconnect reason. + */ +const detachChannels = (transport, reason) => { + if (activeTransport !== transport) { + return; + } + activeTransport = null; + for (const channel of channels.values()) { + channel._detach(reason); + } +}; + +export class ReflexWebSocket extends LocalEmitter { + /** + * Create the transport; it dials on open() or connect(). + * @param url The http(s) endpoint URL of the backend event route. + * @param opts Options: `query` (object) and `protocols` (subprotocol list). + */ + constructor(url, opts) { + super(); + this._url = new URL(url); + // Exposed as io.opts for socket.io API compatibility: state.js refreshes + // io.opts.query before reconnecting. + this.io = { opts }; + // Sent in the connect frame, like the auth of a socket.io CONNECT packet. + this.auth = {}; + this.connected = false; + this._ws = null; + // Set by connect(): from then on, a socket opens its session once open. + this._sessionRequested = false; + // Frames emitted while disconnected, flushed on (re)connect. + this._sendQueue = []; + this._watchdogTimer = null; + // Heartbeat window: 145 seconds (25s ping interval + 120s ping timeout) + // in ms; refined by the server handshake. + this._watchdogMs = (25 + 120) * 1000; + // Give up on a session not acknowledged within 20 seconds, so a + // connect_error always fires and retries proceed. + this._connectTimeoutMs = 20 * 1000; + this._connectTimer = null; + // The backend's inbound message limit, learned from the handshake. + this._maxMessageSize = null; + // The connect frame awaiting its handshake, measured only if refused. + this._connectFrame = null; + // Network emulation and OS offline do not interrupt established + // websockets, so treat the browser's offline event as a disconnect. + // Localhost connections keep working offline. + this._offlineListener = null; + if ( + typeof addEventListener === "function" && + this._url.hostname !== "localhost" + ) { + this._offlineListener = () => + this._dropConnection("transport close", { + description: "network connection lost", + }); + addEventListener("offline", this._offlineListener, false); + } + } + + /** + * Remove registered handlers. With no arguments, also releases the global + * offline listener (transport disposal). + * @param event The event name; omit to remove all handlers. + * @param fn The handler to remove; omit to remove all handlers for event. + */ + off(event, fn) { + super.off(event, fn); + if (event === undefined && this._offlineListener) { + removeEventListener("offline", this._offlineListener, false); + this._offlineListener = null; + } + } + + /** + * Open a session, dialing first if needed. + * + * The connect frame goes out as soon as the socket is open, without waiting + * for the handshake that acknowledges it, and carries `auth`. + */ + connect() { + if ( + this._sessionRequested && + this._ws && + this._ws.readyState <= WebSocket.OPEN + ) { + // The session is already opening or open on this socket. + return; + } + this._sessionRequested = true; + this.open(); + const ws = this._ws; + if (ws === null) { + // The dial failed outright; its connect_error is on the way. + return; + } + if (ws.readyState === WebSocket.OPEN) { + // Dialed ahead by open(). + this._sendConnect(); + } + clearTimeout(this._connectTimer); + this._connectTimer = setTimeout(() => { + if (this._ws === ws && !this.connected) { + ws.close(); + } + }, this._connectTimeoutMs); + } + + /** + * Dial the backend without opening a session, so connect() skips the dial. + */ + open() { + if (this._ws && this._ws.readyState <= WebSocket.OPEN) { + // CONNECTING (0) or OPEN (1): already dialing or connected. + return; + } + const url = new URL(this._url); + // Secure endpoints (https or already-wss) stay secure. + url.protocol = + url.protocol === "https:" || url.protocol === "wss:" ? "wss:" : "ws:"; + // Proxy rules written for Socket.IO route "/_event/*", which the bare + // path misses; the backend serves both. + url.pathname = url.pathname.replace(/\/?$/, "/"); + url.search = new URLSearchParams(this.io.opts.query ?? {}).toString(); + let ws; + try { + ws = new WebSocket(url, this.io.opts.protocols); + } catch (error) { + // Mixed content and CSP violations fail the constructor itself. Report + // it like a failed dial, after the caller has attached its handlers. + this._ws = null; + queueMicrotask(() => this._emitLocal("connect_error", error)); + return; + } + // Channel attachments arrive as binary frames; take them as ArrayBuffers + // so handlers can view them as typed arrays without a copy. + ws.binaryType = "arraybuffer"; + this._ws = ws; + ws.onopen = () => { + if (this._ws === ws && this._sessionRequested) { + this._sendConnect(); + } + }; + ws.onmessage = (msg) => { + if (this._ws === ws) { + // Ignore stragglers from a superseded connection. + this._onMessage(msg.data); + } + }; + ws.onclose = (event) => { + if (this._ws !== ws) { + // A newer connection or an explicit disconnect() superseded this one. + return; + } + const details = { code: event.code, reason: event.reason }; + if (!this._settleClose("transport close", details)) { + // Never handshaked: this was a failed connection attempt. A 1009 + // refused the connect frame for its size, naming the limit unless an + // older backend or a proxy sent it. + this._emitLocal( + "connect_error", + event.code === 1009 + ? hydrationTooLargeError( + this._connectFrame, + event.reason.startsWith(MAX_MESSAGE_SIZE_REASON) && + event.reason.slice(MAX_MESSAGE_SIZE_REASON.length), + ) + : new Error("websocket connection failed"), + ); + } + }; + } + + /** + * Send the connect frame that opens the session on the open socket. + * + * A frame over the limit an earlier handshake announced is refused instead, + * so it is not uploaded only for the backend to close the socket over it. + */ + _sendConnect() { + const frame = stringifyFrame([CONNECT_MESSAGE, this.auth]); + const limit = this._maxMessageSize; + if (exceedsMessageLimit(frame, limit)) { + this._teardown("io client disconnect", undefined); + this._emitLocal("connect_error", hydrationTooLargeError(frame, limit)); + return; + } + this._connectFrame = frame; + this._ws.send(frame); + } + + /** + * Close the connection deliberately (reason "io client disconnect"). + */ + disconnect() { + this._teardown("io client disconnect", undefined); + } + + /** + * Tear down the current connection, reporting the disconnect synchronously + * (onclose may never fire during page unload or while offline). + * @param reason The disconnect reason to report. + * @param details The disconnect details to report. + */ + _teardown(reason, details) { + const ws = this._ws; + if (!ws) { + return; + } + // Detach so the onclose handler does not double-report. + this._ws = null; + if (ws.readyState <= WebSocket.OPEN) { + ws.onclose = null; + ws.onmessage = null; + ws.close(1000); + } + this._settleClose(reason, details); + } + + /** + * Settle the transport's state once its socket is gone, then report it. + * @param reason The disconnect reason to report. + * @param details The disconnect details to report. + * @returns Whether a session was open, so a disconnect was reported. + */ + _settleClose(reason, details) { + clearTimeout(this._connectTimer); + clearTimeout(this._watchdogTimer); + const wasConnected = this.connected; + this.connected = false; + // Report only once the transport's own state is settled: the handlers + // run application code. + detachChannels(this, reason); + if (wasConnected) { + this._emitLocal("disconnect", reason, details); + } + return wasConnected; + } + + /** + * Send an event to the backend, buffering while disconnected. + * @param event The event name. + * @param data The event payload. + */ + emit(event, data) { + const frame = stringifyFrame([event, data]); + const limit = this._maxMessageSize; + if (exceedsMessageLimit(frame, limit)) { + // The backend closes the connection over an oversized frame, losing + // every update in flight with it; drop this one message instead. + report(oversizedEventError(data?.name ?? event, frame, limit)); + return; + } + this._send(frame); + } + + /** + * Write one serialized frame, buffering while disconnected. + * @param frame The serialized text or binary frame. + */ + _send(frame) { + if (this.connected && this._ws?.readyState === WebSocket.OPEN) { + this._ws.send(frame); + } else { + this._sendQueue.push(frame); + } + } + + /** + * Handle one incoming frame. + * @param data The raw frame: text, or an ArrayBuffer for a channel message. + */ + _onMessage(data) { + if (data instanceof ArrayBuffer) { + const frame = decodeChannelFrame(data); + if (frame === undefined) { + console.error("Failed to parse binary websocket message"); + return; + } + const [event, payload, channel, buffers] = frame; + channels.get(channel)?._receive(event, payload, buffers); + return; + } + const text = data; + const message = parseJsonLenient(text, undefined); + if (!Array.isArray(message)) { + console.error("Failed to parse websocket message", text); + return; + } + const [event, payload] = message; + if (event === PING_MESSAGE) { + // The server pings every interval regardless of traffic, so resetting + // the watchdog only here avoids timer churn per data message. + this._resetWatchdog(); + this._ws?.send(stringifyFrame([PONG_MESSAGE])); + return; + } + if (event === HANDSHAKE_MESSAGE) { + // Application-level liveness confirmed; adopt the server's heartbeat + // settings (sent in seconds, converted to ms) for the connection + // watchdog. + clearTimeout(this._connectTimer); + this._watchdogMs = (payload.ping_interval + payload.ping_timeout) * 1000; + this._resetWatchdog(); + this.connected = true; + this._connectFrame = null; + this._maxMessageSize = payload.max_message_size ?? null; + // Open the channels first: a frame queued on the transport for one of + // them has to arrive after its _open, or the backend has no session to + // dispatch it to and answers channel_not_open. + if ((payload.protocol ?? 1) >= CHANNEL_PROTOCOL_VERSION) { + attachChannels(this); + } else { + // A backend older than the channel protocol closes the connection on + // the binary frames channels use, so never open one against it. + disableChannels( + "This backend predates channel support; upgrade Reflex to use channels.", + ); + // A channel frame queued against an earlier connection would reach a + // backend that closes the socket over it, costing the app its state + // updates too. + this._sendQueue = this._sendQueue.filter( + (frame) => !isChannelMessage(decodeFrame(frame)), + ); + } + const queued = this._sendQueue; + this._sendQueue = []; + flushFrames(this, queued); + this._emitLocal("connect"); + return; + } + if (isChannelMessage(message)) { + channels.get(message[2])?._receive(event, payload, []); + return; + } + this._emitLocal(event, payload); + } + + /** + * Drop the socket as a transport failure, which the event loop reconnects + * from -- unlike disconnect(), which reports an intentional close. + * @param reason The disconnect reason to report. + * @param details The disconnect details to report. + */ + _dropConnection(reason, details) { + if (this.connected) { + // Reported now rather than from onclose: on a dead link the browser + // fires that only when its closing handshake times out, up to a minute + // later, and the transport would claim to be connected until then. + this._teardown(reason, details); + } + } + + /** + * (Re)arm the dead-connection watchdog; fires when no message (heartbeat + * included) arrives within the server's ping interval + timeout. + */ + _resetWatchdog() { + clearTimeout(this._watchdogTimer); + this._watchdogTimer = setTimeout( + () => this._dropConnection("ping timeout"), + this._watchdogMs, + ); + } +} diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js index aa79aedf662..3b4cceba6b2 100644 --- a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js +++ b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js @@ -1,5 +1,4 @@ // State management for Reflex web apps. -import io from "socket.io-client"; import { mergician } from "mergician"; import env from "$/env.json"; import reflexEnvironment from "$/reflex.json"; @@ -13,9 +12,18 @@ import { } from "react-router"; import { app, eventLoop } from "$/utils/context-registry"; import debounce from "$/utils/helpers/debounce"; -import { parseJson } from "$/utils/helpers/json"; +import { parseJsonLenient, undefinedToNull } from "$/utils/helpers/json"; import throttle from "$/utils/helpers/throttle"; import { uploadFiles } from "$/utils/helpers/upload"; +import { + ReflexWebSocket, + disableChannels, + getChannel, +} from "$/utils/helpers/websocket"; + +// Re-exported so components can reach a side channel through the module they +// already import for getBackendURL/getToken. +export { getChannel }; // Endpoint URLs. const EVENTURL = env.EVENT; @@ -37,6 +45,9 @@ const ERROR_TYPE_DISPATCH_MISSING = "dispatch_function_missing"; const ERROR_TYPE_STATE_UPDATE = "state_update_processing_error"; const ERROR_TYPE_HYDRATION_PAYLOAD_TOO_LARGE = "hydration_payload_too_large"; +// Connection errors kept for display, oldest dropped first. +const MAX_CONNECT_ERRORS = 10; + // These hostnames indicate that the backend and frontend are reachable via the same domain. const SAME_DOMAIN_HOSTNAMES = ["localhost", "0.0.0.0", "::", "0:0:0:0:0:0:0:0"]; @@ -180,23 +191,27 @@ export const isBackendDisabled = () => { }; /** - * Create a socket without starting its namespace or hydration events. + * Create a plain WebSocket transport without dialing it. * @param endpoint The backend URL. - * @param transports The configured transports. * @param token The session token the backend links the connection to. - * @returns The disconnected socket. + * @returns The idle transport. */ -const createSocket = (endpoint, transports, token) => - io(endpoint.href, { - path: endpoint.pathname, - transports, - protocols: [reflexEnvironment.version], - autoUnref: false, - autoConnect: false, +const createSocket = (endpoint, token) => + new ReflexWebSocket(endpoint.href, { query: { token }, - reconnection: false, + protocols: [reflexEnvironment.version], }); +/** + * Close a transport for good and remove its handlers. + * @param socket The transport. + */ +const disposeSocket = (socket) => { + socket.disconnect(); + // Also releases the plain WebSocket transport's offline listener. + socket.off(); +}; + /** * Measure encoded text the way the Engine.IO server limits a frame. * @param frame The encoded text. @@ -216,21 +231,34 @@ let warmSocket = null; let cancelWarmup = () => {}; let socketStarted = false; -/** Close an unclaimed transport and remove its cleanup handlers. */ -const discardWarmSocket = () => { +/** + * Claim the warmed-up transport, cancelling its cleanup. + * @returns The warm transport, or null if there is none. + */ +const takeWarmSocket = () => { const socket = warmSocket; warmSocket = null; cancelWarmup(); - socket?.disconnect(); + return socket; }; -// Start only the transport while React is still preparing to mount. The -// namespace stays disconnected until connect() installs all its handlers. +/** Close an unclaimed transport and remove its cleanup handlers. */ +const discardWarmSocket = () => { + const socket = takeWarmSocket(); + if (socket !== null) { + disposeSocket(socket); + } +}; + +// Dial the plain WebSocket transport while React is still preparing to mount. +// The session stays closed until connect() installs all its handlers. The +// socket.io transports load their client on connect, so they are not warmed. // Defer past module evaluation because context.js imports this module too. if (typeof window !== "undefined") { queueMicrotask(() => { if ( socketStarted || + env.SOCKETIO || Object.keys(app.initialState ?? {}).length <= 1 || isBackendDisabled() || document.visibilityState === "hidden" @@ -238,11 +266,7 @@ if (typeof window !== "undefined") { return; } try { - warmSocket = createSocket( - getBackendURL(EVENTURL), - [env.TRANSPORT], - peekToken(), - ); + warmSocket = createSocket(getBackendURL(EVENTURL), peekToken()); } catch { // Speculative setup may fail (for example, blocked session storage). // The normal connection path will report failures when the app mounts. @@ -254,9 +278,8 @@ if (typeof window !== "undefined") { clearTimeout(timeout); window.removeEventListener("pagehide", discardWarmSocket); }; - warmSocket.io.open((error) => { - if (error) discardWarmSocket(); - }); + // A dial that fails leaves nothing to clean up: connect() dials again. + warmSocket.open(); }); } @@ -857,7 +880,8 @@ const showNewVersionBanner = () => { /** * Handle the backend reporting that this bundle was compiled for another - * state tree than the one it runs (sent on connect, see EventNamespace.on_connect). + * state tree than the one it runs (sent on connect, see + * BaseEventNamespace.handle_connect). * * In prod the page reloads once per cooldown window. A mismatch inside the * window means the reload landed on the same skew (a cached bundle, or several @@ -887,6 +911,17 @@ const handleStateMismatch = (info) => { showNewVersionBanner(); }; +/** + * Append a connection error, keeping only the most recent ones. + * @param connectErrors The errors so far. + * @param error The new error. + * @returns The updated errors. + */ +const addConnectError = (connectErrors, error) => [ + ...connectErrors.slice(1 - MAX_CONNECT_ERRORS), + error, +]; + /** * Connect to a websocket and set the handlers. * @param socket The socket object to connect. @@ -906,6 +941,8 @@ export const connect = async ( navigate, params, ) => { + // Connecting (again) revokes a pending unmount cancellation. + socket.cancelConnect = false; // Socket already allocated, just reconnect it if needed. if (socket.current) { if (!socket.current.connected) { @@ -913,16 +950,23 @@ export const connect = async ( } return; } + // Another connect() call may be awaiting the socket.io-client import; + // don't create a second transport. + if (socket.connecting) { + return; + } + socket.connecting = true; // Get backend URL object from the endpoint. const endpoint = getBackendURL(EVENTURL); - // The hydrate event rides in the socket.io CONNECT packet, so the backend - // starts loading state as soon as the namespace connects instead of after - // an extra round trip for the connect acknowledgement. The keys are read by - // the backend as CompileVars.CONNECT_AUTH_EVENT and CONNECT_AUTH_STATE_SCHEMA; - // the state schema goes along on every (re)connect, since a reconnect can - // land on another backend version. + // The hydrate event rides in the connect request (the plain WebSocket's + // connect frame, or the socket.io CONNECT packet), so the backend starts + // loading state as soon as the session opens instead of after an extra + // round trip for its acknowledgement. The keys are read by the backend as + // CompileVars.CONNECT_AUTH_EVENT and CONNECT_AUTH_STATE_SCHEMA; the state + // schema goes along on every (re)connect, since a reconnect can land on + // another backend version. const bootAuth = (first) => { const boot_event = withRouterData(app.initialEvents(first)[0], params); recordSentStorageValues(boot_event); @@ -934,62 +978,81 @@ export const connect = async ( // Create the socket. A new session's token is saved here, once the app has // mounted, even when a transport warmed up with it earlier. socketStarted = true; - const session_token = getToken(); - if ( - warmSocket && - (warmSocket.io.opts.query.token !== session_token || - warmSocket.io.opts.transports.length !== transports.length || - transports.some( - (transport, i) => transport !== warmSocket.io.opts.transports[i], - )) - ) { - discardWarmSocket(); + try { + const session_token = getToken(); + if (warmSocket && warmSocket.io.opts.query.token !== session_token) { + discardWarmSocket(); + } + if (!env.SOCKETIO) { + // Default transport: plain WebSocket speaking the Reflex event protocol. + socket.current = + takeWarmSocket() ?? createSocket(endpoint, session_token); + } else { + // Socket.IO, over the engine.io transports given; the client library + // is only loaded when this transport is configured. + const { default: io } = await import("socket.io-client"); + if (socket.cancelConnect) { + // The event loop unmounted while the import was pending. + return; + } + socket.current = io(endpoint.href, { + path: endpoint.pathname, + transports, + protocols: [reflexEnvironment.version], + autoUnref: false, + autoConnect: false, + query: { token: session_token }, + reconnection: false, // Reconnection will be handled manually. + }); + // Ensure undefined fields in events are sent as null instead of removed + socket.current.io.encoder.replacer = undefinedToNull; + // The server drops a transport whose frame exceeds the limit its + // handshake announced, so a CONNECT packet (type 0) that carries an + // oversized hydrate is reported instead of sent. + const { encoder } = socket.current.io; + const encode = encoder.encode.bind(encoder); + encoder.encode = (packet) => { + const encoded = encode(packet); + if (packet.type !== 0) { + return encoded; + } + // socket.io-client is pinned; test_hydration_storage.py covers these internals. + const { engine } = socket.current.io; + const transport = engine.transport.name; + // The Engine.IO message type ("4") prefixes the encoded packet. + const size = 1 + frameSize(encoded[0], transport); + if (size > engine._maxPayload) { + const unit = transport === "websocket" ? "characters" : "bytes"; + const error = new Error( + `Hydration payload (${size} ${unit}) exceeds the server limit (${engine._maxPayload} ${unit}). Reduce browser storage data and reload.`, + ); + error.code = ERROR_TYPE_HYDRATION_PAYLOAD_TOO_LARGE; + socket.current.disconnect(); + retryConnect(error); + return []; + } + return encoded; + }; + // The decoder API expects false (not undefined) for unparsable input. + socket.current.io.decoder.tryParse = (str) => + parseJsonLenient(str, false); + // Channels are a plain-WebSocket protocol feature. + disableChannels( + 'Channels require transport="websocket" in rxconfig.py, not Socket.IO.', + ); + } + } catch (error) { + // Reported rather than rejected: the unhandled rejection handler adds an + // event, which reconnects, which fails again without end -- as a client + // chunk the browser failed to load does on every retry. + console.error("Failed to create the event socket:", error); + setConnectErrors((connectErrors) => addConnectError(connectErrors, error)); + return; + } finally { + socket.connecting = false; } - socket.current = - warmSocket ?? createSocket(endpoint, transports, session_token); - warmSocket = null; - cancelWarmup(); socket.current.auth = bootAuth(true); socket.current.wait_connect = !socket.current.connected; - // Ensure undefined fields in events are sent as null instead of removed - socket.current.io.encoder.replacer = (k, v) => (v === undefined ? null : v); - // The server drops a transport whose frame exceeds the limit its handshake - // announced, so a CONNECT packet (type 0) that carries an oversized hydrate - // is reported instead of sent. The page stays offline until browser storage - // shrinks and it reloads: wait_connect stays set, so reconnect() does not - // rebuild the hydrate. - const { encoder } = socket.current.io; - const encode = encoder.encode.bind(encoder); - encoder.encode = (packet) => { - const encoded = encode(packet); - if (packet.type !== 0) { - return encoded; - } - // socket.io-client is pinned; test_hydration_storage.py covers these internals. - const { engine } = socket.current.io; - const transport = engine.transport.name; - // The Engine.IO message type ("4") prefixes the encoded packet. - const size = 1 + frameSize(encoded[0], transport); - if (size > engine._maxPayload) { - const unit = transport === "websocket" ? "characters" : "bytes"; - const error = new Error( - `Hydration payload (${size} ${unit}) exceeds the server limit (${engine._maxPayload} ${unit}). Reduce browser storage data and reload.`, - ); - error.code = ERROR_TYPE_HYDRATION_PAYLOAD_TOO_LARGE; - console.error(error.message); - setConnectErrors((errors) => [...errors.slice(-9), error]); - socket.current.disconnect(); - return []; - } - return encoded; - }; - socket.current.io.decoder.tryParse = (str) => { - try { - return parseJson(str); - } catch { - return false; - } - }; // Set up a reconnect helper function socket.current.reconnect = () => { if ( @@ -1070,10 +1133,17 @@ export const connect = async ( }); // Record a failed connection attempt and retry it with incremental backoff. + // A hydrate over the server's message limit fails every retry the same way, + // so the page stays offline until browser storage shrinks and it reloads: + // wait_connect stays set, so reconnect() does not rebuild the hydrate. const retryConnect = (error) => { + setConnectErrors((connectErrors) => addConnectError(connectErrors, error)); + if (error.code === ERROR_TYPE_HYDRATION_PAYLOAD_TOO_LARGE) { + console.error(error.message); + return; + } socket.current.wait_connect = false; - failed_attempts = Math.min(failed_attempts + 1, 10); - setConnectErrors((connectErrors) => [...connectErrors.slice(-9), error]); + failed_attempts = Math.min(failed_attempts + 1, MAX_CONNECT_ERRORS); window.clearTimeout(retry_timer); retry_timer = window.setTimeout(() => { retry_timer = undefined; @@ -1088,6 +1158,7 @@ export const connect = async ( socket.current.on("disconnect", (reason, details) => { // A transport closed before the namespace connected is a failed attempt, // e.g. a CONNECT packet over the server or a proxy message size limit. + // Socket.IO reports it here; the plain WebSocket as connect_error. if (socket.current.wait_connect) { retryConnect( new Error(`connection closed before connecting (${reason})`), @@ -1189,6 +1260,9 @@ export const connect = async ( }); socket.current.on(STATE_MISMATCH_EVENT, handleStateMismatch); + // Track the handler on the ref so unmount cleanup can remove it; a + // surviving listener would resurrect a transport for the unmounted hook. + socket.visibilityHandler = checkVisibility; document.addEventListener("visibilitychange", checkVisibility); socket.current.connect(); }; @@ -1564,9 +1638,17 @@ export const useEventLoop = ( // Cleanup function. return () => { mounted.current = false; + // Abort a connect() that is still awaiting the socket.io-client import. + socket.cancelConnect = true; + if (socket.visibilityHandler) { + document.removeEventListener( + "visibilitychange", + socket.visibilityHandler, + ); + socket.visibilityHandler = null; + } if (socket.current) { - socket.current.disconnect(); - socket.current.off(); + disposeSocket(socket.current); socket.current = null; } }; diff --git a/packages/reflex-base/src/reflex_base/config.py b/packages/reflex-base/src/reflex_base/config.py index c3861dc4c3d..3e8895547ca 100644 --- a/packages/reflex-base/src/reflex_base/config.py +++ b/packages/reflex-base/src/reflex_base/config.py @@ -193,7 +193,7 @@ class BaseConfig: hydrate_fallback: Function returning the component shown while the page is hydrating (React Router's HydrateFallback), used when App.hydrate_fallback is not set. Formatted such that `from path_0.path_1... import path[-1]`, and calling it with no arguments would work. For example, "my_app.components.loading". plugins: List of plugins to use in the app. disable_plugins: List of plugin types to disable in the app. - transport: The transport method for client-server communication. + transport: The transport for client-server communication: "websocket" (plain WebSocket, default), or "socketio"/"polling" (Socket.IO; requires the reflex[socketio] extra). """ app_name: str @@ -283,7 +283,7 @@ class BaseConfig: disable_plugins: list[type[Plugin]] = dataclasses.field(default_factory=list) - transport: Literal["websocket", "polling"] = "websocket" + transport: Literal["websocket", "socketio", "polling"] = "websocket" # Whether to skip plugin checks. _skip_plugins_checks: bool = dataclasses.field(default=False, repr=False) @@ -391,6 +391,8 @@ def _post_init(self, **kwargs): self._normalize_frontend_compression_formats() + self._normalize_cors_allowed_origins() + if ( isinstance(self.trusted_proxy_hops, bool) or not isinstance(self.trusted_proxy_hops, int) @@ -592,6 +594,18 @@ def _normalize_disable_plugins(self): ) self.disable_plugins = normalized + def _normalize_cors_allowed_origins(self): + """Read a plain-string cors_allowed_origins as its environment variable is read. + + Left a string, membership in it is a substring test, so every origin + check would accept any origin the configured one contains. + """ + if isinstance(self.cors_allowed_origins, str): + field = self.__dataclass_fields__["cors_allowed_origins"] + self.cors_allowed_origins = interpret_env_var_value( + self.cors_allowed_origins, field.type, field.name + ) + def _normalize_frontend_compression_formats(self): """Normalize and validate configured frontend compression formats. diff --git a/packages/reflex-base/src/reflex_base/environment.py b/packages/reflex-base/src/reflex_base/environment.py index 00e38b9e882..039b227299e 100644 --- a/packages/reflex-base/src/reflex_base/environment.py +++ b/packages/reflex-base/src/reflex_base/environment.py @@ -812,6 +812,12 @@ class EnvironmentVariables: timedelta(seconds=constants.Ping.TIMEOUT) ) + # Whether the websocket server compresses messages with permessage-deflate + # (uvicorn only; granian does not negotiate it). Worth turning off for apps + # sending binary data over a channel: deflate barely shrinks it and costs + # milliseconds of event loop time per message. + REFLEX_SOCKET_PER_MESSAGE_DEFLATE: EnvVar[bool] = env_var(True) + # Whether to run Granian in a spawn process. This enables Reflex to pick up on environment variable changes between hot reloads. REFLEX_STRICT_HOT_RELOAD: EnvVar[bool] = env_var(False) diff --git a/packages/reflex-base/src/reflex_base/event/processor/event_processor.py b/packages/reflex-base/src/reflex_base/event/processor/event_processor.py index 51f7e1231be..0694f1c983b 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/event_processor.py +++ b/packages/reflex-base/src/reflex_base/event/processor/event_processor.py @@ -26,8 +26,8 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: - from reflex.app import EventNamespace from reflex.event import Event, EventSpec + from reflex.event_namespace import BaseEventNamespace if hasattr(asyncio, "QueueShutDown"): @@ -142,7 +142,7 @@ def configure( self, *, state_manager: StateManager | None = None, - event_namespace: EventNamespace | None = None, + event_namespace: BaseEventNamespace | None = None, ) -> Self: """Set up the event processor. diff --git a/packages/reflex-base/src/reflex_base/utils/log.py b/packages/reflex-base/src/reflex_base/utils/log.py index 725313fe889..9245a61f2f2 100644 --- a/packages/reflex-base/src/reflex_base/utils/log.py +++ b/packages/reflex-base/src/reflex_base/utils/log.py @@ -1057,11 +1057,15 @@ def _exclude_paths_from_frame_info() -> list[Path]: import click import granian - import socketio import typing_extensions import reflex_base + try: + import socketio + except ImportError: + socketio = None + try: import reflex as rx except ImportError: diff --git a/packages/reflex-bench/README.md b/packages/reflex-bench/README.md index ddcf95d6880..8c8ba3748ca 100644 --- a/packages/reflex-bench/README.md +++ b/packages/reflex-bench/README.md @@ -339,9 +339,12 @@ the only one so far. (`examples/playground`) answers and how fast. Each benchmark instance starts it once as a production backend (`reflex run --env prod --backend-only`, one granian worker), and each sample drives it for a few seconds with -`reflex_bench.drivers.events`, a socket.io load generator that speaks reflex's -event websocket over `websockets` (no python-socketio). Sessions are browser -tabs: each connects with its own token, hydrates like a page load, then sends +`reflex_bench.drivers.events`, a load generator that speaks reflex's event +websocket over `websockets`: the plain protocol of reflex's default transport, +or Socket.IO (without python-socketio) for releases before it and apps with +`transport="socketio"`. Each pool of sessions detects which by dialing the +plain endpoint, which a Socket.IO backend refuses. Sessions are browser tabs: +each connects with its own token, hydrates like a page load, then sends `BenchState.set_seq*` events whose delta echoes a sequence number. The grid holds the common points, about 30 minutes with the default policy. @@ -592,8 +595,14 @@ app and reflex version, so every metric is exact and a sample is one run. counts, as does any delta of another state. The extra data keeps the echo frame (`reply`), the largest frame, and `delta_bytes`, the bytes of each substate's part of the deltas re-serialized compactly. -- **Keepalives do not count**: an engine.io ping and its pong are timing, not - payload, so a slow run moves the same bytes as a fast one. +- **Keepalives do not count**: a ping and its pong are timing, not payload, so + a slow run moves the same bytes as a fast one. +- **The backend's protocol**: the session speaks the plain protocol or + Socket.IO, whichever the backend does, so a comparison across reflex's switch + includes the framing. On the plain protocol the hydration rides in the + connect frame: `hydrate_and_load` with an empty payload, as on a reconnect, + so the full state rather than the difference from the compiled defaults a + browser's first load asks for. - **`wire.navigate`** sends what the frontend sends on a route change: one `on_load_internal` with the new route's `router_data`. `update_vars_internal` goes first only once the browser holds a value of a client storage var (the diff --git a/packages/reflex-bench/budgets.json b/packages/reflex-bench/budgets.json index b2140bf5c30..4377690692c 100644 --- a/packages/reflex-bench/budgets.json +++ b/packages/reflex-bench/budgets.json @@ -5,7 +5,7 @@ "initial_gzip": 392000, "initial_brotli": 323000, "total_gzip": 4410000, - "chunks": 422, + "chunks": 423, "node_modules": 759000000 } } diff --git a/packages/reflex-bench/src/reflex_bench/drivers/browser.py b/packages/reflex-bench/src/reflex_bench/drivers/browser.py index 877b691416d..ea45522e373 100644 --- a/packages/reflex-bench/src/reflex_bench/drivers/browser.py +++ b/packages/reflex-bench/src/reflex_bench/drivers/browser.py @@ -58,8 +58,9 @@ HYDRATED = "#bench-hydrated" NAV_TIMEOUT_S = 120.0 _ANCHOR_READS = 3 -# A socket.io event frame: "42", an optional namespace, then ["name", ...]. -_SOCKETIO_EVENT = re.compile(r'42(?:/[^,]*,)?\["([^"]+)"') +# A message of the event websocket: ["name", ...], on Socket.IO after "42" and +# an optional namespace. +_MESSAGE = re.compile(r'(?:42(?:/[^,]*,)?)?\["([^"]+)"') # {"perf": ms since navigation start, "epoch": Date.now(), "alive": the # document's tag (see Tab.set_alive) or None}. @@ -133,7 +134,8 @@ class Tab: page: The Playwright page. cdp: A DevTools session of the page, e.g. for ``Network.setCacheDisabled``. ws_bytes: Payload bytes of the page's websocket frames, both directions. - ws_events: How often each socket.io event arrived, by name. + ws_events: How often each message of the event websocket arrived, by + name. transfer_bytes: Bytes received over the wire for the page's HTTP responses, headers included, as the encoded data length of the DevTools ``Network.loadingFinished`` events. @@ -201,14 +203,14 @@ def _on_frame_sent(self, payload: str | bytes) -> None: self.ws_bytes += _payload_size(payload) def _on_frame_received(self, payload: str | bytes) -> None: - """Count a received frame and the socket.io event it carries. + """Count a received frame and the message it carries. Args: payload: The frame's payload. """ self.ws_bytes += _payload_size(payload) - if isinstance(payload, str) and (event := _SOCKETIO_EVENT.match(payload)): - self.ws_events[event[1]] += 1 + if isinstance(payload, str) and (message := _MESSAGE.match(payload)): + self.ws_events[message[1]] += 1 def mark(self, id: str) -> Mark | None: """Read a mark. diff --git a/packages/reflex-bench/src/reflex_bench/drivers/echo_server.py b/packages/reflex-bench/src/reflex_bench/drivers/echo_server.py index 0752c581f76..f9e5b1d839c 100644 --- a/packages/reflex-bench/src/reflex_bench/drivers/echo_server.py +++ b/packages/reflex-bench/src/reflex_bench/drivers/echo_server.py @@ -1,10 +1,12 @@ -"""A socket.io server that answers the event generator like the playground does, without reflex. - -It speaks the frames of reflex's event websocket (see -:mod:`reflex_bench.drivers.events`): the engine.io handshake, the namespace ack, -pings, and one delta per event. The hydration events set ``is_hydrated``; every -other event gets a delta that echoes its ``payload["seq"]``. The generator's -calibration (``selftest.events.calibrate``) runs against it in a separate +"""A server that answers the event generator like the playground does, without reflex. + +It speaks either protocol of reflex's event websocket (see +:mod:`reflex_bench.drivers.events`) and, like reflex, refuses the websocket on +the other one's path, so the generator detects which: the plain protocol's +connect frame and handshake or Socket.IO's engine.io open and namespace ack, +pings, and one delta per event. The hydration events set ``is_hydrated``; +every other event gets a delta that echoes its ``payload["seq"]``. The +generator's calibration (``selftest.events.calibrate``) runs it in a separate process, and the generator's tests subclass it to script delays and faults. """ @@ -13,32 +15,44 @@ import asyncio import contextlib import json -import multiprocessing import os +import urllib.parse import uuid from collections.abc import AsyncIterator, Sequence +from http import HTTPStatus from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess from typing import Any from websockets.asyncio.server import ServerConnection, serve from websockets.exceptions import ConnectionClosed +from websockets.http11 import Request, Response from reflex_bench.drivers.events import ( - CONNECT_FRAME, - DISCONNECT_FRAME, - EVENT_PREFIX, + BOOT_KEY, + CODECS, + CONNECT, + DISCONNECTED, + EVENT, + HANDSHAKE, + HYDRATE_AND_LOAD_EVENT, HYDRATE_EVENT, HYDRATED_VAR, - KILL_GRACE_S, - NAMESPACE, ON_LOAD_EVENT, - PING, ROOT_STATE, - emit_frame, + SIO_OPEN_PREFIX, + WireProtocol, event_loop, + spawn_process, + stop_processes, + wait_readable, ) +# The plain protocol's version, announced in its handshake. +PROTOCOL_VERSION = 2 +# Reflex's default message size limit, announced by either protocol. +MAX_MESSAGE_BYTES = 1_000_000 + class EchoServer: """Answers the generator's frames as a reflex app does, one event at a time per connection.""" @@ -48,6 +62,7 @@ def __init__( *, delta_key: str, seq_var: str, + protocol: WireProtocol = "websocket", ping_interval_s: float = 25.0, ping_timeout_s: float = 120.0, ) -> None: @@ -56,13 +71,16 @@ def __init__( Args: delta_key: The state whose delta echoes the sequence number. seq_var: The var that echoes it. - ping_interval_s: Seconds between engine.io pings. - ping_timeout_s: The ping timeout announced in the handshake. + protocol: The protocol to speak. + ping_interval_s: Seconds between pings, by default reflex's. + ping_timeout_s: The ping timeout announced, by default reflex's. """ self.delta_key = delta_key self.seq_var = seq_var + self.codec = CODECS[protocol] self.ping_interval_s = ping_interval_s self.ping_timeout_s = ping_timeout_s + self._path = urllib.parse.urlsplit(self.codec.url("http://echo", "")).path @contextlib.asynccontextmanager async def serve(self, host: str, port: int) -> AsyncIterator[int]: @@ -83,36 +101,63 @@ async def serve(self, host: str, port: int) -> AsyncIterator[int]: compression=None, max_size=None, ping_interval=None, + process_request=self._route, ) as server: yield next(iter(server.sockets)).getsockname()[1] + def _route(self, connection: ServerConnection, request: Request) -> Response | None: + """Refuse the websocket on any path but the protocol's, as reflex does. + + Args: + connection: The connection. + request: Its upgrade request. + + Returns: + ``None`` to accept, or a 404. + """ + if urllib.parse.urlsplit(request.path).path == self._path: + return None + return connection.respond(HTTPStatus.NOT_FOUND, "Not Found\n") + async def handle(self, ws: ServerConnection) -> None: - """Serve one connection until it closes or leaves the namespace. + """Serve one connection until it closes or leaves the session. + + The codec decodes the client's frames too: Socket.IO's packets are the + same both ways, so its namespace join decodes as ``HANDSHAKE``. Args: ws: The connection. """ sid = uuid.uuid4().hex - await ws.send( - "0" - + json.dumps({ - "sid": sid, - "upgrades": [], - "pingInterval": int(self.ping_interval_s * 1000), - "pingTimeout": int(self.ping_timeout_s * 1000), - "maxPayload": 1_000_000, - }) - ) + codec = self.codec + if codec.greets: + await ws.send( + SIO_OPEN_PREFIX + + json.dumps({ + "sid": sid, + "upgrades": [], + "pingInterval": int(self.ping_interval_s * 1000), + "pingTimeout": int(self.ping_timeout_s * 1000), + "maxPayload": MAX_MESSAGE_BYTES, + }) + ) pinger = asyncio.create_task(self._ping(ws)) try: async for message in ws: - if isinstance(message, str) and message.startswith(EVENT_PREFIX): - name, *args = json.loads(message[len(EVENT_PREFIX) :]) - if name == "event": - await self.on_event(ws, args[0]) - elif message == CONNECT_FRAME: + args = codec.parse(message) + if args is None: + # Socket.IO's pong. + continue + name = args[0] + if name == EVENT: + await self.on_event(ws, args[1]) + elif name in {CONNECT, HANDSHAKE}: await self.on_connect(ws, sid) - elif message == DISCONNECT_FRAME: + # The plain protocol's connect carries the boot event. + auth = args[1] if len(args) > 1 else None + if isinstance(auth, dict) and BOOT_KEY in auth: + await self.on_event(ws, auth[BOOT_KEY]) + elif name == DISCONNECTED: return except ConnectionClosed: pass @@ -120,24 +165,42 @@ async def handle(self, ws: ServerConnection) -> None: pinger.cancel() async def _ping(self, ws: ServerConnection) -> None: - """Send engine.io pings; the generator answers them. + """Send pings; the generator answers them. Args: ws: The connection. """ + ping = self.codec.ping with contextlib.suppress(ConnectionClosed): while True: await asyncio.sleep(self.ping_interval_s) - await ws.send(PING) + await ws.send(ping) async def on_connect(self, ws: ServerConnection, sid: str) -> None: - """Acknowledge the namespace connect. + """Acknowledge the connect. + + A protocol that greets announced the heartbeat in its greeting, so its + ack carries the session id only, as python-socketio's does. Args: ws: The connection. sid: The session id. """ - await ws.send(f"40{NAMESPACE}," + json.dumps({"sid": sid})) + codec = self.codec + if codec.greets: + await ws.send(codec.join + json.dumps({"sid": sid})) + return + await ws.send( + codec.emit( + HANDSHAKE, + { + "ping_interval": self.ping_interval_s, + "ping_timeout": self.ping_timeout_s, + "protocol": PROTOCOL_VERSION, + "max_message_size": MAX_MESSAGE_BYTES, + }, + ) + ) async def on_event(self, ws: ServerConnection, event: dict[str, Any]) -> None: """Answer one event. @@ -155,28 +218,29 @@ def reply(self, event: dict[str, Any]) -> str: event: The event. Returns: - ``hydrate`` and ``on_load_internal`` get the root state's - ``is_hydrated`` (false, then true); any other event gets its - sequence number echoed. + The hydration events get the root state's ``is_hydrated``: false + for ``hydrate``, true for ``on_load_internal`` and + ``hydrate_and_load``; any other event gets its sequence number + echoed. """ name = event["name"] - if name in {HYDRATE_EVENT, ON_LOAD_EVENT}: - delta = {ROOT_STATE: {HYDRATED_VAR: name == ON_LOAD_EVENT}} + if name in {HYDRATE_EVENT, ON_LOAD_EVENT, HYDRATE_AND_LOAD_EVENT}: + delta = {ROOT_STATE: {HYDRATED_VAR: name != HYDRATE_EVENT}} else: delta = {self.delta_key: {self.seq_var: event["payload"]["seq"]}} - return emit_frame("event", {"delta": delta, "events": []}) + return self.codec.emit(EVENT, {"delta": delta, "events": []}) def _serve_process( - conn: Connection, delta_key: str, seq_var: str, cpus: Sequence[int] | None + delta_key: str, seq_var: str, cpus: Sequence[int] | None, conn: Connection ) -> None: """Run an echo server in this process until the parent closes the pipe or terminates it. Args: - conn: Receives the bound port; closing its other end stops the server. delta_key: The state whose delta echoes the sequence number. seq_var: The var that echoes it. cpus: The CPUs to run on (Linux), or ``None``. + conn: Receives the bound port; closing its other end stops the server. """ if cpus is not None and hasattr(os, "sched_setaffinity"): os.sched_setaffinity(0, cpus) @@ -187,9 +251,7 @@ async def main() -> None: conn.send(port) # The parent holds the pipe open while it needs the server, and # its end closes however the parent exits. - parent_gone = asyncio.Event() - asyncio.get_running_loop().add_reader(conn.fileno(), parent_gone.set) - await parent_gone.wait() + await wait_readable(conn) event_loop().run_until_complete(main()) @@ -223,17 +285,10 @@ def start(self, timeout: float = 30.0) -> str: Raises: RuntimeError: When it does not listen in time. """ - context = multiprocessing.get_context("spawn") - conn, child = context.Pipe() - proc = self._proc = context.Process( - target=_serve_process, - args=(child, *self._args), - name="reflex-bench-echo", - daemon=True, + proc, conn = spawn_process( + _serve_process, *self._args, name="reflex-bench-echo" ) - proc.start() - child.close() - self._conn = conn + self._proc, self._conn = proc, conn try: port = conn.recv() if conn.poll(timeout) else None except (EOFError, OSError): @@ -261,13 +316,7 @@ def pid(self) -> int: def stop(self) -> None: """Stop the server; safe to call again.""" - proc, conn = self._proc, self._conn - if conn is not None: - conn.close() - if proc is None: - return - proc.terminate() - proc.join(KILL_GRACE_S) - if proc.is_alive(): - proc.kill() - proc.join(KILL_GRACE_S) + if self._conn is not None: + self._conn.close() + if self._proc is not None: + stop_processes([self._proc]) diff --git a/packages/reflex-bench/src/reflex_bench/drivers/events.py b/packages/reflex-bench/src/reflex_bench/drivers/events.py index 3ca1670373b..b28365bae52 100644 --- a/packages/reflex-bench/src/reflex_bench/drivers/events.py +++ b/packages/reflex-bench/src/reflex_bench/drivers/events.py @@ -1,11 +1,25 @@ -"""An open-loop event generator that speaks reflex's socket.io protocol over ``websockets``. +"""An open-loop event generator that speaks reflex's event websocket over ``websockets``. It measures what a user of a reflex app waits for: the time from the moment an event *should* have been sent until the state delta that answers it arrives. -**Protocol.** reflex mounts python-socketio at ``/_event`` with the websocket -transport only (engine.io 4, socket.io 5), from 0.8.23 to HEAD. One session is -one websocket, one token and one state:: +**Protocols.** One session is one websocket, one token and one state. reflex's +default transport is a plain JSON protocol at ``/_event``; releases before it, +from 0.8.23 on, and apps with ``transport="socketio"`` mount python-socketio at +``/_event/`` with the websocket transport only (engine.io 4, socket.io 5). A +:class:`Codec` frames each; :func:`detect_protocol` tells which one a backend +speaks, as a Socket.IO backend refuses the upgrade of ``/_event``. The plain +protocol:: + + client GET /_event?token= + client ["_connect",{"event":{"name":"...hydrate_and_load",...}}] the boot event + server ["_handshake",{"ping_interval":...,"protocol":2,...}] (a "new_token" may come first) + client ["event",{"name":...,"payload":...,"router_data":...,"token":...}] + server ["event",{"delta":{...},"events":[...],"final":...}] + server ["_ping"] client ["_pong"] + client closes the websocket (the server closes with a code) + +Socket.IO:: client GET /_event/?EIO=4&transport=websocket&token= server 0{"sid":...,"upgrades":[],"pingInterval":25000,"pingTimeout":...} @@ -16,11 +30,13 @@ server 2 client 3 engine.io ping and pong client 41/_event, leave the namespace -A session is primed like a page load, with ``hydrate`` and ``on_load_internal``, -until a delta sets ``is_hydrated``, then with the plan's :class:`LinkEvent` if -it has one (joining a shared state token), until a delta of the linked state -acknowledges it. Every event carries the token: 0.8.23 requires it, HEAD -ignores it. +A session is primed like a page load until a delta sets ``is_hydrated``: on the +plain protocol its connect frame carries ``hydrate_and_load`` as the boot event, +as the frontend's does on a reconnect (an empty payload: the full state); on +Socket.IO it sends ``hydrate`` and ``on_load_internal`` once joined. Then it +sends the plan's :class:`LinkEvent` if it has one (joining a shared state +token), until a delta of the linked state acknowledges it. Every event carries +the token: 0.8.23 requires it, later releases ignore it. **Open, closed and fan-out loop.** In the open loop each session sends on a fixed schedule whatever the server does, late events are sent at once (never @@ -44,7 +60,8 @@ session is primed, and merges what they measured into a :class:`LoadResult`. :func:`open_sessions` primes sessions on the caller's loop without a load, for measurements of idle sessions; :func:`hold_sessions` does the same in a spawned -process and keeps them open until told to close. +process and keeps them open until told to close. Each pool of sessions detects +the protocol once, unless its :class:`Endpoint` names it. """ from __future__ import annotations @@ -68,6 +85,7 @@ import traceback import urllib.parse import uuid +from abc import ABC, abstractmethod from collections.abc import ( AsyncIterator, Awaitable, @@ -80,10 +98,10 @@ from dataclasses import dataclass from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess -from typing import Any, Literal +from typing import Any, ClassVar, Literal from websockets.asyncio.client import ClientConnection, connect -from websockets.exceptions import ConnectionClosed +from websockets.exceptions import ConnectionClosed, InvalidStatus from reflex_bench import stats @@ -91,22 +109,39 @@ import resource Mode = Literal["open", "closed", "fanout"] +WireProtocol = Literal["websocket", "socketio"] NAMESPACE = "/_event" ROOT_STATE = "reflex___state____state" HYDRATE_EVENT = f"{ROOT_STATE}.hydrate" +HYDRATE_AND_LOAD_EVENT = f"{ROOT_STATE}.hydrate_and_load" ON_LOAD_EVENT = ( f"{ROOT_STATE}.reflex___state____on_load_internal_state.on_load_internal" ) HYDRATED_VAR = "is_hydrated_rx_state_" -OPEN_PREFIX = "0" -CLOSE = "1" -PING = "2" -PONG = "3" -CONNECT_FRAME = f"40{NAMESPACE}," -DISCONNECT_FRAME = f"41{NAMESPACE}," -EVENT_PREFIX = f"42{NAMESPACE}," -CONNECT_ERROR_PREFIX = f"44{NAMESPACE}," +# Message names of the plain protocol (reflex.event_namespace). +EVENT = "event" +CONNECT = "_connect" +HANDSHAKE = "_handshake" +PING = "_ping" +PONG = "_pong" +# The key of the boot event in the connect frame (CompileVars.CONNECT_AUTH_EVENT). +BOOT_KEY = "event" +# Socket.IO packets the plain protocol has no message for decode to these names. +OPENED = "_sio_open" +REFUSED = "_sio_refused" +DISCONNECTED = "_sio_disconnect" +CONTROL = frozenset({CONNECT, HANDSHAKE, PING, PONG, OPENED, REFUSED, DISCONNECTED}) +# Socket.IO packets on reflex's namespace. +SIO_OPEN_PREFIX = "0" +SIO_CLOSE = "1" +SIO_PING = "2" +SIO_PONG = "3" +SIO_CONNECT_FRAME = f"40{NAMESPACE}," +SIO_DISCONNECT_FRAME = f"41{NAMESPACE}," +SIO_EVENT_PREFIX = f"42{NAMESPACE}," +SIO_CONNECT_ERROR_PREFIX = f"44{NAMESPACE}," +_COMPACT = (",", ":") PRIME_TIMEOUT_S = 30.0 PRIME_CONCURRENCY = 64 @@ -138,27 +173,15 @@ class LoadError(RuntimeError): """The load did not run: sessions could not start, a generator process failed, or it was stopped.""" -def emit_frame(*args: Any) -> str: - """Encode a socket.io event on reflex's namespace. - - Args: - *args: The event name and its arguments. - - Returns: - ``42/_event,[name, ...]``. - """ - return EVENT_PREFIX + json.dumps(list(args), separators=(",", ":")) - - -def event_frame( +def event( name: str, payload: Mapping[str, Any] | None, *, token: str, pathname: str = "/", query: Mapping[str, str] | None = None, -) -> str: - """Encode an event as the reflex frontend sends it. +) -> dict[str, Any]: + """Build an event as the reflex frontend sends it. Args: name: The full event handler name. @@ -169,36 +192,389 @@ def event_frame( ``/item/42`` of a ``/item/[item_id]`` page. Returns: - The frame. + The event. """ router_data = {"pathname": pathname, "asPath": pathname, "query": dict(query or {})} - return emit_frame( - "event", - {"name": name, "payload": payload, "router_data": router_data, "token": token}, - ) + return { + "name": name, + "payload": payload, + "router_data": router_data, + "token": token, + } + + +def _message(text: str) -> list[Any] | None: + """Decode a JSON message ``[name, *args]``. + + Args: + text: The JSON. + + Returns: + The message, or ``None`` when the text is not one. + """ + try: + message = json.loads(text) + except ValueError: + return None + if isinstance(message, list) and message and isinstance(message[0], str): + return message + return None -def event_url(backend_url: str, token: str) -> str: - """Build the websocket URL of a session. +def _ws_url(backend_url: str, path: str, query: Mapping[str, Any]) -> str: + """Build a websocket URL on a backend. Args: backend_url: The backend's URL, e.g. ``http://localhost:8000``. - token: The session's token. + path: The path below the backend's. + query: The query parameters. Returns: - E.g. ``ws://localhost:8000/_event/?EIO=4&transport=websocket&token=``. + The URL, ``wss`` for an ``https`` backend. """ parts = urllib.parse.urlsplit(backend_url) - query = urllib.parse.urlencode({"EIO": 4, "transport": "websocket", "token": token}) return urllib.parse.urlunsplit(( "wss" if parts.scheme == "https" else "ws", parts.netloc, - f"{parts.path.rstrip('/')}{NAMESPACE}/", - query, + f"{parts.path.rstrip('/')}{path}", + urllib.parse.urlencode(query), "", )) +class Codec(ABC): + """Frames reflex's event websocket in one protocol. + + Received frames decode to messages ``[name, *args]`` named as in the plain + protocol; Socket.IO's open, refusal and disconnect packets, which it has + no message for, decode to :data:`OPENED`, :data:`REFUSED` and + :data:`DISCONNECTED`. + """ + + protocol: ClassVar[WireProtocol] + # Whether the server greets with an engine.io open before the connect. + greets: ClassVar[bool] + # The hydration events to send once the server acknowledged the connect. + load_events: ClassVar[tuple[str, ...]] + # The frame that opens a session without hydrating it. + join: ClassVar[str] + # The server's keepalive. + ping: ClassVar[str] + # The answer to a ping. + pong: ClassVar[str] + # The frame that leaves the session before the websocket closes, if any. + leave: ClassVar[str | None] + + @abstractmethod + def url(self, backend_url: str, token: str) -> str: + """Build the websocket URL of a session. + + Args: + backend_url: The backend's URL, e.g. ``http://localhost:8000``. + token: The session's token. + + Returns: + The URL. + """ + + @abstractmethod + def emit(self, *args: Any) -> str: + """Encode a message. + + Args: + *args: The message name and its arguments. + + Returns: + The frame. + """ + + @abstractmethod + def parse(self, message: str | bytes) -> list[Any] | None: + """Decode a received frame. + + Args: + message: The frame. + + Returns: + The message ``[name, *args]``, or ``None`` for a frame the + protocol does not have. + """ + + @abstractmethod + def connect_frame( + self, *, token: str, pathname: str, query: Mapping[str, str] | None = None + ) -> str: + """Encode the frame that opens a session. + + Args: + token: The session's token. + pathname: The page route the session loads. + query: The route's parameters. + + Returns: + The frame; the plain protocol's carries the hydration as its boot + event. + """ + + def event_frame( + self, + name: str, + payload: Mapping[str, Any] | None, + *, + token: str, + pathname: str = "/", + query: Mapping[str, str] | None = None, + ) -> str: + """Encode an event as the reflex frontend sends it. + + Args: + name: The full event handler name. + payload: The handler's arguments. + token: The session's token. + pathname: The page route the event comes from. + query: The route's parameters. + + Returns: + The frame. + """ + return self.emit( + EVENT, event(name, payload, token=token, pathname=pathname, query=query) + ) + + +class WebsocketCodec(Codec): + """reflex's plain protocol: every frame is a JSON message ``[name, *args]``.""" + + protocol = "websocket" + greets = False + load_events = () + join = json.dumps([CONNECT, {}], separators=_COMPACT) + ping = json.dumps([PING]) + pong = json.dumps([PONG]) + leave = None + + def url(self, backend_url: str, token: str) -> str: + """Build the websocket URL of a session. + + Args: + backend_url: The backend's URL, e.g. ``http://localhost:8000``. + token: The session's token. + + Returns: + E.g. ``ws://localhost:8000/_event?token=``. + """ + return _ws_url(backend_url, NAMESPACE, {"token": token}) + + def emit(self, *args: Any) -> str: + """Encode a message. + + Args: + *args: The message name and its arguments. + + Returns: + ``[name, ...]``. + """ + return json.dumps(args, separators=_COMPACT) + + def parse(self, message: str | bytes) -> list[Any] | None: + """Decode a received frame. + + Args: + message: The frame. + + Returns: + The message, or ``None`` for a binary frame or one that is not a + message. + """ + return _message(message) if isinstance(message, str) else None + + def connect_frame( + self, *, token: str, pathname: str, query: Mapping[str, str] | None = None + ) -> str: + """Encode the connect frame, with ``hydrate_and_load`` as its boot event. + + Args: + token: The session's token. + pathname: The page route the session loads. + query: The route's parameters. + + Returns: + ``["_connect",{"event":{...}}]``. + """ + boot = event( + HYDRATE_AND_LOAD_EVENT, {}, token=token, pathname=pathname, query=query + ) + return self.emit(CONNECT, {BOOT_KEY: boot}) + + +class SocketIOCodec(Codec): + """python-socketio on reflex's namespace: engine.io 4 and socket.io 5 packets.""" + + protocol = "socketio" + greets = True + load_events = (HYDRATE_EVENT, ON_LOAD_EVENT) + join = SIO_CONNECT_FRAME + ping = SIO_PING + pong = SIO_PONG + leave = SIO_DISCONNECT_FRAME + + def url(self, backend_url: str, token: str) -> str: + """Build the websocket URL of a session. + + Args: + backend_url: The backend's URL, e.g. ``http://localhost:8000``. + token: The session's token. + + Returns: + E.g. ``ws://localhost:8000/_event/?EIO=4&transport=websocket&token=``. + """ + query = {"EIO": 4, "transport": "websocket", "token": token} + return _ws_url(backend_url, f"{NAMESPACE}/", query) + + def emit(self, *args: Any) -> str: + """Encode a socket.io event on reflex's namespace. + + Args: + *args: The event name and its arguments. + + Returns: + ``42/_event,[name, ...]``. + """ + return SIO_EVENT_PREFIX + json.dumps(args, separators=_COMPACT) + + def parse(self, message: str | bytes) -> list[Any] | None: + """Decode a received packet. + + Args: + message: The frame. + + Returns: + A socket.io event's arguments; :data:`PING` for an engine.io ping, + :data:`HANDSHAKE` for the namespace ack, :data:`OPENED`, + :data:`REFUSED` and :data:`DISCONNECTED` with the packet's data; + ``None`` for anything else. + """ + if not isinstance(message, str): + return None + if message.startswith(SIO_EVENT_PREFIX): + return _message(message[len(SIO_EVENT_PREFIX) :]) + if message == SIO_PING: + return [PING] + if message.startswith(SIO_CONNECT_FRAME): + return [HANDSHAKE, message[len(SIO_CONNECT_FRAME) :]] + if message in {SIO_DISCONNECT_FRAME, SIO_CLOSE}: + return [DISCONNECTED] + if message.startswith(SIO_CONNECT_ERROR_PREFIX): + return [REFUSED, message[len(SIO_CONNECT_ERROR_PREFIX) :]] + if message.startswith(SIO_OPEN_PREFIX): + return [OPENED, message[len(SIO_OPEN_PREFIX) :]] + return None + + def connect_frame( + self, *, token: str, pathname: str, query: Mapping[str, str] | None = None + ) -> str: + """Encode the namespace join; the hydration events follow its ack. + + Args: + token: The session's token. + pathname: The page route the session loads. + query: The route's parameters. + + Returns: + ``40/_event,``. + """ + return self.join + + +CODECS: dict[WireProtocol, Codec] = { + codec.protocol: codec for codec in (WebsocketCodec(), SocketIOCodec()) +} + + +def end_reason(args: list[Any], frame: str | bytes) -> str | None: + """Tell why a message ends the session, if it does. + + Args: + args: The message. + frame: Its raw frame. + + Returns: + Why the server refused or disconnected the session, or ``None``. + """ + name = args[0] + if name == REFUSED: + return f"the server refused {NAMESPACE}: {args[1]}" + if name == DISCONNECTED: + return f"the server disconnected the session ({frame!r})" + return None + + +async def leave(ws: ClientConnection, codec: Codec) -> None: + """Leave a session before its websocket closes, on a protocol that has a leave frame. + + Args: + ws: The websocket. + codec: Its protocol. + """ + if codec.leave is not None: + with contextlib.suppress(ConnectionClosed): + await ws.send(codec.leave) + + +def is_hydrated(delta: Mapping[str, Any]) -> bool: + """Tell whether a delta completes a hydration. + + Args: + delta: The delta. + + Returns: + Whether it sets the root state's ``is_hydrated``. + """ + root = delta.get(ROOT_STATE) + return root is not None and root.get(HYDRATED_VAR) is True + + +async def detect_protocol( + backend_url: str, timeout: float = PRIME_TIMEOUT_S +) -> WireProtocol: + """Find the protocol of a backend's event websocket. + + The plain protocol's endpoint accepts the websocket before the client + sends anything; a Socket.IO backend refuses the upgrade of ``/_event``, + its mount being ``/_event/``. + + Args: + backend_url: The backend's URL, e.g. ``http://localhost:8000``. + timeout: Seconds to wait for the websocket to open. + + Returns: + ``websocket`` or ``socketio``. + """ + try: + async with connect( + CODECS["websocket"].url(backend_url, str(uuid.uuid4())), + proxy=None, + open_timeout=timeout, + ping_interval=None, + close_timeout=1, + ): + return "websocket" + except InvalidStatus: + return "socketio" + + +async def resolve_codec(endpoint: Endpoint) -> Codec: + """Find the codec of an endpoint's protocol, detecting it if the endpoint names none. + + Args: + endpoint: The endpoint. + + Returns: + The codec. + """ + return CODECS[endpoint.protocol or await detect_protocol(endpoint.backend_url)] + + def seq_payload(seq: int) -> dict[str, int]: """Build the payload of the playground's ``set_seq*`` handlers. @@ -300,16 +676,16 @@ class LinkEvent: class Endpoint: """Where sessions connect: a reflex backend and the page they load. - The protocol is the same from 0.8.23 to HEAD (every event carries the - token), so nothing here depends on the reflex version. - Attributes: backend_url: The backend's URL, e.g. ``http://localhost:8000``. pathname: The page route the sessions hydrate and send events from. + protocol: The protocol of its event websocket, or ``None`` to detect + it (:func:`detect_protocol`) once per pool of sessions. """ backend_url: str pathname: str = "/" + protocol: WireProtocol | None = None @dataclass(frozen=True, kw_only=True) @@ -524,11 +900,12 @@ class _Frames: """Encodes a session's events; only the payload changes between them.""" def __init__( - self, shape: EventShape, token: str, pathname: str, index: int + self, codec: Codec, shape: EventShape, token: str, pathname: str, index: int ) -> None: """Split the event frame around its payload. Args: + codec: The protocol's codec. shape: The event. token: The session's token. pathname: The page route. @@ -537,7 +914,7 @@ def __init__( """ self._payload = shape.payload self._client = {"client": index} if shape.client_var is not None else {} - template = event_frame(shape.name, None, token=token, pathname=pathname) + template = codec.event_frame(shape.name, None, token=token, pathname=pathname) head, marker, self._tail = template.partition('"payload":null') self._head = head + marker.removesuffix("null") @@ -551,7 +928,7 @@ def event(self, seq: int) -> str: The frame. """ payload = json.dumps( - {**self._payload(seq), **self._client}, separators=(",", ":") + {**self._payload(seq), **self._client}, separators=_COMPACT ) return f"{self._head}{payload}{self._tail}" @@ -584,7 +961,8 @@ def __init__( self.shape = shape self.link = link self.token = str(uuid.uuid4()) - self.frames = self._frames() + self.codec: Codec | None = None + self.frames: _Frames | None = None self.ws: ClientConnection | None = None self.error: str | None = None self.prime_ns: int | None = None @@ -610,7 +988,10 @@ def _frames(self) -> _Frames | None: """ if self.shape is None: return None - return _Frames(self.shape, self.token, self.endpoint.pathname, self.index) + assert self.codec is not None + return _Frames( + self.codec, self.shape, self.token, self.endpoint.pathname, self.index + ) def _fail(self, reason: str) -> None: """Mark the session failed and wake whatever waits on it. @@ -624,74 +1005,81 @@ def _fail(self, reason: str) -> None: self._linked.set() self._idle.set() - async def prime(self, gate: asyncio.Semaphore) -> None: - """Connect, join the namespace, hydrate like a page load and link; failures are recorded. + async def prime(self, gate: asyncio.Semaphore, codec: Codec) -> None: + """Connect, open the session, hydrate like a page load and link; failures are recorded. Args: gate: Limits how many sessions connect at once. + codec: The endpoint's protocol. """ + self.codec = codec + self.frames = self._frames() async with gate: try: - await asyncio.wait_for(self._open(), PRIME_TIMEOUT_S) + await asyncio.wait_for(self._open(codec), PRIME_TIMEOUT_S) except TimeoutError: step = "linked" if self._hydrated.is_set() else "hydrated" self._fail(f"not {step} within {PRIME_TIMEOUT_S:g} s") except Exception as exc: self._fail(f"could not connect: {type(exc).__name__}: {exc}") - async def _open(self) -> None: - """Run the handshake, the hydration events and the link event. + async def _open(self, codec: Codec) -> None: + """Run the handshake, the hydration and the link event. + + Args: + codec: The endpoint's protocol. Raises: - ProtocolError: When the server does not open or refuses the namespace. + ProtocolError: When the server does not open or refuses the session. ConnectionError: When the session failed while hydrating or linking. """ started = time.perf_counter_ns() ws = self.ws = await connect( - event_url(self.endpoint.backend_url, self.token), + codec.url(self.endpoint.backend_url, self.token), proxy=None, open_timeout=None, max_size=MAX_FRAME_BYTES, ping_interval=None, close_timeout=1, ) - opened = await ws.recv() - if not (isinstance(opened, str) and opened.startswith(OPEN_PREFIX)): - msg = f"the server did not open an engine.io session: {opened[:40]!r}" - raise ProtocolError(msg) - await ws.send(CONNECT_FRAME) + pathname = self.endpoint.pathname + if codec.greets: + greeting = await ws.recv() + opened = codec.parse(greeting) + if opened is None or opened[0] != OPENED: + msg = f"the server did not open an engine.io session: {greeting[:40]!r}" + raise ProtocolError(msg) + await ws.send(codec.connect_frame(token=self.token, pathname=pathname)) while True: message = await ws.recv() - if isinstance(message, str): - if message.startswith(EVENT_PREFIX): - # python-socketio sends events of the connect handler first. - self._on_emit(json.loads(message[len(EVENT_PREFIX) :]), message, 0) - continue - if message.startswith(CONNECT_FRAME): - break - if message.startswith(CONNECT_ERROR_PREFIX): - reason = message[len(CONNECT_ERROR_PREFIX) :] - msg = f"the server refused {NAMESPACE}: {reason}" - raise ProtocolError(msg) - if message == PING: - await ws.send(PONG) - continue - msg = f"unexpected frame in the handshake: {message[:40]!r}" - raise ProtocolError(msg) - self._reader = asyncio.create_task(self._read()) - pathname = self.endpoint.pathname - await ws.send( - event_frame(HYDRATE_EVENT, {}, token=self.token, pathname=pathname) - ) - await ws.send( - event_frame(ON_LOAD_EVENT, {}, token=self.token, pathname=pathname) - ) + args = codec.parse(message) + if args is None: + msg = f"unexpected frame in the handshake: {message[:40]!r}" + raise ProtocolError(msg) + name = args[0] + if name == HANDSHAKE: + break + if name == PING: + await ws.send(codec.pong) + elif name in CONTROL: + msg = end_reason(args, message) or ( + f"unexpected frame in the handshake: {message[:40]!r}" + ) + raise ProtocolError(msg) + else: + # Events of the connect come first, e.g. a new token. + self._on_emit(args, message, 0) + self._reader = asyncio.create_task(self._read(ws, codec)) + for name in codec.load_events: + await ws.send( + codec.event_frame(name, {}, token=self.token, pathname=pathname) + ) await self._hydrated.wait() if self.error is not None: raise ConnectionError(self.error) if (link := self.link) is not None: await ws.send( - event_frame( + codec.event_frame( link.name, link.payload, token=self.token, pathname=pathname ) ) @@ -700,43 +1088,49 @@ async def _open(self) -> None: raise ConnectionError(self.error) self.prime_ns = time.perf_counter_ns() - started - async def _read(self) -> None: - """Receive frames until the websocket closes: answers, pings, tokens and disconnects.""" - ws = self.ws - assert ws is not None - start = len(EVENT_PREFIX) + async def _read(self, ws: ClientConnection, codec: Codec) -> None: + """Receive frames until the websocket closes: answers, pings, tokens and disconnects. + + Args: + ws: The websocket. + codec: Its protocol. + """ + parse = codec.parse try: async for message in ws: now = time.perf_counter_ns() - if isinstance(message, str) and message.startswith(EVENT_PREFIX): - self._on_emit(json.loads(message[start:]), message, now) - elif message == PING: - await ws.send(PONG) - elif message in {DISCONNECT_FRAME, CLOSE}: - self._fail(f"the server disconnected the session ({message!r})") - return - else: + args = parse(message) + if args is None: self._fail(f"unexpected frame: {message[:40]!r}") return + name = args[0] + if name == PING: + await ws.send(codec.pong) + elif name in CONTROL: + self._fail( + end_reason(args, message) + or f"unexpected frame: {message[:40]!r}" + ) + return + else: + self._on_emit(args, message, now) except ConnectionClosed as exc: if not self._closing: self._fail(f"the websocket closed: {exc}") - except ValueError as exc: - self._fail(f"{type(exc).__name__}: {exc}") else: if not self._closing: self._fail("the websocket closed") - def _on_emit(self, args: list[Any], frame: str, now: int) -> None: - """Handle a socket.io event from the server. + def _on_emit(self, args: list[Any], frame: str | bytes, now: int) -> None: + """Handle a message from the server. Args: - args: The event name and arguments. + args: The message name and arguments. frame: The raw frame. now: When it arrived. """ name = args[0] - if name == "event": + if name == EVENT: delta = args[1].get("delta") if not delta: return @@ -750,7 +1144,7 @@ def _on_emit(self, args: list[Any], frame: str, now: int) -> None: or state.get(shape.client_var) == self.index ) ): - if self.reply_frame is None: + if self.reply_frame is None and isinstance(frame, str): self.reply_frame = frame if self.fanout is None: self._answer(seq, now) @@ -765,11 +1159,7 @@ def _on_emit(self, args: list[Any], frame: str, now: int) -> None: and self.link.delta_key in delta ): self._linked.set() - if ( - not self._hydrated.is_set() - and (root := delta.get(ROOT_STATE)) - and root.get(HYDRATED_VAR) is True - ): + if not self._hydrated.is_set() and is_hydrated(delta): self._hydrated.set() elif name == "new_token": # The token was in use by another session; reflex hands out a new one. @@ -922,13 +1312,13 @@ async def _drain(self, until: int) -> None: await asyncio.wait_for(self._idle.wait(), remaining) async def close(self) -> None: - """Leave the namespace and close the websocket.""" + """Leave the session and close the websocket.""" self._closing = True - ws = self.ws + ws, codec = self.ws, self.codec if ws is not None: + assert codec is not None if self.error is None: - with contextlib.suppress(ConnectionClosed): - await ws.send(DISCONNECT_FRAME) + await leave(ws, codec) await ws.close() if self._reader is not None: self._reader.cancel() @@ -959,6 +1349,7 @@ def __init__( shape: The event they send, or ``None`` for idle sessions. link: The event they send once hydrated, or ``None``. """ + self.endpoint = endpoint self.sessions = [_Session(index, endpoint, shape, link) for index in indices] @property @@ -971,12 +1362,18 @@ def errors(self) -> list[str]: return [f"session {s.index}: {s.error}" for s in self.sessions if s.error] async def prime(self) -> None: - """Connect, hydrate and link every session, up to ``PRIME_CONCURRENCY`` at a time.""" + """Find the protocol, then connect, hydrate and link every session, up to ``PRIME_CONCURRENCY`` at a time.""" + try: + codec = await resolve_codec(self.endpoint) + except Exception as exc: + for session in self.sessions: + session._fail(f"could not connect: {type(exc).__name__}: {exc}") + return gate = asyncio.Semaphore(PRIME_CONCURRENCY) - await asyncio.gather(*(session.prime(gate) for session in self.sessions)) + await asyncio.gather(*(session.prime(gate, codec) for session in self.sessions)) async def close(self) -> None: - """Leave the namespace and close every websocket.""" + """Leave the sessions and close every websocket.""" await asyncio.gather(*(session.close() for session in self.sessions)) @@ -1350,6 +1747,120 @@ async def wait_readable(conn: Connection) -> None: loop.remove_reader(fd) +def spawn_process( + target: Callable[..., None], *args: Any, name: str +) -> tuple[BaseProcess, Connection]: + """Start a daemon process with a pipe to it. + + It is spawned, not forked: the harness runs hooks on threads. + + Args: + target: The process's function; the child's end of the pipe is its + last argument. + *args: Its other arguments. + name: The process's name. + + Returns: + The started process and the parent's end of the pipe. + """ + context = multiprocessing.get_context("spawn") + conn, child = context.Pipe() + proc = context.Process(target=target, args=(*args, child), name=name, daemon=True) + proc.start() + # Only the process holds the child's end now, so its exit closes the pipe. + child.close() + return proc, conn + + +def stop_processes(procs: Sequence[BaseProcess]) -> None: + """Terminate started processes together, killing each one that does not exit in time. + + Args: + procs: The processes. + """ + for proc in procs: + if proc.is_alive(): + proc.terminate() + for proc in procs: + proc.join(KILL_GRACE_S) + if proc.is_alive(): + proc.kill() + proc.join(KILL_GRACE_S) + + +def _ready_timeout(sessions: int) -> float: + """Bound the wait for a process to prime its sessions. + + Args: + sessions: The process's sessions. + + Returns: + Seconds: a margin plus the prime timeout per batch of + ``PRIME_CONCURRENCY`` sessions. + """ + return READY_MARGIN_S + -(-sessions // PRIME_CONCURRENCY) * PRIME_TIMEOUT_S + + +def _check_started(errors: Sequence[str], sessions: int) -> None: + """Fail when sessions could not start. + + Args: + errors: Why sessions failed to prime. + sessions: How many were to start. + + Raises: + LoadError: With how many failed and the first reason. + """ + if errors: + msg = f"{len(errors)} of {sessions} sessions could not start; {errors[0]}" + raise LoadError(msg) + + +def _report(conn: Connection, proc: BaseProcess, who: str) -> Any: + """Receive a process's next report. + + Args: + conn: The pipe to it, readable or closed. + proc: The process. + who: The process as messages name it, e.g. ``generator process 0``. + + Returns: + The report's value. + + Raises: + LoadError: When the process exited or reports a failure. + """ + try: + kind, value = conn.recv() + except (EOFError, OSError): + msg = f"{who} exited with code {proc.exitcode}" + raise LoadError(msg) from None + if kind == "error": + msg = f"{who} failed:\n{value}" + raise LoadError(msg) + return value + + +def _serve( + conn: Connection, sessions: int, main: Callable[[], Awaitable[None]] +) -> None: + """Run a spawned process's work on a fresh loop, reporting a failure to the parent. + + Args: + conn: The pipe to the parent. + sessions: The sessions the work opens. + main: Starts the work, which reports its results itself. + """ + # The parent decides when to stop, also on Ctrl-C. + signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + raise_fd_limit(sessions) + event_loop().run_until_complete(main()) + except BaseException: + with contextlib.suppress(OSError): + conn.send(("error", traceback.format_exc())) + + async def _handshake(conn: Connection, errors: list[str]) -> int | None: """Report the primed sessions to the parent and wait for the start. @@ -1382,28 +1893,29 @@ def abort() -> None: return value +async def _load(plan: LoadPlan, indices: Sequence[int], conn: Connection) -> None: + """Run a generator process's sessions on its CPUs and send the parent what they measured. + + Args: + plan: The load plan. + indices: The sessions of this process. + conn: The pipe to the parent. + """ + if plan.cpus is not None and hasattr(os, "sched_setaffinity"): + os.sched_setaffinity(0, plan.cpus) + shard = await _run_sessions(plan, indices, lambda errors: _handshake(conn, errors)) + conn.send(("result", shard)) + + def _worker(plan: LoadPlan, indices: Sequence[int], conn: Connection) -> None: - """Run one generator process: its sessions, on its CPUs, reporting over a pipe. + """Run one generator process, reporting over a pipe. Args: plan: The load plan. indices: The sessions of this process. conn: The pipe to the parent. """ - # The parent decides when to stop, also on Ctrl-C. - signal.signal(signal.SIGINT, signal.SIG_IGN) - try: - if plan.cpus is not None and hasattr(os, "sched_setaffinity"): - os.sched_setaffinity(0, plan.cpus) - raise_fd_limit(len(indices)) - loop = event_loop() - shard = loop.run_until_complete( - _run_sessions(plan, indices, lambda errors: _handshake(conn, errors)) - ) - conn.send(("result", shard)) - except BaseException: - with contextlib.suppress(OSError): - conn.send(("error", traceback.format_exc())) + _serve(conn, len(indices), lambda: _load(plan, indices, conn)) class LoadRunner: @@ -1429,8 +1941,6 @@ def __init__(self, plan: LoadPlan) -> None: def run(self, on_window: Callable[[str], None] | None = None) -> LoadResult: """Start the processes, prime every session, run the schedule and merge the results. - Processes are spawned, not forked: the harness runs hooks on threads. - Args: on_window: Called with ``"start"`` and ``"end"`` on this thread when the measured window starts and ends, e.g. to read the server's @@ -1444,30 +1954,25 @@ def run(self, on_window: Callable[[str], None] | None = None) -> LoadResult: out, or the runner was stopped. """ plan = self.plan - context = multiprocessing.get_context("spawn") try: with self._lock: if self._stopped: msg = "the load was stopped before it started" raise LoadError(msg) for index in range(plan.processes): - conn, child = context.Pipe() - self._conns.append(conn) - proc = context.Process( - target=_worker, - args=(plan, range(index, plan.sessions, plan.processes), child), + proc, conn = spawn_process( + _worker, + plan, + range(index, plan.sessions, plan.processes), name=f"reflex-bench-load-{index}", - daemon=True, ) - proc.start() - child.close() self._procs.append(proc) + self._conns.append(conn) per_process = -(-plan.sessions // plan.processes) - batches = -(-per_process // PRIME_CONCURRENCY) - ready = self._collect("ready", READY_MARGIN_S + batches * PRIME_TIMEOUT_S) - if errors := [error for errors in ready for error in errors]: - msg = f"{len(errors)} of {plan.sessions} sessions could not start; {errors[0]}" - raise LoadError(msg) + ready = self._collect("ready", _ready_timeout(per_process)) + _check_started( + [error for errors in ready for error in errors], plan.sessions + ) t0 = time.perf_counter_ns() + int(START_DELAY_S * 1e9) for conn in self._conns: conn.send(("go", t0)) @@ -1539,18 +2044,9 @@ def _collect(self, kind: str, timeout: float) -> list[Any]: msg = f"generator processes sent no {kind} within {timeout:g} s" raise LoadError(msg) for index, conn in pending: - if conn not in ready and self._procs[index].sentinel not in ready: - continue - try: - got, value = conn.recv() - except (EOFError, OSError): - code = self._procs[index].exitcode - msg = f"generator process {index} exited with code {code}" - raise LoadError(msg) from None - if got == "error": - msg = f"generator process {index} failed:\n{value}" - raise LoadError(msg) - values[index] = value + proc = self._procs[index] + if conn in ready or proc.sentinel in ready: + values[index] = _report(conn, proc, f"generator process {index}") return [values[index] for index in range(len(self._conns))] def stop(self) -> None: @@ -1560,14 +2056,7 @@ def stop(self) -> None: """ with self._lock: self._stopped = True - for proc in self._procs: - if proc.is_alive(): - proc.terminate() - for proc in self._procs: - proc.join(KILL_GRACE_S) - if proc.is_alive(): - proc.kill() - proc.join(KILL_GRACE_S) + stop_processes(self._procs) def run_load( @@ -1609,22 +2098,15 @@ def _hold_worker(endpoint: Endpoint, count: int, conn: Connection) -> None: count: How many. conn: The pipe to the parent. """ - # The parent decides when to stop, also on Ctrl-C. - signal.signal(signal.SIGINT, signal.SIG_IGN) - try: - raise_fd_limit(count) - event_loop().run_until_complete(_hold(endpoint, count, conn)) - except BaseException: - with contextlib.suppress(OSError): - conn.send(("error", traceback.format_exc())) + _serve(conn, count, lambda: _hold(endpoint, count, conn)) class SessionHold: """Idle sessions in a spawned process: connected, hydrated and answering pings. The sessions send no events after priming, like browser tabs left open. - :meth:`ready` waits until every session is primed, :meth:`close` leaves the - namespace and closes each websocket, and :meth:`kill` ends the process from + :meth:`ready` waits until every session is primed, :meth:`close` leaves + each session and closes its websocket, and :meth:`kill` ends the process from any thread (a benchmark's ``conclude`` after a timed-out ``sample``). A hold whose parent exits closes its sessions and ends. """ @@ -1652,21 +2134,13 @@ def start(self) -> SessionHold: Raises: LoadError: When the hold was killed before it started. """ - context = multiprocessing.get_context("spawn") with self._lock: if self._stopped: msg = "the hold was stopped before it started" raise LoadError(msg) - conn, child = context.Pipe() - self._conn = conn - proc = self._proc = context.Process( - target=_hold_worker, - args=(self.endpoint, self.count, child), - name="reflex-bench-hold", - daemon=True, + self._proc, self._conn = spawn_process( + _hold_worker, self.endpoint, self.count, name="reflex-bench-hold" ) - proc.start() - child.close() return self def ready(self) -> None: @@ -1676,14 +2150,11 @@ def ready(self) -> None: LoadError: When sessions could not start, with how many, or the process failed, timed out or was killed. """ - batches = -(-self.count // PRIME_CONCURRENCY) - errors = self._receive("ready", READY_MARGIN_S + batches * PRIME_TIMEOUT_S) - if errors: - msg = f"{len(errors)} of {self.count} sessions could not start; {errors[0]}" - raise LoadError(msg) + errors = self._receive("ready", _ready_timeout(self.count)) + _check_started(errors, self.count) def close(self) -> list[str]: - """Leave the namespace and close every websocket, then end the process. + """Leave every session and close its websocket, then end the process. Returns: Why sessions failed while they were held, e.g. because the server @@ -1728,15 +2199,7 @@ def _receive(self, kind: str, timeout: float) -> list[str]: if not ready: msg = f"the hold sent no {kind} report within {timeout:g} s" raise LoadError(msg) - try: - got, value = conn.recv() - except (EOFError, OSError): - msg = f"the hold's process exited with code {proc.exitcode}" - raise LoadError(msg) from None - if got == "error": - msg = f"the hold's process failed:\n{value}" - raise LoadError(msg) - return value + return _report(conn, proc, "the hold's process") def kill(self) -> None: """Terminate the process, killing it if it does not exit. @@ -1746,14 +2209,8 @@ def kill(self) -> None: with self._lock: self._stopped = True proc = self._proc - if proc is None: - return - if proc.is_alive(): - proc.terminate() - proc.join(KILL_GRACE_S) - if proc.is_alive(): - proc.kill() - proc.join(KILL_GRACE_S) + if proc is not None: + stop_processes([proc]) def hold_sessions(endpoint: Endpoint, count: int) -> SessionHold: diff --git a/packages/reflex-bench/src/reflex_bench/suites/wire.py b/packages/reflex-bench/src/reflex_bench/suites/wire.py index 15d7d15c98b..077edb0fe74 100644 --- a/packages/reflex-bench/src/reflex_bench/suites/wire.py +++ b/packages/reflex-bench/src/reflex_bench/suites/wire.py @@ -6,7 +6,10 @@ payload bytes of every websocket frame it sends and receives: - ``wire.hydrate``: from the connect until the delta that sets ``is_hydrated``, - the page load of the playground's index route. + the page load of the playground's index route. On the plain protocol the + connect frame carries ``hydrate_and_load`` with an empty payload, as the + frontend's does on a reconnect: the full state, not the difference from the + compiled defaults that a browser's first load asks for. - ``wire.event[shape=...]``: after that, one event of a playground shape (:data:`~reflex_bench.suites.events.SHAPES`): the request frame, and the reply until the delta that echoes its sequence number, every frame in @@ -25,7 +28,10 @@ as the control that changes no collection. Every handler also sets ``last_seq``, so the deltas differ only by the collection they carry. -A byte count is a text frame's UTF-8 length; the HTTP handshake that opens the +The session speaks whichever protocol the backend does +(:func:`~reflex_bench.drivers.events.detect_protocol`), so a comparison across +reflex's switch from Socket.IO to the plain protocol includes the framing. A +byte count is a text frame's UTF-8 length; the HTTP handshake that opens the websocket is not a frame. Sizes are deterministic for a given app and reflex version, so every metric is exact and a sample is one run. The playground is the fixture of ``hydrate`` and ``event``; the generated app's source hash keys @@ -35,7 +41,6 @@ from __future__ import annotations import asyncio -import contextlib import dataclasses import hashlib import json @@ -47,30 +52,26 @@ from typing import Any, Protocol, TypeVar from websockets.asyncio.client import connect -from websockets.exceptions import ConnectionClosed from reflex_bench.context import Context from reflex_bench.drivers.app_process import AppProcess, run_cli from reflex_bench.drivers.events import ( - CLOSE, - CONNECT_ERROR_PREFIX, - CONNECT_FRAME, - DISCONNECT_FRAME, - EVENT_PREFIX, - HYDRATE_EVENT, - HYDRATED_VAR, + CODECS, + EVENT, + HANDSHAKE, MAX_FRAME_BYTES, ON_LOAD_EVENT, - OPEN_PREFIX, + OPENED, PING, - PONG, - ROOT_STATE, + Codec, Endpoint, EventShape, ProtocolError, - event_frame, + end_reason, event_loop, - event_url, + is_hydrated, + leave, + resolve_codec, seq_payload, ) from reflex_bench.fixtures import describe_fixture @@ -143,10 +144,12 @@ class Hydration: """What a page load moved, from the connect until ``is_hydrated``. Attributes: - sent_bytes: Frame bytes sent: the namespace join and the hydration - events. - received_bytes: Frame bytes received: the engine.io open, the namespace - ack, a ``new_token`` if any, and the deltas. + sent_bytes: Frame bytes sent: the connect frame, which carries the + hydration as its boot event, or on Socket.IO the namespace join + and the hydration events. + received_bytes: Frame bytes received: the handshake (on Socket.IO the + engine.io open and the namespace ack), a ``new_token`` if any, and + the deltas. sent_frames: The frames sent. received_frames: The frames received. largest_frame_bytes: The largest frame received. @@ -189,15 +192,23 @@ class Exchange: class WireSession: """One browser tab on a metered websocket: it hydrates, then sends events one at a time.""" - def __init__(self, ws: Transport, token: str, pathname: str = "/") -> None: + def __init__( + self, + ws: Transport, + token: str, + pathname: str = "/", + codec: Codec = CODECS["websocket"], + ) -> None: """Wrap a connected websocket; nothing is sent yet. Args: ws: The websocket, connected with ``token`` in its URL. token: The session's token. pathname: The page route the session loads. + codec: The protocol the websocket speaks. """ self._ws = ws + self.codec = codec self.token = token self.pathname = pathname self.query: dict[str, str] = {} @@ -221,61 +232,55 @@ async def _send(self, frame: str) -> int: await self._ws.send(frame) return size - async def _receive(self) -> tuple[str | bytes, int, list[Any] | None]: + async def _receive(self) -> tuple[str | bytes, int, list[Any]]: """Receive a frame and count it; pings are answered, disconnects raise. A keepalive ping and its pong are timing, not payload: neither counts. Returns: - The frame, its bytes and, for a socket.io event, its arguments. + The frame, its bytes and the message it decodes to. Raises: ProtocolError: When the server disconnects the session, refuses - the namespace, asks the page to reload or sends a frame the - session does not speak. + it, asks the page to reload or sends a frame the session does + not speak. """ + codec = self.codec while True: message = await self._ws.recv() - if message == PING: - await self._ws.send(PONG) + args = codec.parse(message) + if args is not None and args[0] == PING: + await self._ws.send(codec.pong) continue size = frame_bytes(message) self.received_bytes += size self.received_frames += 1 - if isinstance(message, str): - if message.startswith(EVENT_PREFIX): - args = json.loads(message[len(EVENT_PREFIX) :]) - if args[0] == "new_token": - # The token was in use; reflex hands out a new one. - self.token = args[1] - elif args[0] == "reload": - msg = "the server asked the page to reload" - raise ProtocolError(msg) - return message, size, args - if message in {DISCONNECT_FRAME, CLOSE}: - msg = f"the server disconnected the session ({message!r})" - raise ProtocolError(msg) - if message.startswith(CONNECT_ERROR_PREFIX): - reason = message[len(CONNECT_ERROR_PREFIX) :] - msg = f"the server refused the namespace: {reason}" - raise ProtocolError(msg) - if message.startswith((OPEN_PREFIX, CONNECT_FRAME)): - return message, size, None - msg = f"unexpected frame: {message[:40]!r}" - raise ProtocolError(msg) + if args is None: + msg = f"unexpected frame: {message[:40]!r}" + raise ProtocolError(msg) + name = args[0] + if name == "new_token": + # The token was in use; reflex hands out a new one. + self.token = args[1] + elif name == "reload": + msg = "the server asked the page to reload" + raise ProtocolError(msg) + elif (reason := end_reason(args, message)) is not None: + raise ProtocolError(reason) + return message, size, args @staticmethod - def _delta(args: list[Any] | None, sizes: dict[str, int]) -> dict[str, Any]: - """Take the delta of an event frame, adding its substates' bytes. + def _delta(args: list[Any], sizes: dict[str, int]) -> dict[str, Any]: + """Take the delta of an event message, adding its substates' bytes. Args: - args: The frame's arguments, or ``None`` for a control frame. + args: The message. sizes: Bytes per substate, added to. Returns: - The delta, empty when the frame carries none. + The delta, empty when the message carries none. """ - if args is None or args[0] != "event": + if args[0] != EVENT: return {} delta = args[1].get("delta") or {} for substate, values in delta.items(): @@ -285,32 +290,39 @@ def _delta(args: list[Any] | None, sizes: dict[str, int]) -> dict[str, Any]: return delta async def hydrate(self) -> Hydration: - """Join the namespace and hydrate like a page load. + """Open the session and hydrate it like a page load. Returns: What the page load moved. Raises: - ProtocolError: When the server does not open an engine.io session. + ProtocolError: When a Socket.IO server does not open an engine.io + session. """ sent0, sent_frames0 = self.sent_bytes, self.sent_frames received0, received_frames0 = self.received_bytes, self.received_frames - opened, largest, _ = await self._receive() - if not (isinstance(opened, str) and opened.startswith(OPEN_PREFIX)): - msg = f"the server did not open an engine.io session: {opened[:40]!r}" - raise ProtocolError(msg) - await self._send(CONNECT_FRAME) + codec = self.codec + largest = 0 + if codec.greets: + greeting, largest, args = await self._receive() + if args[0] != OPENED: + msg = f"the server did not open an engine.io session: {greeting[:40]!r}" + raise ProtocolError(msg) + await self._send( + codec.connect_frame( + token=self.token, pathname=self.pathname, query=self.query + ) + ) sizes: dict[str, int] = {} while True: - message, size, args = await self._receive() + _, size, args = await self._receive() largest = max(largest, size) - if isinstance(message, str) and message.startswith(CONNECT_FRAME): - # Joined; a new_token, when the server sends one, came before. - for name in (HYDRATE_EVENT, ON_LOAD_EVENT): + if args[0] == HANDSHAKE: + # Opened; a new_token, when the server sends one, came before. + for name in codec.load_events: await self._send(self._event(name, {})) continue - delta = self._delta(args, sizes) - if delta.get(ROOT_STATE, {}).get(HYDRATED_VAR) is True: + if is_hydrated(self._delta(args, sizes)): break return Hydration( sent_bytes=self.sent_bytes - sent0, @@ -331,7 +343,7 @@ def _event(self, name: str, payload: dict[str, Any]) -> str: Returns: The frame. """ - return event_frame( + return self.codec.event_frame( name, payload, token=self.token, pathname=self.pathname, query=self.query ) @@ -371,10 +383,7 @@ async def navigate(self, pathname: str, query: dict[str, str]) -> Exchange: self.pathname = pathname self.query = query request = await self._send(self._event(ON_LOAD_EVENT, {})) - return await self._reply( - request, - lambda delta: delta.get(ROOT_STATE, {}).get(HYDRATED_VAR) is True, - ) + return await self._reply(request, is_hydrated) async def _reply( self, request: int, done: Callable[[dict[str, Any]], bool] @@ -409,7 +418,7 @@ async def _reply( def measure(endpoint: Endpoint, work: Callable[[WireSession], Awaitable[_T]]) -> _T: - """Connect one session, run work on it, and leave the namespace. + """Connect one session, run work on it, and leave the session. Args: endpoint: The backend and the page route. @@ -424,8 +433,9 @@ def measure(endpoint: Endpoint, work: Callable[[WireSession], Awaitable[_T]]) -> async def run() -> _T: token = str(uuid.uuid4()) + codec = await resolve_codec(endpoint) async with connect( - event_url(endpoint.backend_url, token), + codec.url(endpoint.backend_url, token), # A browser sends the page's origin; the router's page host and # full paths in the deltas come from it. additional_headers={"Origin": endpoint.backend_url}, @@ -435,15 +445,14 @@ async def run() -> _T: ping_interval=None, close_timeout=1, ) as ws: - session = WireSession(ws, token, endpoint.pathname) + session = WireSession(ws, token, endpoint.pathname, codec) try: return await asyncio.wait_for(work(session), WIRE_TIMEOUT_S) except TimeoutError: msg = f"the session did not finish within {WIRE_TIMEOUT_S:g} s" raise TimeoutError(msg) from None finally: - with contextlib.suppress(ConnectionClosed): - await ws.send(DISCONNECT_FRAME) + await leave(ws, codec) loop = event_loop() try: diff --git a/pyi_hashes.json b/pyi_hashes.json index 9df583b7fa1..3d732f5e2ba 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -118,7 +118,7 @@ "packages/reflex-components-recharts/src/reflex_components_recharts/polar.pyi": "99ebcfc07868061bdc3c2010d85a153f", "packages/reflex-components-recharts/src/reflex_components_recharts/recharts.pyi": "f2a9cf6db58d169289fbc76de6813300", "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", - "reflex/__init__.pyi": "23c090b0294d3a91e59f447f36a7a37f", + "reflex/__init__.pyi": "e9b4f4bfded698840bc7f12ae5b6ec75", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", "reflex/experimental/memo.pyi": "87e485038dc79f66b2ab15144845ef1e" } diff --git a/pyproject.toml b/pyproject.toml index 0934af7a6e8..d6fb5d0d138 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,12 +21,11 @@ keywords = ["web", "framework"] requires-python = ">=3.11,<4.0" dependencies = [ "click >=8.2", - "granian[reload] >=2.7.4", + "granian[reload] >=2.8.1", "httpx >=0.26,<1.0", "packaging >=24.2,<27", "psutil >=7.0.0,<8.0; sys_platform == 'win32'", "python-multipart >=0.0.32,<1.0", - "python-socketio >=5.12.0,<6.0", "redis >=6.4,<8.0", "rich >=13,<16", "starlette >=1.3.1", @@ -71,10 +70,12 @@ db = [ "sqlmodel >=0.0.24", ] pydantic = ["reflex-base[pydantic]"] +socketio = ["python-socketio >=5.12.0,<6.0"] +# The uvicorn backend: server, websocket support, and the posix prod runner. +uvicorn = ["uvicorn >=0.20.0", "websockets >=13.0", "gunicorn >=23.0"] testing = [ "psutil >=7.0.0,<8.0", "selenium >=4.0.0,<5.0", - "uvicorn >=0.34.0,<1.0", ] [project.urls] @@ -121,6 +122,7 @@ dev = [ "pytest-split", "pytest", "python-dotenv", + "python-socketio", "pyyaml", "reflex-bench", "reflex-docgen", diff --git a/reflex/__init__.py b/reflex/__init__.py index ea1775c275f..4504772b4ea 100644 --- a/reflex/__init__.py +++ b/reflex/__init__.py @@ -229,6 +229,7 @@ _SUBMODULES: set[str] = { "components", "app", + "channels", "style", "admin", "base", diff --git a/reflex/app.py b/reflex/app.py index 97d7d7472e5..32578d246f1 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -17,20 +17,16 @@ import tempfile import time import traceback -import urllib.parse from collections.abc import ( AsyncIterator, Callable, Collection, Coroutine, - Iterable, Mapping, Sequence, ) from contextvars import Token -from datetime import timedelta from pathlib import Path -from types import SimpleNamespace from typing import TYPE_CHECKING, Any, overload from reflex_base import constants, otel @@ -38,23 +34,15 @@ from reflex_base.config import get_config, reload_config from reflex_base.context.base import BaseContext from reflex_base.environment import auto_reload_cooldown, environment -from reflex_base.event import ( - _EVENT_FIELDS, - Event, - EventSpec, - EventType, - IndividualEventType, - noop, -) +from reflex_base.event import Event, EventSpec, EventType, IndividualEventType, noop from reflex_base.event.context import EventContext from reflex_base.event.processor import BaseStateEventProcessor, EventProcessor from reflex_base.registry import RegistrationContext from reflex_base.telemetry_context import CompileTrigger, TelemetryContext -from reflex_base.utils import memo_paths +from reflex_base.utils import console, memo_paths from reflex_base.utils.imports import ImportVar from reflex_base.utils.types import ( ASGIApp, - Message, Receive, Scope, Send, @@ -71,12 +59,11 @@ from reflex_components_core.core.breakpoints import set_breakpoints from reflex_components_core.core.sticky import sticky from reflex_components_sonner.toast import toast -from socketio import ASGIApp as EngineIOApp -from socketio import AsyncNamespace, AsyncServer from starlette.applications import Starlette from starlette.middleware import cors from starlette.requests import Request from starlette.responses import JSONResponse, Response +from starlette.routing import WebSocketRoute from starlette.staticfiles import StaticFiles from typing_extensions import Unpack @@ -84,9 +71,10 @@ from reflex._upload import UploadFile as UploadFile from reflex.admin import AdminDash from reflex.app_mixins import AppMixin, LifespanMixin, MiddlewareMixin +from reflex.channels import Channel from reflex.compiler import compiler from reflex.compiler.compiler import readable_name_from_component -from reflex.istate.data import SessionData +from reflex.event_namespace import BaseEventNamespace, WebsocketEventNamespace from reflex.istate.manager import StateManager, StateModificationContext from reflex.istate.manager.token import BaseStateToken from reflex.route import ( @@ -111,7 +99,6 @@ should_prerender_routes, ) from reflex.utils.misc import is_page_meta_set, run_in_thread -from reflex.utils.token_manager import RedisTokenManager, TokenManager logger = logging.getLogger(__name__) @@ -124,11 +111,16 @@ from reflex_base.plugins import Plugin from reflex_base.plugins.base import AddPageProtocol from reflex_base.vars import Var + from socketio import AsyncServer # Define custom types. ComponentCallable = Callable[[], Component | tuple[Component, ...] | str | Var] + _SocketIOServer = AsyncServer else: ComponentCallable = Callable[[], Component | tuple[Component, ...] | str] + # Lets annotations resolve without the optional python-socketio package; + # the public name resolves lazily, through __getattr__. + _SocketIOServer = Any Reducer = Callable[[Event], Coroutine[Any, Any, StateUpdate]] @@ -457,7 +449,7 @@ class App(MiddlewareMixin, LifespanMixin): head_components: list[Component] = dataclasses.field(default_factory=list) - sio: AsyncServer | None = None + sio: _SocketIOServer | None = None html_lang: str | None = None @@ -503,7 +495,10 @@ class App(MiddlewareMixin, LifespanMixin): admin_dash: AdminDash | None = None # The async server name space. - _event_namespace: EventNamespace | None = None + _event_namespace: BaseEventNamespace | None = None + + # Side channels multiplexed onto the event websocket, by channel name. + _channels: dict[str, Channel] = dataclasses.field(default_factory=dict) # The processor queue for handling events. _event_processor: EventProcessor | None = None @@ -533,7 +528,7 @@ class App(MiddlewareMixin, LifespanMixin): ) = None @property - def event_namespace(self) -> EventNamespace | None: + def event_namespace(self) -> BaseEventNamespace | None: """The event namespace. Returns: @@ -541,6 +536,43 @@ def event_namespace(self) -> EventNamespace | None: """ return self._event_namespace + def register_channel(self, channel: Channel) -> None: + """Register a side channel multiplexed onto the event websocket. + + Channels run only over the plain WebSocket transport with state + enabled; elsewhere the channel is skipped with a warning. + + Args: + channel: The channel to serve. + + Raises: + RuntimeError: If the name is taken. + """ + name = type(channel).name + # Packages register their channels from a plugin, unaware of how the + # app is set up; where channels cannot run, the browser-side handle + # reports it, so the app keeps serving rather than failing to start. + if self.event_namespace is None: + # A supplied `_state` is not enough: without enable_state the app + # never sets up a transport, so the channel would be unreachable. + logger.warning( + f"Channel {name!r} is not served: it needs the event websocket, " + "which exists only when state is enabled (rx.App(enable_state=True), " + "the default)." + ) + return + if not self.event_namespace.supports_channels: + logger.warning( + f"Channel {name!r} is not served: it requires the plain WebSocket " + "transport; remove the transport setting in rxconfig.py or set " + 'transport="websocket".' + ) + return + if name in self._channels: + msg = f"A channel named {name!r} is already registered." + raise RuntimeError(msg) + self._channels[name] = channel + @property def event_processor(self) -> EventProcessor: """The event processor. @@ -607,7 +639,8 @@ def _setup_state(self) -> None: """Set up the state for the app. Raises: - RuntimeError: If the socket server is invalid. + RuntimeError: If the socket server is invalid, or the Socket.IO + transport is requested without python-socketio installed. """ if not self._state: return @@ -621,80 +654,73 @@ def _setup_state(self) -> None: # rather than on the first frontend error that consults it. auto_reload_cooldown() - # Set up the Socket.IO AsyncServer. - if not self.sio: - self.sio = AsyncServer( - async_mode="asgi", - cors_allowed_origins=( - ( - "*" - if config.cors_allowed_origins == ("*",) - else list(config.cors_allowed_origins) - ) - if config.transport == "websocket" - else [] - ), - cors_credentials=config.transport == "websocket", - max_http_buffer_size=environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE.get(), - ping_interval=environment.REFLEX_SOCKET_INTERVAL.get().total_seconds(), - ping_timeout=environment.REFLEX_SOCKET_TIMEOUT.get().total_seconds(), - json=SimpleNamespace( - dumps=staticmethod(_sio_dumps), - loads=staticmethod(_sio_loads), - ), - allow_upgrades=False, - transports=[config.transport], - # Handlers here only parse and enqueue (or emit a pong), so run - # them inline on the socket's receive loop instead of paying a - # task creation and a loop hop per incoming message. - async_handlers=False, - ) - elif getattr(self.sio, "async_mode", "") != "asgi": - msg = f"Custom `sio` must use `async_mode='asgi'`, not '{self.sio.async_mode}'." - raise RuntimeError(msg) - - # Create the socket app. Note event endpoint constant replaces the default 'socket.io' path. - socket_app = EngineIOApp(self.sio, socketio_path="") namespace = config.get_event_namespace() + event_path = config.prepend_backend_path(str(constants.Endpoint.EVENT)) - # Create the event namespace and attach the main app. Not related to any paths. - self._event_namespace = EventNamespace(namespace, self) - - # Register the event namespace with the socket. - self.sio.register_namespace(self.event_namespace) - # Mount the socket app with the API. - if self._api: - - class HeaderMiddleware: - def __init__(self, app: ASGIApp): - self.app = app + if self.sio is not None and config.transport == "websocket": + if "transport" in config._non_default_attributes: + msg = ( + "A custom `sio` server requires the Socket.IO transport; " + 'set transport="socketio" (or "polling") in rxconfig.py.' + ) + raise RuntimeError(msg) + # The default transport used to be Socket.IO over a websocket, so + # an app passing its own server never had to choose one. + console.deprecate( + feature_name="A custom `sio` server without a transport setting", + reason='set transport="socketio" in rxconfig.py; the default ' + "transport is now a plain WebSocket", + deprecation_version="0.10.0", + removal_version="1.0", + ) + # Persisted, so the compiled frontend and the workers use it too. + config._set_persistent(transport="socketio") - async def __call__(self, scope: Scope, receive: Receive, send: Send): - original_send = send + if config.transport in ("socketio", "polling"): + # Legacy Socket.IO transport, kept behind the optional dependency. + try: + from reflex import socketio_namespace - async def modified_send(message: Message): - if message["type"] == "websocket.accept": - if scope.get("subprotocols"): - # The following *does* say "subprotocol" instead of "subprotocols", intentionally. - message["subprotocol"] = scope["subprotocols"][0] + # Through the module, so an app that replaced the class gets + # its own. + event_namespace_class = sys.modules[__name__].EventNamespace + except ImportError as ex: + msg = f"transport={config.transport!r} requires {_SOCKETIO_PACKAGE}." + raise RuntimeError(msg) from ex - headers = dict(message.get("headers", [])) - header_key = b"sec-websocket-protocol" - if subprotocol := headers.get(header_key): - message["headers"] = [ - *message.get("headers", []), - (header_key, subprotocol), - ] + if ( + event_namespace_class.on_event + is not socketio_namespace.EventNamespace.on_event + ): + console.deprecate( + feature_name="Overriding EventNamespace.on_event", + reason="Override handle_event instead; it runs for every " + "event, including the hydrate on connect, on both transports", + deprecation_version="0.10.1", + removal_version="1.0", + ) - return await original_send(message) + socket_app = socketio_namespace.create_socketio_app(self, config) - return await self.app(scope, receive, modified_send) + # Create the event namespace and attach the main app. Not related to any paths. + self._event_namespace = event_namespace_class(namespace, self) - socket_app_with_headers = HeaderMiddleware(socket_app) - self._api.mount( - config.prepend_backend_path(str(constants.Endpoint.EVENT)), - socket_app_with_headers, - ) + # Register the event namespace with the socket. + self.sio.register_namespace(self._event_namespace) # pyright: ignore[reportOptionalMemberAccess] + # Mount the socket app with the API. + if self._api: + self._api.mount(event_path, socket_app) + else: + # Default transport: plain WebSocket served by the API itself. Read + # from the module here, so an app that replaced the class gets its own. + self._event_namespace = WebsocketEventNamespace(namespace, self) + if self._api: + # Also with a trailing slash, the form the client dials: proxy + # rules written for Socket.IO route "/_event/*". + for path in (event_path, event_path.rstrip("/") + "/"): + self._api.router.routes.append( + WebSocketRoute(path, self._event_namespace.handle_websocket) + ) # Check the exception handlers self._validate_exception_handlers() @@ -2051,530 +2077,43 @@ async def health(_request: Request) -> JSONResponse: return JSONResponse(content=health_status, status_code=status_code) -def _utf8_size(data: str) -> int: - """Size of a serialized message in UTF-8 bytes. - - ASCII payloads (the common case) are sized without encoding a copy. - - Args: - data: The serialized message. - - Returns: - The number of bytes the message occupies on the wire. - """ - return len(data) if data.isascii() else len(data.encode()) - - -def _sio_dumps(obj: Any, **kwargs: Any) -> str: - """Serialize an outgoing Socket.IO packet, recording its size when telemetry is on. - - Args: - obj: The packet payload. - **kwargs: Options forwarded to the JSON encoder. +# What the Socket.IO transport needs installed. +_SOCKETIO_PACKAGE = ( + "the optional python-socketio package; install it with: " + "pip install 'reflex[socketio]'" +) - Returns: - The JSON string. - """ - data = format.json_dumps(obj, **kwargs) - if otel.enabled: - otel.record_message_size(_utf8_size(data), "transmit") - return data +# Socket.IO names this module exported before python-socketio became optional: +# name -> (module, attribute). +_SOCKETIO_EXPORTS = { + "AsyncNamespace": ("socketio", "AsyncNamespace"), + "AsyncServer": ("socketio", "AsyncServer"), + "EngineIOApp": ("socketio", "ASGIApp"), + "EventNamespace": ("reflex.socketio_namespace", "EventNamespace"), +} -def _sio_loads(data: str | bytes, **kwargs: Any) -> Any: - """Deserialize an incoming Socket.IO packet, recording its size when telemetry is on. +def __getattr__(name: str) -> Any: + """Resolve the optional Socket.IO exports lazily. Args: - data: The JSON string. - **kwargs: Options forwarded to the JSON decoder. + name: The attribute name. Returns: - The decoded payload. - """ - if otel.enabled: - otel.record_message_size( - _utf8_size(data) if isinstance(data, str) else len(data), "receive" - ) - return json.loads(data, **kwargs) - + The resolved attribute. -def _decode_asgi_headers(headers: Iterable[tuple[bytes, bytes]]) -> dict[str, str]: - """Decode raw ASGI scope header pairs into a str-keyed dict. - - ASGI keeps a repeated header as separate pairs, so their values are combined in - order: cookies with "; " (RFC 9113 section 8.2.3), any other field with ", " - (RFC 9110 section 5.3). Header values that are not valid UTF-8 are decoded as - Latin-1. - - Args: - headers: Raw (name, value) byte pairs from the ASGI scope. - - Returns: - A dict mapping decoded header names to decoded values. + Raises: + AttributeError: If the attribute is unknown, or needs python-socketio + and it is not installed. """ - decoded: dict[str, str] = {} - for raw_name, raw_value in headers: - name = raw_name.decode() - try: - value = raw_value.decode() - except UnicodeDecodeError: - value = raw_value.decode("latin-1") - if name in decoded: - decoded[name] += ("; " if name == "cookie" else ", ") + value - else: - decoded[name] = value - return decoded - - -class EventNamespace(AsyncNamespace): - """The event namespace.""" - - # The application object. - app: App - - # Maximum error-level log entries a single session may produce via the - # client_error event before further reports from it are dropped. - _MAX_CLIENT_ERRORS_PER_SID = 5 - - # Process-wide bound on error-level client_error log entries per time - # window; per-SID budgets alone reset on reconnect, so scripted - # reconnects could otherwise flood the logs. - _CLIENT_ERROR_WINDOW_SECONDS = 60.0 - _MAX_CLIENT_ERRORS_PER_WINDOW = 20 - - def __init__(self, namespace: str, app: App): - """Initialize the event namespace. - - Args: - namespace: The namespace. - app: The application object. - """ - super().__init__(namespace) - self.app = app - - # Use TokenManager for distributed duplicate tab prevention - self._token_manager = TokenManager.create() - - # Number of client_error reports logged per SID, for rate limiting. - self._client_error_counts: dict[str, int] = {} - - # Connection-scoped router_data entries per SID, computed once at - # connect time instead of for every event on the connection. - self._static_router_data: dict[str, dict[str, Any]] = {} - - # Start time and count of the current process-wide client_error window. - self._client_error_window_start = 0.0 - self._client_error_window_count = 0 - - @property - def token_to_sid(self) -> Mapping[str, str]: - """Token to SID mapping for backward compatibility. - - Note: this mapping is read-only. - - Returns: - The token to SID mapping. - """ - # For backward compatibility, expose the underlying dict - return self._token_manager.token_to_sid - - @property - def sid_to_token(self) -> dict[str, str]: - """SID to token mapping for backward compatibility. - - Returns: - The SID to token mapping dict. - """ - # For backward compatibility, expose the underlying dict - return self._token_manager.sid_to_token - - async def on_connect(self, sid: str, environ: dict, auth: Any = None): - """Event for when the websocket is connected. - - Args: - sid: The Socket.IO session id. - environ: The request information, including HTTP headers. - auth: The payload of the socket.io CONNECT packet. The frontend - puts its hydrate event here so it is processed without waiting - for the connect acknowledgement round trip. - """ - if isinstance(self._token_manager, RedisTokenManager): - # Make sure this instance is watching for updates from other instances. - self._token_manager.ensure_lost_and_found_task(self.emit_update) - boot_event = frontend_schema = None - if isinstance(auth, dict): - boot_event = auth.get(constants.CompileVars.CONNECT_AUTH_EVENT) - frontend_schema = auth.get(constants.CompileVars.CONNECT_AUTH_STATE_SCHEMA) - # A frontend compiled for another state tree is told so before it - # hydrates; an older frontend sends no hash and is left alone. The - # backend hash is cached on the state class, so this is one lookup. - if ( - frontend_schema is not None - and self.app._state is not None - and frontend_schema - != (backend_schema := self.app._state._to_frontend_schema()) - ): - logger.warning( - "Frontend state schema " - f"{format.sanitize_client_log_value(frontend_schema, max_length=64)} " - f"for session {sid} does not match the backend state schema " - f"{backend_schema}: the frontend was built for another version of " - "the app and is told to reload." - ) - await self.emit( - str(constants.SocketEvent.STATE_MISMATCH), - { - "backend": backend_schema, - "reload_cooldown_ms": auto_reload_cooldown() - // timedelta(milliseconds=1), - }, - to=sid, - ) - query_params = urllib.parse.parse_qs(environ.get("QUERY_STRING", "")) - token_list = query_params.get("token", []) - if token_list: - # The boot event is the connection's first event: processing it - # records the new sid and token on the state under the state lock. - await self.link_token_to_sid( - sid, token_list[0], update_state=boot_event is None - ) - else: - logger.warning(f"No token provided in connection for session {sid}") - - subprotocol = environ.get("HTTP_SEC_WEBSOCKET_PROTOCOL") - if subprotocol and subprotocol != constants.Reflex.VERSION: - logger.warning( - f"Frontend version {subprotocol} for session {sid} does not match the backend version {constants.Reflex.VERSION}." - ) - if otel.enabled: - otel.record_connection(1) - - # Headers, client IP, and session id cannot change for the lifetime of - # the connection; compute them once instead of on every event. - self._static_router_data[sid] = self._build_static_router_data(sid, environ) - - if boot_event is not None: - try: - await self.on_event(sid, boot_event) - except Exception: - # A refused connect never reaches on_disconnect, so drop the - # token link and connection data made above before the error - # refuses the connect. - self._static_router_data.pop(sid, None) - if (linked_token := self.sid_to_token.get(sid)) is not None: - await self._token_manager.disconnect_token(linked_token, sid) - raise - - def _build_static_router_data(self, sid: str, environ: dict) -> dict[str, Any]: - """Build the connection-scoped router_data entries for a socket. - - Args: - sid: The Socket.IO session id. - environ: The request information, including HTTP headers. - - Returns: - The router_data entries that are constant for the connection. - """ - asgi_scope = environ.get("asgi.scope", {}) - - # Get the client headers. - headers = _decode_asgi_headers(asgi_scope.get("headers", [])) - - # Get the client IP - try: - client_ip: str = asgi_scope["client"][0] - headers["asgi-scope-client"] = client_ip - except (KeyError, IndexError): - client_ip = environ.get("REMOTE_ADDR", "0.0.0.0") - - # X-Forwarded-For is client-controlled unless the request arrived through - # a configured chain of trusted reverse proxies. - trusted_proxy_hops = get_config().trusted_proxy_hops - if trusted_proxy_hops: - forwarded_for = headers.get("x-forwarded-for", "") - forwarded_ips = [ - ip.strip() for ip in forwarded_for.split(",") if ip.strip() - ] - if len(forwarded_ips) >= trusted_proxy_hops: - client_ip = forwarded_ips[-trusted_proxy_hops] - return { - constants.RouteVar.SESSION_ID: sid, - constants.RouteVar.HEADERS: headers, - constants.RouteVar.CLIENT_IP: client_ip, - } - - def on_disconnect(self, sid: str) -> asyncio.Task | None: - """Event for when the websocket disconnects. - - Args: - sid: The Socket.IO session id. - - Returns: - An asyncio Task for cleaning up the token, or None. - """ - if otel.enabled: - otel.record_connection(-1) - self._client_error_counts.pop(sid, None) - self._static_router_data.pop(sid, None) - # Get token before cleaning up - disconnect_token = self.sid_to_token.get(sid) - if disconnect_token: - socket_record = self._token_manager.token_to_socket.get(disconnect_token) - if ( - socket_record is not None - and socket_record.sid == sid - and socket_record.instance_id == self._token_manager.instance_id - and (processor := self.app._event_processor) is not None - ): - processor._on_disconnect(disconnect_token) - # Use async cleanup through token manager - task = asyncio.create_task( - self._token_manager.disconnect_token(disconnect_token, sid), - name=f"reflex_disconnect_token|{disconnect_token}|{time.time()}", - ) - # Don't await to avoid blocking disconnect, but handle potential errors - task.add_done_callback( - lambda t: ( - t.exception() - and logger.error(f"Token cleanup error: {t.exception()}") - ) - ) - return task - return None - - async def emit_update(self, update: StateUpdate, token: str) -> None: - """Emit an update to the client. - - Args: - update: The state update to send. - token: The client token (tab) associated with the event. - """ - socket_record = self._token_manager.token_to_socket.get(token) - if ( - socket_record is None - or socket_record.instance_id != self._token_manager.instance_id - ): - if isinstance(self._token_manager, RedisTokenManager): - # The socket belongs to another instance of the app, send it to the lost and found. - await self._token_manager.emit_lost_and_found(token, update) - else: - # If the socket record is None, we are not connected to a client. Prevent sending - # updates to all clients. - logger.warning( - f"Attempting to send delta to disconnected client {token!r}" - ) - return - # Await the emit directly: wrapping it in a task does not unblock the - # caller (awaiting the task blocks just the same) and only adds task - # creation/scheduling overhead on every update. - await self.emit(str(constants.SocketEvent.EVENT), update, to=socket_record.sid) - # The emit may complete without suspending (the packet is queued, not - # sent). Yield one loop tick so the websocket writer can flush the - # packet before the caller potentially blocks the event loop (e.g. a - # sync event handler resuming after a yield). - await asyncio.sleep(0) - - async def on_event(self, sid: str, data: Any): - """Event for receiving front-end websocket events. - - Args: - sid: The Socket.IO session id. - data: The event data. - - Raises: - RuntimeError: If the Socket.IO is badly initialized. - EventDeserializationError: If the event data is not a dictionary. - """ - # Determine the token for this SID - if (token := self.sid_to_token.get(sid)) is None: - logger.warning( - f"Received event from session {sid} with no associated token. This may indicate a bug. Event data: {data}" - ) - return - - fields = data - - if isinstance(fields, str): - logger.warning( - "Received event data as a string. This generally should not happen and may indicate a bug." - f" Event data: {fields}" - ) - try: - fields = json.loads(fields) - except json.JSONDecodeError as ex: - msg = f"Failed to deserialize event data: {fields}." - raise exceptions.EventDeserializationError(msg) from ex - - if not isinstance(fields, dict): - msg = f"Event data must be a dictionary, but received {fields} of type {type(fields)}." - raise exceptions.EventDeserializationError(msg) - - try: - # Get the event. - event = Event(**{k: v for k, v in fields.items() if k in _EVENT_FIELDS}) - except (TypeError, ValueError) as ex: - msg = f"Failed to deserialize event data: {fields}." - raise exceptions.EventDeserializationError(msg) from ex - - static_router_data = self._static_router_data.get(sid) - if static_router_data is None: - # The connection was not seen by on_connect (e.g. namespace created - # after the socket connected); fall back to the connection environ. - if self.app.sio is None: - msg = "Socket.IO is not initialized." - raise RuntimeError(msg) - environ = self.app.sio.get_environ(sid, self.namespace) - if environ is None: - msg = "Socket.IO environ is not initialized." - raise RuntimeError(msg) - static_router_data = self._static_router_data[sid] = ( - self._build_static_router_data(sid, environ) - ) - router_data = event.router_data - router_data.update(static_router_data) - # The cached headers reach the event, and from there `state.router_data`, - # which is a plain mutable dict: sharing the mapping would let a handler - # mutating `self.router_data["headers"]` corrupt the connection cache for - # every later event on this socket. The shallow copy is ~17x cheaper than - # the per-event header decode it replaced, so the cache still pays off. - router_data[constants.RouteVar.HEADERS] = static_router_data[ - constants.RouteVar.HEADERS - ].copy() - router_data.update({ - constants.RouteVar.QUERY: format.format_query_params(event.router_data), - constants.RouteVar.CLIENT_TOKEN: token, - }) - router_data[constants.RouteVar.PATH] = "/" + ( - self.app.router(path) or "404" - if (path := router_data.get(constants.RouteVar.PATH)) - else "404" - ).removeprefix("/") - if not otel.enabled: - await self.app.event_processor.enqueue(token, event) - return - with otel.remote_context(fields): - await self.app.event_processor.enqueue(token, event) - - async def on_ping(self, sid: str): - """Event for testing the API endpoint. - - Args: - sid: The Socket.IO session id. - """ - # Emit the test event. - await self.emit(str(constants.SocketEvent.PING), "pong", to=sid) - - async def on_client_error(self, sid: str, data: Any = None): - """Handle errors reported by the frontend. - - This is a dedicated socket event rather than a state event - (``FrontendEventExceptionState.handle_frontend_exception``) because a - state event is addressed by a handler name the frontend derives from - its own state definitions. When those definitions are what disagree - with the backend -- the case this handler exists to report -- the name - may not resolve and the report is lost. A fixed socket event name - cannot drift, and it still gets through after the frontend has stopped - sending events on detecting the mismatch. - - Reports are routed through the app's ``frontend_exception_handler``, - so frontend errors (especially state update processing errors) are - visible in backend logs and reach custom exception handlers. - - Args: - sid: The Socket.IO session id. - data: The error data from the client. Defaults to None because - python-socketio dispatches a payload-less emit as - ``on_client_error(sid)``; the malformed-payload guard below - then drops it without raising. - """ - if not isinstance(data, dict): - logger.debug(f"Ignoring malformed client_error payload from SID {sid}.") - return - - # Check the sender and the rate limits before sanitizing: sanitizing is - # linear in the size of the client-supplied values, and reports that are - # dropped here must not cost more than the check itself. - if sid not in self.sid_to_token: - # Sockets without a linked token are not known clients; don't let - # them write error-level entries into the backend logs. - logger.debug(f"Ignoring client_error report from unknown SID {sid}.") - return - - # Rate limit per session so a client cannot flood the backend logs. - error_count = self._client_error_counts.get(sid, 0) - if error_count >= self._MAX_CLIENT_ERRORS_PER_SID: - return - - # Also bound total entries per time window: per-SID budgets reset on - # reconnect, so they alone do not stop scripted reconnect loops. - now = time.monotonic() - if now - self._client_error_window_start > self._CLIENT_ERROR_WINDOW_SECONDS: - self._client_error_window_start = now - self._client_error_window_count = 0 - if self._client_error_window_count >= self._MAX_CLIENT_ERRORS_PER_WINDOW: - if self._client_error_window_count == self._MAX_CLIENT_ERRORS_PER_WINDOW: - # Warn once per window so suppression is visible in the logs - # and a flooding client cannot silently starve reports from - # other sessions. - self._client_error_window_count += 1 - logger.warning( - f"Received more than {self._MAX_CLIENT_ERRORS_PER_WINDOW} " - f"client_error reports in {self._CLIENT_ERROR_WINDOW_SECONDS:.0f}s; " - "suppressing further reports for this window." - ) - return - self._client_error_window_count += 1 - self._client_error_counts[sid] = error_count + 1 - - error_type = format.sanitize_client_log_value(data.get("error_type", "unknown")) - if error_type == constants.ClientErrorType.DISPATCH_MISSING: - substate = format.sanitize_client_log_value(data.get("substate", "")) - report = ( - f"[SID: {sid}] Skipped state update for unknown substate(s) " - f"'{substate}': no dispatch function registered. " - "This indicates a frontend/backend state mismatch. " - "Rebuild the frontend or check that api_url points to the matching backend." - ) - else: - message = format.sanitize_client_log_value( - data.get("message", "No error message provided") - ) - report = f"[SID: {sid}] {error_type}: {message}" - # Route through the app's frontend exception handler so custom - # handlers (e.g. error trackers) receive client errors too. - self.app.frontend_exception_handler(Exception(report)) - - async def link_token_to_sid( - self, sid: str, token: str, *, update_state: bool = True - ): - """Link a token to a session id. - - Args: - sid: The Socket.IO session id. - token: The client token. - update_state: Whether to record the new sid and token on the state - now. The connect skips it when the CONNECT packet carries the - boot event, which records them as the connection's first event - without an extra load and save of the whole state tree. - """ - # Use TokenManager for duplicate detection and Redis support - new_token = await self._token_manager.link_token_to_sid(token, sid) - - if new_token: - # Duplicate detected, emit new token to client - await self.emit("new_token", new_token, to=sid) - - # Update client state to apply new sid/token for running background tasks. - if update_state and self.app._state is not None: - async with self.app.state_manager.modify_state( - BaseStateToken(ident=new_token or token, cls=self.app._state) - ) as state: - state.router_data[constants.RouteVar.SESSION_ID] = sid - # Record the identity the state was loaded under; duplicate-token - # handling can hand back a fresh one here. - state.router_data[constants.RouteVar.CLIENT_TOKEN] = new_token or token - # Rebuild from router_data to keep the session var in step with it. - if ( - session := SessionData.from_router_data(state.router_data) - ) != state.rx_router_session: - state.rx_router_session = session + if (target := _SOCKETIO_EXPORTS.get(name)) is None: + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) + module_name, attribute = target + try: + module = importlib.import_module(module_name) + except ImportError as ex: + # An AttributeError, so hasattr() and getattr() with a default work. + msg = f"{__name__}.{name} needs {_SOCKETIO_PACKAGE}." + raise AttributeError(msg) from ex + return getattr(module, attribute) diff --git a/reflex/channels.py b/reflex/channels.py new file mode 100644 index 00000000000..06336ad80e1 --- /dev/null +++ b/reflex/channels.py @@ -0,0 +1,374 @@ +"""Multiplexed side channels on the Reflex event websocket. + +A channel carries an application-defined message stream -- including binary +payloads -- over the connection that already delivers state deltas, so a +package needing its own data plane inherits the app's origin checks, client +token, reconnect handling and proxy configuration instead of reimplementing +them on a second socket. + +Register a channel on the app and messages from the matching client-side +channel are dispatched to it:: + + class Ticks(rx.channels.Channel): + name = "ticks" + + async def on_message(self, session, event, data, buffers): + if event == "subscribe": + session.join(data["symbol"]) + + app.register_channel(Ticks()) +""" + +from __future__ import annotations + +import dataclasses +import re +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable, Sequence +from typing import Any, ClassVar, TypeVar + +from typing_extensions import Buffer + +# Channel names ride in every frame and name a client-side channel handle; +# keep them short and free of JSON or URL escaping. +_NAME_PATTERN = re.compile(r"[A-Za-z0-9_./:-]{1,64}") + +# Maximum binary attachments one channel message may carry, in either +# direction. Bounds the work a single inbound frame can ask for. +MAX_MESSAGE_BUFFERS = 64 + +# Names a client-side channel handle reports its own lifecycle under. A +# message may not use them, or a consumer could not tell an application +# message from the transport event it is named after. Names starting with an +# underscore are reserved too: the channel protocol's own messages use them. +RESERVED_EVENTS = frozenset({"connect", "disconnect", "error"}) + +# Sends one channel message to connected sessions: (sids, channel, event, +# data, buffers), returning whether it accepted any of them. Supplied by the +# transport when the session opens. It takes every recipient at once so a +# fan-out serializes the frame only once. +ChannelSender = Callable[ + [Sequence[str], str, str, Any, Sequence[Buffer]], Awaitable[bool] +] + + +def is_reserved_event(event: str) -> bool: + """Whether a message name is reserved for the handle's lifecycle or the protocol. + + Args: + event: The message name. + + Returns: + Whether an application message may not use it. + """ + return event in RESERVED_EVENTS or event.startswith("_") + + +def _validate_message(event: str, buffers: Sequence[Buffer]) -> None: + """Check a message against what a frame may carry, before any is sent. + + Args: + event: The message name. + buffers: The binary attachments. + + Raises: + ValueError: If the message name is not a usable string or is + reserved, or the message carries more attachments than a frame + may hold. + """ + if not isinstance(event, str) or not event: + msg = f"Channel message name {event!r} is not a non-empty string." + raise ValueError(msg) + if is_reserved_event(event): + msg = ( + f"Channel message name {event!r} is reserved: the client-side " + "handle reports its own lifecycle under it, and names starting " + "with '_' belong to the channel protocol." + ) + raise ValueError(msg) + if len(buffers) > MAX_MESSAGE_BUFFERS: + msg = ( + f"Channel message {event!r} carries {len(buffers)} attachments, " + f"over the {MAX_MESSAGE_BUFFERS} a frame may hold." + ) + raise ValueError(msg) + + +_K = TypeVar("_K") +_V = TypeVar("_V") + + +def _discard(index: dict[_K, set[_V]], key: _K, member: _V) -> None: + """Remove a member from an index of sets, dropping the key once it is empty. + + Args: + index: The index. + key: The key whose set holds the member. + member: The member to remove. + """ + members = index.get(key) + if members is not None: + members.discard(member) + if not members: + del index[key] + + +def validate_channel_name(name: Any) -> None: + """Check a channel name against the wire format. + + Args: + name: The channel name declared by a Channel subclass. + + Raises: + ValueError: If the name is missing or unusable on the wire. + """ + if not isinstance(name, str) or not _NAME_PATTERN.fullmatch(name): + msg = ( + f"Invalid channel name {name!r}: expected 1-64 characters from " + "[A-Za-z0-9_./:-]." + ) + raise ValueError(msg) + + +@dataclasses.dataclass(eq=False, slots=True) +class ChannelSession: + """One client connection's participation in a channel.""" + + # The transport session id, shared with the app's event session. + sid: str + + # The Reflex client token (tab) the connection authenticated with. + client_token: str + + # The channel this session belongs to. + channel: Channel + + # Per-connection scratch space owned by the channel implementation. + data: dict[str, Any] + + _send: ChannelSender + + _rooms: set[str] + + # Whether the client is still connected. A handler that awaits (a rebuild, + # a thread hop) can come back to a session whose socket is gone; anything + # expensive or long-lived should check before proceeding. + open: bool = True + + async def send( + self, event: str, data: Any = None, buffers: Sequence[Buffer] = () + ) -> None: + """Send one message to this client. + + Args: + event: The message name. + data: The JSON-serializable metadata. + buffers: Binary attachments delivered alongside the metadata. + + Raises: + ValueError: If the message name is not a usable string or is + reserved, or the message carries more attachments than a + frame may hold. + """ + _validate_message(event, buffers) + await self._send((self.sid,), self.channel.name, event, data, buffers) + + def join(self, room: str) -> None: + """Add this session to a room for fan-out, unless it is closed. + + A handler resuming after the socket went away would otherwise put the + session back in the room map, where nothing removes it again and every + fan-out keeps addressing the dead connection. + + Args: + room: The room name. + """ + if not self.open: + return + self._rooms.add(room) + self.channel._rooms.setdefault(room, set()).add(self) + + def leave(self, room: str) -> None: + """Remove this session from a room. + + Args: + room: The room name. + """ + self._rooms.discard(room) + _discard(self.channel._rooms, room, self) + + +class Channel(ABC): + """A named message stream multiplexed onto the event websocket. + + Rooms and sessions are local to the worker holding the connection, which + is what a socket-bound side channel can offer: a client reconnecting to + another worker opens its session there. State that must outlive a + connection belongs in Reflex state, not in the channel. + """ + + # The channel name, matching the name the client opens. + name: ClassVar[str] + + # Whether inbound messages from this channel may carry binary + # attachments. Off by default: a channel that never expects binary + # answers an error instead of handing one to its handler. + accepts_binary: ClassVar[bool] = False + + def __init__(self): + """Initialize the channel's session and room bookkeeping.""" + validate_channel_name(getattr(type(self), "name", None)) + # Sessions by client token, for token-addressed sends. + self._sessions: dict[str, set[ChannelSession]] = {} + # Room name to member sessions. + self._rooms: dict[str, set[ChannelSession]] = {} + + async def on_open(self, session: ChannelSession) -> None: + """Handle a client opening the channel. + + Args: + session: The opening session. + """ + return + + @abstractmethod + async def on_message( + self, + session: ChannelSession, + event: str, + data: Any, + buffers: list[bytes], + ) -> None: + """Handle one message from a client. + + Every field is client-controlled and unvalidated. + + Args: + session: The sending session. + event: The message name. + data: The JSON metadata, as decoded from the frame. + buffers: Binary attachments, empty unless ``accepts_binary``. + """ + + async def on_close(self, session: ChannelSession) -> None: + """Handle a session going away. + + Args: + session: The closing session. + """ + return + + async def send_to_room( + self, + room: str, + event: str, + data: Any = None, + buffers: Sequence[Buffer] = (), + ) -> None: + """Send one message to every session in a room. + + Args: + room: The room name. + event: The message name. + data: The JSON-serializable metadata. + buffers: Binary attachments delivered alongside the metadata. + + Raises: + ValueError: If the message name is not a usable string or is + reserved, or the message carries more attachments than a + frame may hold. Raised before anything is sent, so a fan-out + fails whole rather than reaching some clients. + """ + await self._fan_out(self._rooms.get(room), event, data, buffers) + + async def send_to_token( + self, + client_token: str, + event: str, + data: Any = None, + buffers: Sequence[Buffer] = (), + ) -> bool: + """Send one message to every session of a client token on this worker. + + Args: + client_token: The Reflex client token (tab). + event: The message name. + data: The JSON-serializable metadata. + buffers: Binary attachments delivered alongside the metadata. + + Returns: + Whether a session received the message. + + Raises: + ValueError: If the message name is not a usable string or is + reserved, or the message carries more attachments than a + frame may hold. Raised before anything is sent, so a fan-out + fails whole rather than reaching some clients. + """ + return await self._fan_out( + self._sessions.get(client_token), event, data, buffers + ) + + async def _fan_out( + self, + sessions: set[ChannelSession] | None, + event: str, + data: Any, + buffers: Sequence[Buffer], + ) -> bool: + """Send one message to a group of sessions, validated before any is sent. + + Args: + sessions: The recipients, if any. + event: The message name. + data: The JSON-serializable metadata. + buffers: Binary attachments delivered alongside the metadata. + + Returns: + Whether a session received the message. + """ + _validate_message(event, buffers) + if not sessions: + return False + # Every session of a channel is opened by the app's one transport, so + # any member can carry the frame for all of them -- serialized once + # rather than once per recipient. The recipients are collected first: + # a send can close a session and mutate the group. + return await next(iter(sessions))._send( + [session.sid for session in sessions], self.name, event, data, buffers + ) + + def open_session( + self, sid: str, client_token: str, send: ChannelSender + ) -> ChannelSession: + """Create and track a session for a connection. Called by the transport. + + Args: + sid: The transport session id. + client_token: The client token the connection authenticated with. + send: The transport's send callable. + + Returns: + The new session. + """ + session = ChannelSession( + sid=sid, + client_token=client_token, + channel=self, + data={}, + _send=send, + _rooms=set(), + ) + self._sessions.setdefault(client_token, set()).add(session) + return session + + def forget_session(self, session: ChannelSession) -> None: + """Drop a session's rooms and tracking. Called by the transport. + + Args: + session: The session going away. + """ + session.open = False + for room in tuple(session._rooms): + session.leave(room) + _discard(self._sessions, session.client_token, session) diff --git a/reflex/event_namespace.py b/reflex/event_namespace.py new file mode 100644 index 00000000000..35c1e90ca7b --- /dev/null +++ b/reflex/event_namespace.py @@ -0,0 +1,1683 @@ +"""Event namespaces bridging client sessions to the Reflex event loop.""" + +from __future__ import annotations + +import asyncio +import collections +import contextlib +import dataclasses +import json +import logging +import time +import urllib.parse +import uuid +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Iterable, Mapping, MutableMapping, Sequence +from datetime import timedelta +from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple + +from reflex_base import constants, otel +from reflex_base.config import get_config +from reflex_base.environment import auto_reload_cooldown, environment +from reflex_base.event import _EVENT_FIELDS, Event +from starlette.websockets import WebSocket, WebSocketDisconnect +from typing_extensions import Buffer + +from reflex.channels import MAX_MESSAGE_BUFFERS, ChannelSession, is_reserved_event +from reflex.istate.data import SessionData +from reflex.istate.manager.token import BaseStateToken +from reflex.state import StateUpdate +from reflex.utils import exceptions, format +from reflex.utils.token_manager import RedisTokenManager, TokenManager + +if TYPE_CHECKING: + from reflex.app import App + +logger = logging.getLogger(__name__) + +# Protocol-level message names for the plain WebSocket transport. These are +# reserved (underscore-prefixed) and never dispatched as application events. +# They must match the names in .templates/web/utils/helpers/websocket.js. +CONNECT_MESSAGE = "_connect" +HANDSHAKE_MESSAGE = "_handshake" +PING_MESSAGE = "_ping" +PONG_MESSAGE = "_pong" +OPEN_MESSAGE = "_open" +OPENED_MESSAGE = "_opened" +CLOSE_MESSAGE = "_close" +CHANNEL_ERROR_MESSAGE = "_error" + +# Close reason prefix of a connect frame refused for its size; the limit +# follows it. It must match the client's in websocket.js. +MAX_MESSAGE_SIZE_REASON = "max_message_size=" + +# Wire protocol version, announced in the handshake. The client gates channel +# frames on it: a backend that predates channels closes the connection on the +# binary frames they use. +PROTOCOL_VERSION = 2 + +# Binary channel frames align every attachment to this boundary so the client +# can view them as typed arrays without copying. +_FRAME_ALIGNMENT = 8 + +# Application-level socket event names, resolved once for the hot paths. +_EVENT = str(constants.SocketEvent.EVENT) +_PING = str(constants.SocketEvent.PING) +_CLIENT_ERROR = str(constants.SocketEvent.CLIENT_ERROR) + +# Frames carry no separator whitespace: it is pure overhead on the wire. +_COMPACT = (",", ":") + +# The heartbeat frame is static; serialize it once. +_PING_FRAME = json.dumps([PING_MESSAGE], separators=_COMPACT) + +# ASGI scope key holding the connection-scoped router_data. +_STATIC_ROUTER_DATA = "_reflex_static_router_data" + + +def _decode_header_octets(octets: bytes) -> str: + """Decode a header name or value, as UTF-8 where it is valid. + + Args: + octets: The raw header name or value. + + Returns: + The decoded text. + """ + try: + return octets.decode() + except UnicodeDecodeError: + # Header octets need not be UTF-8; Latin-1 maps every byte, as ASGI + # servers and Starlette decode them. + return octets.decode("latin-1") + + +def _decode_asgi_headers(headers: Iterable[tuple[bytes, bytes]]) -> dict[str, str]: + """Decode raw ASGI scope header pairs into a str-keyed dict. + + ASGI keeps a repeated header as separate pairs, so their values are combined in + order: cookies with "; " (RFC 9113 section 8.2.3), any other field with ", " + (RFC 9110 section 5.3). + + Args: + headers: Raw (name, value) byte pairs from the ASGI scope. + + Returns: + A dict mapping decoded header names to decoded values. + """ + decoded: dict[str, str] = {} + for raw_name, raw_value in headers: + # The common all-UTF-8 case decodes without a function call per header. + try: + name = raw_name.decode() + value = raw_value.decode() + except UnicodeDecodeError: + name = _decode_header_octets(raw_name) + value = _decode_header_octets(raw_value) + if name in decoded: + decoded[name] += ("; " if name == "cookie" else ", ") + value + else: + decoded[name] = value + return decoded + + +def build_static_router_data(sid: str, asgi_scope: Mapping[str, Any]) -> dict[str, Any]: + """Build the router_data entries that are constant for a connection. + + Args: + sid: The session id. + asgi_scope: The ASGI scope of the client connection. + + Returns: + The connection-scoped router_data entries. + """ + headers = _decode_asgi_headers(asgi_scope["headers"]) + + # Get the client IP. + if client := asgi_scope.get("client"): + client_ip = client[0] + headers["asgi-scope-client"] = client_ip + else: + client_ip = "0.0.0.0" + + # X-Forwarded-For is client-controlled unless the request arrived through + # a configured chain of trusted reverse proxies. + trusted_proxy_hops = get_config().trusted_proxy_hops + if trusted_proxy_hops: + forwarded_for = headers.get("x-forwarded-for", "") + forwarded_ips = [ip.strip() for ip in forwarded_for.split(",") if ip.strip()] + if len(forwarded_ips) >= trusted_proxy_hops: + client_ip = forwarded_ips[-trusted_proxy_hops] + return { + constants.RouteVar.SESSION_ID: sid, + constants.RouteVar.HEADERS: headers, + constants.RouteVar.CLIENT_IP: client_ip, + } + + +def connect_auth(auth: Any) -> tuple[Any, Any]: + """Get the boot event and frontend state schema a session's connect carries. + + The frontend sends its hydrate event with the connect, so the backend + processes it without waiting for the connect acknowledgement round trip. + + Args: + auth: The connect payload. + + Returns: + The boot event and the state schema hash, each None if not carried. + """ + if not isinstance(auth, dict): + return None, None + return ( + auth.get(constants.CompileVars.CONNECT_AUTH_EVENT), + auth.get(constants.CompileVars.CONNECT_AUTH_STATE_SCHEMA), + ) + + +class _SocketPolicy(NamedTuple): + """The event socket's heartbeat and message size settings.""" + + # Seconds between heartbeat pings. + ping_interval: float + + # Seconds past the interval a client may stay silent. + ping_timeout: float + + # The message size limit in bytes. + max_message_size: int + + @property + def silence_limit(self) -> float: + """Seconds a client may stay silent before its session ends. + + Returns: + The ping interval plus the ping timeout. + """ + return self.ping_interval + self.ping_timeout + + +def _socket_policy() -> _SocketPolicy: + """Read the event socket's settings, per connection. + + Returns: + The settings currently configured. + """ + return _SocketPolicy( + environment.REFLEX_SOCKET_INTERVAL.get().total_seconds(), + environment.REFLEX_SOCKET_TIMEOUT.get().total_seconds(), + environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE.get(), + ) + + +def _closing(sid: str, reason: str, code: int) -> int: + """Log why a session is closing, at debug level. + + Every reason is client-driven, so logging louder would let a hostile + client flood the logs. + + Args: + sid: The session id. + reason: Why the session closes. + code: The websocket close code. + + Returns: + The close code, to close the session with. + """ + logger.debug("Closing session %s: %s.", sid, reason) + return code + + +def _dumps(obj: Any) -> str: + """Serialize a frame compactly, as text that is valid UTF-8. + + Args: + obj: The frame to serialize. + + Returns: + The JSON text. + """ + text = format.json_dumps(obj, separators=_COMPACT) + if not text.isascii(): + try: + text.encode() + except UnicodeEncodeError: + # An unpaired surrogate (a non-UTF-8 file name decodes to one) + # cannot travel as text; escape it, as ensure_ascii would. + return text.encode(errors="backslashreplace").decode() + return text + + +def _parse_frame(text: str) -> Any: + """Parse a text frame. + + Args: + text: The raw frame text. + + Returns: + The decoded JSON, or None if the text is not JSON. + """ + try: + return json.loads(text) + except (ValueError, RecursionError): + # Besides a JSONDecodeError, the decoder raises a plain ValueError for + # an integer past the digit limit and exhausts its stack on deeply + # nested JSON; all are just a malformed frame here. + return None + + +def utf8_size(data: str | bytes) -> int: + """Size of a serialized message in UTF-8 bytes. + + ASCII payloads (the common case) are sized without encoding a copy. + + Args: + data: The serialized message, text or binary. + + Returns: + The number of bytes the message occupies on the wire. + """ + if isinstance(data, bytes): + return len(data) + return len(data) if data.isascii() else len(data.encode()) + + +def exceeds_message_limit(data: str | bytes, max_size: int) -> bool: + """Whether a received message is over the policy limit. + + The limit counts UTF-8 bytes, and UTF-8 encodes 1-4 bytes per character: + more characters than the limit is certainly over, a quarter or fewer + certainly under, so only the range between is measured exactly -- and an + ASCII payload, which is most of them, is measured without a copy. Mirrors + ``exceedsMessageLimit`` in .templates/web/utils/helpers/websocket.js. + + Args: + data: The received message, text or binary. + max_size: The limit in bytes. + + Returns: + Whether the message exceeds it. + """ + if isinstance(data, bytes): + return len(data) > max_size + length = len(data) + return length > max_size or (length * 4 > max_size and utf8_size(data) > max_size) + + +def encode_channel_frame( + event: str, data: Any, channel: str, buffers: Sequence[Buffer] +) -> bytes: + """Serialize a channel message carrying binary attachments. + + The frame is a 4-byte little-endian header length, the JSON header + ``[event, data, channel, [lengths]]``, then the attachments, each padded + so every payload starts on an 8-byte boundary. + + Args: + event: The message name. + data: The JSON-serializable metadata. + channel: The channel name. + buffers: The binary attachments. + + Returns: + The frame bytes. + """ + # In bytes: len() counts items, which is not the size of a typed array. + sizes = [memoryview(buffer).nbytes for buffer in buffers] + header = _dumps([event, data, channel, sizes]).encode() + parts: list[Buffer] = [len(header).to_bytes(4, "little"), header] + offset = 4 + len(header) + for buffer, size in zip(buffers, sizes, strict=True): + padding = -offset % _FRAME_ALIGNMENT + if padding: + parts.append(bytes(padding)) + parts.append(buffer) + offset += padding + size + return b"".join(parts) + + +def decode_channel_frame(frame: bytes) -> tuple[str, Any, str, list[bytes]]: + """Deserialize a binary channel frame. + + Args: + frame: The raw frame bytes. + + Returns: + The message name, metadata, channel name and binary attachments. + + Raises: + ValueError: If the frame does not follow the binary channel format. + """ + if len(frame) < 4: + msg = "Binary frame is too short to hold a header length." + raise ValueError(msg) + header_size = int.from_bytes(frame[:4], "little") + if 4 + header_size > len(frame): + # The frame is the only bound the header needs: the transport rejects + # one over the size limit before decoding it. + msg = f"Binary frame declares an unusable header size {header_size}." + raise ValueError(msg) + try: + header = json.loads(frame[4 : 4 + header_size]) + except RecursionError as ex: + # The JSON decoder recurses per nesting level, so a header nested + # deeper than its stack allows never parses. That is malformed input + # like any other, and callers should not have to know that the + # decoder signals it differently. + msg = "Binary frame header is nested too deeply." + raise ValueError(msg) from ex + match header: + case [str(event), data, str(channel), [*lengths]] if all( + isinstance(length, int) and not isinstance(length, bool) and length >= 0 + for length in lengths + ): + pass + case _: + msg = "Binary frame header is malformed." + raise ValueError(msg) + if len(lengths) > MAX_MESSAGE_BUFFERS: + msg = f"Binary frame carries more than {MAX_MESSAGE_BUFFERS} attachments." + raise ValueError(msg) + buffers: list[bytes] = [] + offset = 4 + header_size + for length in lengths: + offset += -offset % _FRAME_ALIGNMENT + end = offset + length + if end > len(frame): + msg = "Binary frame is shorter than its declared attachments." + raise ValueError(msg) + buffers.append(frame[offset:end]) + offset = end + return event, data, channel, buffers + + +# Size of the frames one connection may have waiting behind the one being +# written. Past it the client has stopped reading, and is dropped: it would +# otherwise buffer without bound until its heartbeat times out. +_MAX_SEND_BACKLOG = 16 * 1024 * 1024 + + +class _Connection: + """One client websocket, as the transport sends to and closes it. + + Sending only queues a frame for the connection's own writer task, so a + client that stops reading cannot stall the coroutines emitting to it. + Closing ends the session's receive loop, which closes the socket itself: + a server may not complete a close sent while a receive is pending on it. + """ + + def __init__(self, websocket: WebSocket): + """Start the connection's writer. + + Args: + websocket: The client websocket connection. + """ + self.websocket = websocket + # When a frame last arrived, for the heartbeat's liveness check. + self.last_received = time.monotonic() + # The code to close with, once the connection is ending. + self.close_code: int | None = None + self._frames: collections.deque[str | bytes] = collections.deque() + # Size of the frames waiting behind the one being written. + self._backlog = 0 + self._ready = asyncio.Event() + self._lifetime: asyncio.Timeout | None = None + self._writer = asyncio.create_task(self._write()) + + def send(self, frame: str | bytes) -> bool: + """Queue a frame for the client, closing a connection that stopped reading. + + Args: + frame: The serialized text or binary frame. + + Returns: + Whether the frame was queued on a connection that is not closing. + """ + if self.close_code is not None: + return False + frames = self._frames + backlog = self._backlog + len(frame) + # One frame may exceed the bound; frames piling up behind it may not. + # Checked before queueing: the writer runs on until teardown stops it. + if backlog > _MAX_SEND_BACKLOG and frames: + logger.debug("Closing a session whose client stopped reading.") + self.close(1008) + return False + frames.append(frame) + self._backlog = backlog + self._ready.set() + return True + + async def _write(self) -> None: + """Write the queued frames in order until the connection ends.""" + websocket = self.websocket + frames = self._frames + try: + while True: + await self._ready.wait() + self._ready.clear() + while frames: + frame = frames.popleft() + self._backlog -= len(frame) + if isinstance(frame, str): + await websocket.send_text(frame) + else: + await websocket.send_bytes(frame) + except Exception: + # The connection went away mid-write; its receive loop cleans up. + logger.debug("Failed to write to a client websocket.", exc_info=True) + + def attach(self, lifetime: asyncio.Timeout) -> None: + """Let close() end the receive loop running under a lifetime. + + Args: + lifetime: The timeout the session's receive loop runs under. + """ + self._lifetime = lifetime + if self.close_code is not None: + lifetime.reschedule(0) + + def close(self, code: int) -> None: + """End the session with a close code, dropping further sends. + + Args: + code: The websocket close code. + """ + if self.close_code is not None: + return + self.close_code = code + lifetime = self._lifetime + if lifetime is not None and not lifetime.expired(): + # Fires at the receive loop's next await, which ends it. + lifetime.reschedule(0) + + def stop(self) -> None: + """Stop writing, dropping the frames not written yet.""" + self._lifetime = None + self._writer.cancel() + + +@dataclasses.dataclass +class _ErrorBudget: + """Bounds error-level logging driven by one kind of client traffic. + + Per session, because one client must not fill the log; and per time + window, because per-session budgets reset on reconnect and so do not stop + scripted reconnect loops. + """ + + # What the budget covers, for the message that announces suppression. + label: str + + max_per_session: int + max_per_window: int + window_seconds: float + + counts: dict[str, int] = dataclasses.field(default_factory=dict) + window_start: float = 0.0 + window_count: int = 0 + + def allows(self, sid: str) -> bool: + """Whether another record for this session fits the budget. + + Args: + sid: The session id. + + Returns: + Whether the record may be written. + """ + session_count = self.counts.get(sid, 0) + if session_count >= self.max_per_session: + return False + now = time.monotonic() + if now - self.window_start > self.window_seconds: + self.window_start = now + self.window_count = 0 + if self.window_count >= self.max_per_window: + if self.window_count == self.max_per_window: + # Warn once per window so suppression is visible in the logs + # and a flooding client cannot silently starve reports from + # other sessions. + self.window_count += 1 + logger.warning( + f"More than {self.max_per_window} {self.label} in " + f"{self.window_seconds:.0f}s; suppressing further reports " + "for this window." + ) + return False + self.window_count += 1 + self.counts[sid] = session_count + 1 + return True + + def forget(self, sid: str) -> None: + """Drop a disconnected session's counter. + + Args: + sid: The session id. + """ + self.counts.pop(sid, None) + + +class BaseEventNamespace(ABC): + """Transport-agnostic handler for client event sessions.""" + + # The application object. + app: App + + # Whether the transport carries channels (reflex.channels). + supports_channels: ClassVar[bool] = False + + # Maximum error-level log entries a single session may produce, per kind + # of client-triggered error, before further ones from it are dropped. + _MAX_CLIENT_ERRORS_PER_SID = 5 + + # Process-wide bound on those entries per time window; per-SID budgets + # alone reset on reconnect, so scripted reconnects could otherwise flood + # the logs. + _CLIENT_ERROR_WINDOW_SECONDS = 60.0 + _MAX_CLIENT_ERRORS_PER_WINDOW = 20 + + def __init__(self, namespace: str, app: App): + """Initialize the event namespace. + + Args: + namespace: The namespace. + app: The application object. + """ + self.namespace = namespace + self.app = app + + # Use TokenManager for distributed duplicate tab prevention + self._token_manager = TokenManager.create() + + # Client-reported errors and server-side handler failures are both + # driven by client traffic, but a client chooses how many reports it + # sends while a handler traceback means a real bug, so one cannot be + # allowed to suppress the other. + self._client_error_budget = self._error_budget("client_error reports") + self._handler_error_budget = self._error_budget("handler errors") + + @classmethod + def _error_budget(cls, label: str) -> _ErrorBudget: + """Build a budget for one kind of client-triggered error logging. + + Args: + label: What the budget covers, for the suppression message. + + Returns: + The budget. + """ + return _ErrorBudget( + label=label, + max_per_session=cls._MAX_CLIENT_ERRORS_PER_SID, + max_per_window=cls._MAX_CLIENT_ERRORS_PER_WINDOW, + window_seconds=cls._CLIENT_ERROR_WINDOW_SECONDS, + ) + + @property + def token_to_sid(self) -> Mapping[str, str]: + """Token to SID mapping for backward compatibility. + + Note: this mapping is read-only. + + Returns: + The token to SID mapping. + """ + # For backward compatibility, expose the underlying dict + return self._token_manager.token_to_sid + + @property + def sid_to_token(self) -> dict[str, str]: + """SID to token mapping for backward compatibility. + + Returns: + The SID to token mapping dict. + """ + # For backward compatibility, expose the underlying dict + return self._token_manager.sid_to_token + + @abstractmethod + async def emit(self, event: str, data: Any = None, to: str | None = None) -> None: + """Emit an event to a connected client session, or to all of them. + + Args: + event: The event name. + data: The event payload. + to: The session id to emit to; every connected session if None. + """ + + async def handle_connect( + self, + sid: str, + query_string: str, + subprotocol: str | None, + frontend_schema: Any = None, + *, + update_state: bool = True, + ) -> None: + """Handle a new client session connecting. + + Args: + sid: The session id. + query_string: The raw query string of the connection request. + subprotocol: The websocket subprotocol offered by the client. + frontend_schema: The state schema hash the frontend was compiled + for; None from a frontend too old to send one. + update_state: Whether to record the new sid and token on the state; + see link_token_to_sid. + """ + if otel.enabled: + otel.record_connection(1) + if isinstance(self._token_manager, RedisTokenManager): + # Make sure this instance is watching for updates from other instances. + self._token_manager.ensure_lost_and_found_task(self.emit_update) + query_params = urllib.parse.parse_qs(query_string) + token_list = query_params.get("token", []) + if not token_list: + # A Reflex client always sends a token; the transport closes the + # session, so a warning per hostile connect would only flood logs. + logger.debug(f"No token provided in connection for session {sid}.") + return + await self.link_token_to_sid(sid, token_list[0], update_state=update_state) + # Only report the version for linked sessions; the value is + # client-controlled, so sanitize it before it reaches the logs. + if subprotocol and subprotocol != constants.Reflex.VERSION: + logger.warning( + f"Frontend version {format.sanitize_client_log_value(subprotocol)} " + f"for session {sid} does not match the backend version {constants.Reflex.VERSION}." + ) + # A frontend compiled for another state tree is told so before it + # hydrates. The backend hash is cached on the state class. + if ( + frontend_schema is not None + and self.app._state is not None + and frontend_schema + != (backend_schema := self.app._state._to_frontend_schema()) + ): + logger.warning( + "Frontend state schema " + f"{format.sanitize_client_log_value(frontend_schema, max_length=64)} " + f"for session {sid} does not match the backend state schema " + f"{backend_schema}: the frontend was built for another version of " + "the app and is told to reload." + ) + await self.emit( + str(constants.SocketEvent.STATE_MISMATCH), + { + "backend": backend_schema, + "reload_cooldown_ms": auto_reload_cooldown() + // timedelta(milliseconds=1), + }, + to=sid, + ) + + async def _open_session( + self, sid: str, auth: Any, query_string: str, subprotocol: str | None + ) -> tuple[bool, Any]: + """Link a connecting session, reading the boot event its connect carries. + + Args: + sid: The session id. + auth: The connect payload. + query_string: The raw query string of the connection request. + subprotocol: The websocket subprotocol offered by the client. + + Returns: + Whether the session linked a token, and its boot event if any. + """ + boot_event, frontend_schema = connect_auth(auth) + await self.handle_connect( + sid, + query_string, + subprotocol, + frontend_schema, + update_state=boot_event is None, + ) + return sid in self._token_manager.sid_to_token, boot_event + + async def _handle_boot_event( + self, sid: str, dispatch: Awaitable[None] + ) -> int | None: + """Await the dispatch of a connect's boot event, the session's first event. + + A session whose boot event failed is ended: kept open, it would leave + the client waiting for a hydrate that never comes instead of retrying. + + Args: + sid: The session id. + dispatch: The transport's dispatch of the boot event, raising + whatever made it fail. + + Returns: + None once dispatched, else the close code to end the session with. + """ + try: + await dispatch + except Exception as exc: + return self._handler_failed(sid, "the boot event", exc) or 1011 + return None + + def _handler_failed(self, sid: str, what: str, error: Exception) -> int | None: + """Apply the failure policy to a handler that raised. + + Undeserializable input is client-controlled and a Reflex client never + sends it, so it ends the session quietly; anything else is a + server-side bug, logged within the session's error budget. + + Args: + sid: The session id. + what: What was being handled, for the log. + error: The exception the handler raised. + + Returns: + The close code the session must end with, or None to keep it. + """ + if isinstance(error, exceptions.EventDeserializationError): + return _closing(sid, f"{what} did not deserialize", 1002) + self._log_handler_failure( + sid, f"Error handling {what} for session {sid}.", error + ) + return None + + def handle_disconnect(self, sid: str) -> asyncio.Task | None: + """Handle a client session disconnecting. + + Args: + sid: The session id. + + Returns: + An asyncio Task for cleaning up the token, or None. + """ + if otel.enabled: + otel.record_connection(-1) + self._client_error_budget.forget(sid) + self._handler_error_budget.forget(sid) + # Get token before cleaning up + disconnect_token = self.sid_to_token.get(sid) + if disconnect_token: + socket_record = self._token_manager.token_to_socket.get(disconnect_token) + if ( + socket_record is not None + and socket_record.sid == sid + and socket_record.instance_id == self._token_manager.instance_id + and (processor := self.app._event_processor) is not None + ): + processor._on_disconnect(disconnect_token) + # Use async cleanup through token manager + task = asyncio.create_task( + self._token_manager.disconnect_token(disconnect_token, sid), + name=f"reflex_disconnect_token|{disconnect_token}|{time.time()}", + ) + # Don't await to avoid blocking disconnect, but handle potential errors + task.add_done_callback( + lambda t: ( + t.exception() + and logger.error(f"Token cleanup error: {t.exception()}") + ) + ) + return task + return None + + async def emit_update(self, update: StateUpdate, token: str) -> None: + """Emit an update to the client. + + Args: + update: The state update to send. + token: The client token (tab) associated with the event. + """ + socket_record = self._token_manager.token_to_socket.get(token) + if ( + socket_record is None + or socket_record.instance_id != self._token_manager.instance_id + ): + if isinstance(self._token_manager, RedisTokenManager): + # The socket belongs to another instance of the app, send it to the lost and found. + await self._token_manager.emit_lost_and_found(token, update) + else: + # If the socket record is None, we are not connected to a client. Prevent sending + # updates to all clients. + logger.warning( + f"Attempting to send delta to disconnected client {token!r}" + ) + return + # Awaiting a task wrapping the emit blocks just the same, so await it + # directly and skip the task overhead. + await self.emit(_EVENT, update, to=socket_record.sid) + # The emit only queues the packet; yield a tick so the writer can flush + # it before the caller potentially blocks the loop. + await asyncio.sleep(0) + + async def handle_event( + self, sid: str, data: Any, asgi_scope: MutableMapping[str, Any] + ) -> None: + """Handle an incoming front-end event. + + Args: + sid: The session id. + data: The event data. + asgi_scope: The ASGI scope of the client connection. + + Raises: + EventDeserializationError: If the event data is malformed. + """ + # Determine the token for this SID, without the property's call. + if (token := self._token_manager.sid_to_token.get(sid)) is None: + # The mapping is dropped when a token moves to another socket or + # its record goes stale, so a live connection can reach this and + # keep sending. Log it per frame at debug, without the + # client-controlled payload: at warning level it would be a log + # flood and an injection vector both. + logger.debug(f"Ignoring event from session {sid} with no linked token.") + return + + # Both transports JSON-decode the frame, so a Reflex client's event + # arrives as a dict; anything else (including a JSON-encoded string) + # is rejected rather than logged per frame. + if not isinstance(data, dict): + msg = f"Event data must be a dictionary, but received {data} of type {type(data)}." + raise exceptions.EventDeserializationError(msg) + + try: + # Get the event. + event = Event(**{k: v for k, v in data.items() if k in _EVENT_FIELDS}) + except (TypeError, ValueError) as ex: + msg = f"Failed to deserialize event data: {data}." + raise exceptions.EventDeserializationError(msg) from ex + + # The dataclass does not validate field types. A router_data that is + # no dict fails its normalization below. + if not isinstance(event.name, str) or not isinstance(event.payload, dict): + msg = "Event fields have invalid types." + raise exceptions.EventDeserializationError(msg) + + # Headers, client IP, and session id cannot change for the lifetime of + # the connection; derive them once and cache them on its ASGI scope, + # which is per-connection state, instead of on every event. + static_router_data = asgi_scope.get(_STATIC_ROUTER_DATA) + if static_router_data is None: + static_router_data = asgi_scope[_STATIC_ROUTER_DATA] = ( + build_static_router_data(sid, asgi_scope) + ) + + router_data = event.router_data + try: + router_data.update(static_router_data) + # The cached headers reach the event, and from there + # `state.router_data`, a plain mutable dict: sharing the mapping + # would let a handler mutating `self.router_data["headers"]` + # corrupt the connection cache for every later event on this + # socket. The shallow copy is far cheaper than the per-event + # header decode it replaced, so the cache still pays off. + router_data[constants.RouteVar.HEADERS] = static_router_data[ + constants.RouteVar.HEADERS + ].copy() + # The nested values are still client-controlled. + router_data.update({ + constants.RouteVar.QUERY: format.format_query_params(event.router_data), + constants.RouteVar.CLIENT_TOKEN: token, + }) + router_data[constants.RouteVar.PATH] = "/" + ( + self.app.router(path) or "404" + if (path := router_data.get(constants.RouteVar.PATH)) + else "404" + ).removeprefix("/") + except (AttributeError, LookupError, TypeError, ValueError) as ex: + msg = "Failed to normalize event router_data." + raise exceptions.EventDeserializationError(msg) from ex + if not otel.enabled: + await self.app.event_processor.enqueue(token, event) + return + with otel.remote_context(data): + await self.app.event_processor.enqueue(token, event) + + async def handle_ping(self, sid: str) -> None: + """Handle an application-level ping test event. + + Args: + sid: The session id. + """ + # Emit the test event. + await self.emit(_PING, "pong", to=sid) + + def _log_handler_failure( + self, sid: str, message: str, error: BaseException + ) -> None: + """Report a handler that raised, within the session's error budget. + + A client can keep sending whatever made the handler raise, so the + traceback is budgeted like any other client-triggered error. + + Args: + sid: The session id. + message: What failed. + error: The exception to attach. + """ + if self._handler_error_budget.allows(sid): + logger.error(message, exc_info=error) + else: + logger.debug(f"Suppressed a repeated handler error for session {sid}.") + + async def handle_client_error(self, sid: str, data: Any) -> None: + """Handle errors reported by the frontend. + + This is a dedicated socket event rather than a state event + (``FrontendEventExceptionState.handle_frontend_exception``) because a + state event is addressed by a handler name the frontend derives from + its own state definitions. When those definitions are what disagree + with the backend -- the case this handler exists to report -- the name + may not resolve and the report is lost. A fixed socket event name + cannot drift, and it still gets through after the frontend has stopped + sending events on detecting the mismatch. + + Reports are routed through the app's ``frontend_exception_handler``, + so frontend errors (especially state update processing errors) are + visible in backend logs and reach custom exception handlers. + + Args: + sid: The session id. + data: The error data from the client. + """ + if not isinstance(data, dict): + logger.debug(f"Ignoring malformed client_error payload from SID {sid}.") + return + + # Check the sender and the rate limits before sanitizing: sanitizing is + # linear in the size of the client-supplied values, and reports that are + # dropped here must not cost more than the check itself. + if sid not in self.sid_to_token: + # Sockets without a linked token are not known clients; don't let + # them write error-level entries into the backend logs. + logger.debug(f"Ignoring client_error report from unknown SID {sid}.") + return + + if not self._client_error_budget.allows(sid): + return + + error_type = format.sanitize_client_log_value(data.get("error_type", "unknown")) + if error_type == constants.ClientErrorType.DISPATCH_MISSING: + substate = format.sanitize_client_log_value(data.get("substate", "")) + report = ( + f"[SID: {sid}] Skipped state update for unknown substate(s) " + f"'{substate}': no dispatch function registered. " + "This indicates a frontend/backend state mismatch. " + "Rebuild the frontend or check that api_url points to the matching backend." + ) + else: + message = format.sanitize_client_log_value( + data.get("message", "No error message provided") + ) + report = f"[SID: {sid}] {error_type}: {message}" + # Route through the app's frontend exception handler so custom + # handlers (e.g. error trackers) receive client errors too. + self.app.frontend_exception_handler(Exception(report)) + + async def link_token_to_sid( + self, sid: str, token: str, *, update_state: bool = True + ): + """Link a token to a session id. + + Args: + sid: The session id. + token: The client token. + update_state: Whether to record the new sid and token on the state + now. A connect that carries the boot event skips it: processed + as the session's first event, that records them without an + extra load and save of the whole state tree. + """ + # Use TokenManager for duplicate detection and Redis support + new_token = await self._token_manager.link_token_to_sid(token, sid) + + if new_token: + # Duplicate detected, emit new token to client + await self.emit("new_token", new_token, to=sid) + + # Update client state to apply new sid/token for running background tasks. + if update_state and self.app._state is not None: + async with self.app.state_manager.modify_state( + BaseStateToken(ident=new_token or token, cls=self.app._state) + ) as state: + state.router_data[constants.RouteVar.SESSION_ID] = sid + # Record the identity the state was loaded under; duplicate-token + # handling can hand back a fresh one here. + state.router_data[constants.RouteVar.CLIENT_TOKEN] = new_token or token + # Rebuild from router_data to keep the session var in step with it. + if ( + session := SessionData.from_router_data(state.router_data) + ) != state.rx_router_session: + state.rx_router_session = session + + +class WebsocketEventNamespace(BaseEventNamespace): + """Default event transport over a plain WebSocket. + + Frames are JSON arrays ``[event_name, payload]``. + """ + + supports_channels = True + + def __init__(self, namespace: str, app: App): + """Initialize the websocket event namespace. + + Args: + namespace: The namespace. + app: The application object. + """ + super().__init__(namespace, app) + self._connections: dict[str, _Connection] = {} + # Open channel sessions per connection, by session id and channel name. + self._channel_sessions: dict[str, dict[str, ChannelSession]] = {} + + def _deliver(self, sids: Iterable[str], payload: str | bytes, label: str) -> bool: + """Queue one serialized frame for connected client sessions. + + Args: + sids: The session ids to send to. + payload: The serialized text or binary frame. + label: The message name, for diagnostics. + + Returns: + Whether any session's connection queued the frame. + """ + size = utf8_size(payload) if otel.enabled else 0 + queued = False + for sid in sids: + if (connection := self._connections.get(sid)) is None: + # Routine race: the client disconnected while an event was + # still being processed, so its remaining updates have nowhere + # to go. + logger.debug(f"Attempted to emit {label!r} to unknown session {sid!r}.") + continue + if otel.enabled: + otel.record_message_size(size, "transmit") + if connection.send(payload): + queued = True + return queued + + async def emit(self, event: str, data: Any = None, to: str | None = None) -> None: + """Emit an event to a connected client session, or to all of them. + + Args: + event: The event name. + data: The event payload. + to: The session id to emit to; every connected session if None, + like Socket.IO's emit without a recipient. + """ + self._deliver( + self._connections if to is None else (to,), _dumps([event, data]), event + ) + + async def _send_channel_message( + self, + sids: Sequence[str], + channel: str, + event: str, + data: Any, + buffers: Sequence[Buffer], + ) -> bool: + """Send one channel message to connected client sessions. + + The frame is serialized once for every recipient: a room broadcast + costs one encode and one buffer, not one per member. + + Args: + sids: The session ids to send to. + channel: The channel name. + event: The message name. + data: The JSON-serializable metadata. + buffers: The binary attachments. + + Returns: + Whether a session still linked to its token queued the frame. + """ + payload = ( + encode_channel_frame(event, data, channel, buffers) + if buffers + else _dumps([event, data, channel]) + ) + linked = self._token_manager.sid_to_token + recipients = [] + for sid in sids: + if sid in linked: + recipients.append(sid) + elif (connection := self._connections.get(sid)) is not None: + # Its token moved to another socket, which drops the link but + # not the connection. Close it, as its next frame would; the + # teardown forgets its channel sessions. + connection.close(1008) + return self._deliver(recipients, payload, event) + + async def _send_channel_error( + self, sid: str, channel: str, code: str, message: str + ) -> None: + """Report a channel-level failure to a client session. + + Args: + sid: The session id. + channel: The channel name. + code: The machine-readable error code. + message: The human-readable explanation. + """ + await self._send_channel_message( + (sid,), + channel, + CHANNEL_ERROR_MESSAGE, + {"code": code, "message": message}, + (), + ) + + async def _open_channel_session(self, sid: str, channel_name: str) -> None: + """Open a channel session for a connection, answering the client. + + Args: + sid: The session id. + channel_name: The channel the client is opening. + """ + if channel_name in self._channel_sessions.get(sid, ()): + # Opening twice would orphan the first session in its rooms; the + # client only opens once per connection, so answer and move on. + await self._send_channel_message( + (sid,), channel_name, OPENED_MESSAGE, None, () + ) + return + channel = self.app._channels.get(channel_name) + if channel is None: + await self._send_channel_error( + sid, + channel_name, + "unknown_channel", + f"No channel named {channel_name!r} is registered.", + ) + return + token = self.sid_to_token.get(sid) + if token is None: + # The token was unlinked while the frame was in flight. + logger.debug(f"Ignoring channel open from session {sid} with no token.") + return + session = channel.open_session(sid, token, self._send_channel_message) + # Track the session before the hook runs: on_open may join rooms, and + # if it is interrupted the disconnect cleanup must still find it. + sessions = self._channel_sessions.setdefault(sid, {}) + sessions[channel_name] = session + try: + await channel.on_open(session) + except Exception as exc: + self._drop_channel_session(sid, channel_name) + self._log_handler_failure( + sid, f"Error opening channel {channel_name!r} for session {sid}.", exc + ) + await self._send_channel_error( + sid, channel_name, "open_failed", "The channel failed to open." + ) + return + await self._send_channel_message((sid,), channel_name, OPENED_MESSAGE, None, ()) + + async def _close_channel_session(self, sid: str, channel_name: str) -> None: + """Close one open channel session. + + Args: + sid: The session id. + channel_name: The channel to close. + """ + session = self._drop_channel_session(sid, channel_name) + if session is not None: + await self._notify_channel_close(sid, channel_name, session) + + def _drop_channel_session( + self, sid: str, channel_name: str + ) -> ChannelSession | None: + """Remove one session from the connection, along with its rooms. + + Args: + sid: The session id. + channel_name: The channel to drop. + + Returns: + The dropped session, or None if it was not open. + """ + sessions = self._channel_sessions.get(sid) + session = sessions.pop(channel_name, None) if sessions is not None else None + if session is None: + return None + if not sessions: + del self._channel_sessions[sid] + session.channel.forget_session(session) + return session + + async def _close_channel_sessions(self, sid: str) -> None: + """Close every channel session of a disconnected connection. + + Args: + sid: The session id. + """ + sessions = self._channel_sessions.pop(sid, None) + if not sessions: + return + # Drop rooms and tracking for all of them first: this runs during + # teardown, where a cancellation at the first await would otherwise + # leave the remaining sessions reachable by fan-out forever. + for session in sessions.values(): + session.channel.forget_session(session) + for channel_name, session in sessions.items(): + await self._notify_channel_close(sid, channel_name, session) + + async def _notify_channel_close( + self, sid: str, channel_name: str, session: ChannelSession + ) -> None: + """Run a channel's close hook, logging a failure instead of raising. + + Args: + sid: The session id. + channel_name: The channel being closed. + session: The session being closed. + """ + try: + await session.channel.on_close(session) + except Exception as exc: + self._log_handler_failure( + sid, f"Error closing channel {channel_name!r} for session {sid}.", exc + ) + + async def _handle_channel_message( + self, sid: str, channel_name: str, event: str, data: Any, buffers: list[bytes] + ) -> None: + """Dispatch one inbound channel frame. + + Never raises: a channel failing is a bug in that channel, not a reason + to drop the app's connection. + + Args: + sid: The session id. + channel_name: The channel the frame is addressed to. + event: The message name. + data: The message metadata. + buffers: The binary attachments. + """ + try: + if event == OPEN_MESSAGE: + await self._open_channel_session(sid, channel_name) + return + sessions = self._channel_sessions.get(sid) + session = sessions.get(channel_name) if sessions is not None else None + if session is None: + await self._send_channel_error( + sid, + channel_name, + "channel_not_open", + f"Channel {channel_name!r} is not open on this connection.", + ) + return + if event == CLOSE_MESSAGE: + await self._close_channel_session(sid, channel_name) + return + if is_reserved_event(event): + # The reservation holds in both directions, so a handler that + # relays what it receives cannot be made to attempt a send the + # channel API refuses. + await self._send_channel_error( + sid, + channel_name, + "reserved_event", + f"Channel message name {event!r} is reserved.", + ) + return + if buffers and not session.channel.accepts_binary: + await self._send_channel_error( + sid, + channel_name, + "binary_not_accepted", + f"Channel {channel_name!r} does not accept binary attachments.", + ) + return + await session.channel.on_message(session, event, data, buffers) + except Exception as exc: + self._log_handler_failure( + sid, + f"Error handling {event!r} on channel {channel_name!r} " + f"for session {sid}.", + exc, + ) + + async def _handle_binary_frame( + self, sid: str, frame: bytes, max_size: int + ) -> int | None: + """Validate and dispatch one inbound binary channel frame. + + Args: + sid: The session id. + frame: The raw frame bytes. + max_size: The message size limit in bytes. + + Returns: + The websocket close code the session must end with, or None to + keep serving it. + """ + if (close_code := self._accept_inbound(sid, frame, max_size)) is not None: + return close_code + try: + event, data, channel_name, buffers = decode_channel_frame(frame) + except ValueError: + # A Reflex client never sends malformed frames; close instead of + # logging per frame. + return _closing(sid, "malformed binary frame", 1002) + await self._handle_channel_message(sid, channel_name, event, data, buffers) + return None + + @staticmethod + async def _close_quietly( + websocket: WebSocket, code: int, reason: str | None = None + ) -> None: + """Close a connection, tolerating one that is already closed. + + The client can close its side while the server decides to close, and + starlette refuses a close after the connection has ended. + + Args: + websocket: The client websocket connection. + code: The close code. + reason: The close reason, at most 123 UTF-8 bytes. + """ + try: + await websocket.close(code=code, reason=reason) + except RuntimeError: + logger.debug("Connection was already closed.", exc_info=True) + + @staticmethod + def _accept_inbound(sid: str, payload: str | bytes, max_size: int) -> int | None: + """Apply the message size policy to a received frame, and account for it. + + ASGI delivers complete messages, so the server has already buffered + the frame; its protocol-level caps (enforced during frame reassembly) + bound that allocation. This applies the Reflex policy limit on top. + + Args: + sid: The session id. + payload: The received frame, text or binary. + max_size: The message size limit in bytes. + + Returns: + The close code the session must end with, or None to dispatch it. + """ + if exceeds_message_limit(payload, max_size): + return _closing(sid, f"message over {max_size} bytes", 1009) + if otel.enabled: + otel.record_message_size(utf8_size(payload), "receive") + return None + + @staticmethod + def _origin_allowed(origin: str | None) -> bool: + """Check a connection's Origin header against the CORS config. + + Args: + origin: The Origin header value, if any. + + Returns: + Whether the connection is allowed. + """ + if origin is None: + # Non-browser clients don't send an Origin header. + return True + allowed_origins = get_config().cors_allowed_origins + return "*" in allowed_origins or origin in allowed_origins + + async def _handle_frame( + self, sid: str, text: str, scope: MutableMapping[str, Any], max_size: int + ) -> int | None: + """Validate and dispatch one inbound text frame. + + Args: + sid: The session id. + text: The raw frame text. + scope: The ASGI scope of the client connection. + max_size: The message size limit in bytes. + + Returns: + The websocket close code the session must end with, or None to + keep serving it. + """ + if (close_code := self._accept_inbound(sid, text, max_size)) is not None: + return close_code + message = _parse_frame(text) + if ( + not isinstance(message, list) + or not message + or not isinstance(message[0], str) + ): + # A Reflex client never sends malformed frames; close instead of + # logging per frame, which a hostile client could use to flood + # the logs. + return _closing(sid, "malformed frame", 1002) + event = message[0] + data = message[1] if len(message) > 1 else None + if len(message) > 2: + if not isinstance(message[2], str): + return _closing(sid, "malformed channel frame", 1002) + await self._handle_channel_message(sid, message[2], event, data, []) + return None + try: + # Ordered by frequency: events are the hot path, heartbeat pongs + # arrive once per ping interval. + if event == _EVENT: + await self.handle_event(sid, data, scope) + elif event == PONG_MESSAGE: + # Receiving it already refreshed the liveness deadline. + pass + elif event == _PING: + await self.handle_ping(sid) + elif event == _CLIENT_ERROR: + await self.handle_client_error(sid, data) + else: + logger.debug( + f"Ignoring unknown socket event {event!r} from session {sid}." + ) + except Exception as exc: + return self._handler_failed(sid, f"socket event {event!r}", exc) + return None + + async def _receive_connect( + self, sid: str, websocket: WebSocket, timeout: float, max_size: int + ) -> list[Any] | None: + """Read the frame that opens a session, closing the socket on anything else. + + Args: + sid: The session id. + websocket: The client websocket connection. + timeout: Seconds to wait for the frame. + max_size: The message size limit in bytes. + + Returns: + The connect frame, or None once the socket is closed. + """ + try: + received = await asyncio.wait_for(websocket.receive(), timeout) + except TimeoutError: + await self._close_quietly( + websocket, _closing(sid, "no connect frame", 1008) + ) + return None + if received["type"] == "websocket.disconnect": + return None + if (text := received.get("text")) is not None: + if (close_code := self._accept_inbound(sid, text, max_size)) is not None: + # The client reports the oversized hydrate this carries, with + # the limit named here. + await self._close_quietly( + websocket, close_code, f"{MAX_MESSAGE_SIZE_REASON}{max_size}" + ) + return None + message = _parse_frame(text) + if isinstance(message, list) and message and message[0] == CONNECT_MESSAGE: + return message + # A Reflex client opens every session with a connect frame. + await self._close_quietly( + websocket, _closing(sid, "expected a connect frame", 1002) + ) + return None + + async def handle_websocket(self, websocket: WebSocket) -> None: + """Serve one client websocket connection for its full lifetime. + + Args: + websocket: The client websocket connection. + """ + if not self._origin_allowed(websocket.headers.get("origin")): + # Reject cross-origin connections before accepting. + await websocket.close(code=1008) + return + if b"EIO" in urllib.parse.parse_qs(websocket.scope.get("query_string", b"")): + # A Socket.IO client -- a tab or bundle from before this transport + # -- cannot speak it; refuse it at once rather than leave it + # waiting for an engine.io handshake. + await websocket.close(code=1008) + return + subprotocols = websocket.scope.get("subprotocols") or [] + # Echo the client's offered subprotocol (the Reflex version); browsers + # abort the connection if the server selects none. + await websocket.accept(subprotocol=subprotocols[0] if subprotocols else None) + + sid = str(uuid.uuid4()) + policy = _socket_policy() + # The client sends its connect frame on open, without waiting for the + # handshake, so the boot event it carries saves that round trip. A + # socket that never sends one ends within the heartbeat window. + connect_frame = await self._receive_connect( + sid, websocket, policy.silence_limit, policy.max_message_size + ) + if connect_frame is None: + return + connection = _Connection(websocket) + self._connections[sid] = connection + heartbeat_task = asyncio.create_task( + self._heartbeat(connection, policy), + name=f"reflex_heartbeat|{sid}", + ) + lifetime = asyncio.timeout(None) + try: + async with lifetime: + connection.attach(lifetime) + await self._serve_session( + sid, + connection, + connect_frame[1] if len(connect_frame) > 1 else None, + subprotocols[0] if subprotocols else None, + policy, + ) + except TimeoutError: + # The lifetime firing is how connection.close() ends the loop. + if not lifetime.expired(): + raise + except WebSocketDisconnect: + pass + finally: + heartbeat_task.cancel() + connection.stop() + self._connections.pop(sid, None) + # Start the token cleanup before any teardown await: a cancelled + # shutdown must not leave the token linked to a dead session. + cleanup_task = self.handle_disconnect(sid) + await self._close_channel_sessions(sid) + # A close hook that raised logged through the handler budget, + # which recreated the counter handle_disconnect had just dropped. + self._handler_error_budget.forget(sid) + if cleanup_task is not None: + # Await the token cleanup so the client's reconnect after the + # close is not treated as a duplicate tab; shielded so + # cancellation (e.g. server shutdown) cannot abort it. Errors + # are logged by the task's done callback. + with contextlib.suppress(Exception): + await asyncio.shield(cleanup_task) + if connection.close_code is not None: + await self._close_quietly(websocket, connection.close_code) + + @staticmethod + async def _heartbeat(connection: _Connection, policy: _SocketPolicy) -> None: + """Ping a session's client, closing the session once it stops answering. + + Args: + connection: The client connection. + policy: The socket's heartbeat settings. + """ + while True: + await asyncio.sleep(policy.ping_interval) + if time.monotonic() - connection.last_received > policy.silence_limit: + connection.close(1001) + return + connection.send(_PING_FRAME) + + async def _serve_session( + self, + sid: str, + connection: _Connection, + auth: Any, + subprotocol: str | None, + policy: _SocketPolicy, + ) -> None: + """Open a session and dispatch its frames until it ends. + + Args: + sid: The session id. + connection: The client connection. + auth: The payload its connect frame carried, if any. + subprotocol: The websocket subprotocol offered by the client. + policy: The socket's heartbeat and message size settings. + """ + websocket = connection.websocket + linked, boot_event = await self._open_session( + sid, auth, websocket.scope.get("query_string", b"").decode(), subprotocol + ) + if not linked: + # No token was linked; not a Reflex client. + connection.close(1008) + return + # The handshake acknowledges the connect and carries the heartbeat + # settings for the client's connection watchdog. + connection.send( + _dumps([ + HANDSHAKE_MESSAGE, + { + "ping_interval": policy.ping_interval, + "ping_timeout": policy.ping_timeout, + "protocol": PROTOCOL_VERSION, + # So a client can refuse an oversized frame itself + # rather than lose the connection to one. + "max_message_size": policy.max_message_size, + }, + ]) + ) + if ( + boot_event is not None + and ( + close_code := await self._handle_boot_event( + sid, self.handle_event(sid, boot_event, websocket.scope) + ) + ) + is not None + ): + connection.close(close_code) + return + while connection.close_code is None: + received = await websocket.receive() + if received["type"] == "websocket.disconnect": + return + connection.last_received = time.monotonic() + if sid not in self._token_manager.sid_to_token: + # The token moved to another socket or its record went stale. + # Nothing this session sends can be served -- not events, not + # channel messages, which would otherwise keep invoking + # handlers under a token that has moved on -- and a reconnect + # is how it gets a working session back. + connection.close(_closing(sid, "its token is gone", 1008)) + return + text = received.get("text") + if text is not None: + close_code = await self._handle_frame( + sid, text, websocket.scope, policy.max_message_size + ) + elif (frame := received.get("bytes")) is not None and self.app._channels: + close_code = await self._handle_binary_frame( + sid, frame, policy.max_message_size + ) + else: + # Binary frame with no channel to carry it. + close_code = _closing(sid, "received a binary frame", 1003) + if close_code is not None: + connection.close(close_code) diff --git a/reflex/socketio_namespace.py b/reflex/socketio_namespace.py new file mode 100644 index 00000000000..8ab206bd07f --- /dev/null +++ b/reflex/socketio_namespace.py @@ -0,0 +1,323 @@ +"""Socket.IO event transport (requires the optional python-socketio dependency).""" + +from __future__ import annotations + +import json +from collections.abc import Awaitable, MutableMapping +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +from reflex_base import otel +from reflex_base.utils.types import ASGIApp, Message, Receive, Scope, Send +from socketio import ASGIApp as EngineIOApp +from socketio import AsyncNamespace, AsyncServer + +from reflex.event_namespace import ( + _CLIENT_ERROR, + _EVENT, + _PING, + BaseEventNamespace, + _dumps, + _socket_policy, + utf8_size, +) + +if TYPE_CHECKING: + import asyncio + + from reflex_base.config import Config + + from reflex.app import App + + +def _sio_dumps(obj: Any, **kwargs: Any) -> str: + """Serialize an outgoing Socket.IO packet, recording its size when telemetry is on. + + Serialized like a plain transport frame: compact, with the non-finite + float tokens the frontend revives, and valid UTF-8. + + Args: + obj: The packet payload. + **kwargs: Encoder options; socket.io only asks for compact separators. + + Returns: + The JSON string. + """ + data = _dumps(obj) + if otel.enabled: + otel.record_message_size(utf8_size(data), "transmit") + return data + + +def _sio_loads(data: str | bytes, **kwargs: Any) -> Any: + """Deserialize an incoming Socket.IO packet, recording its size when telemetry is on. + + Args: + data: The JSON string. + **kwargs: Options forwarded to the JSON decoder. + + Returns: + The decoded payload. + """ + if otel.enabled: + otel.record_message_size(utf8_size(data), "receive") + return json.loads(data, **kwargs) + + +# The JSON codec socket.io serializes packets with. +_SOCKET_JSON_CODEC = SimpleNamespace( + dumps=staticmethod(_sio_dumps), + loads=staticmethod(_sio_loads), +) + + +class EventNamespace(AsyncNamespace, BaseEventNamespace): + """The Socket.IO event namespace.""" + + def __init__(self, namespace: str, app: App): + """Initialize the event namespace. + + Args: + namespace: The namespace. + app: The application object. + """ + AsyncNamespace.__init__(self, namespace) + BaseEventNamespace.__init__(self, namespace, app) + # ASGI scope per session, so the per-event path reaches the connection + # state it caches without going through the Socket.IO server. + self._scopes: dict[str, MutableMapping[str, Any]] = {} + # Sessions whose boot event on_event is handling, which raises its failure. + self._booting: set[str] = set() + + async def on_connect( + self, sid: str, environ: dict, auth: Any = None + ) -> bool | None: + """Event for when the websocket is connected. + + Args: + sid: The Socket.IO session id. + environ: The request information, including HTTP headers. + auth: The payload of the CONNECT packet, carrying the boot event + and the frontend state schema. + + Returns: + False to refuse a session that linked no token or whose boot event + failed, else None. + """ + if (scope := environ.get("asgi.scope")) is not None: + self._scopes[sid] = scope + linked, boot_event = await self._open_session( + sid, + auth, + environ.get("QUERY_STRING", ""), + environ.get("HTTP_SEC_WEBSOCKET_PROTOCOL"), + ) + # A failed boot event refuses the connect rather than raising, which + # would keep the connect registered yet unanswered, or disconnecting, + # which would acknowledge a session that never hydrated. + if linked and (boot_event is None or await self._boot(sid, boot_event) is None): + return None + # Socket.IO runs no disconnect handler for a refused connect, so undo + # this one here. + self.on_disconnect(sid) + return False + + async def _boot(self, sid: str, boot_event: Any) -> int | None: + """Dispatch a connect's boot event through on_event, so its overrides see it. + + Args: + sid: The Socket.IO session id. + boot_event: The boot event. + + Returns: + None once dispatched, else the close code the session failed with. + """ + # Marked as booting, on_event raises a failure for the boot event's + # policy instead of applying the per-event one. + self._booting.add(sid) + try: + return await self._handle_boot_event(sid, self.on_event(sid, boot_event)) + finally: + self._booting.discard(sid) + + def on_disconnect(self, sid: str) -> asyncio.Task | None: + """Event for when the websocket disconnects. + + Args: + sid: The Socket.IO session id. + + Returns: + An asyncio Task for cleaning up the token, or None. + """ + self._scopes.pop(sid, None) + return self.handle_disconnect(sid) + + async def on_event(self, sid: str, data: Any): + """Event for receiving front-end websocket events. + + A failing boot event raises, for on_connect to refuse the connect. + + Args: + sid: The Socket.IO session id. + data: The event data. + """ + # Inline rather than through _run_handler, which costs a coroutine + # frame: this runs for every event. + try: + await self.handle_event(sid, data, self._scope(sid)) + except Exception as exc: + if sid in self._booting: + raise + if self._handler_failed(sid, f"socket event {_EVENT!r}", exc): + await self.disconnect(sid) + + async def _run_handler(self, sid: str, event: str, handler: Awaitable[None]): + """Await a socket event handler under the shared failure policy. + + Raising into python-socketio would log a traceback per frame, outside + the budget that bounds client-triggered error logs. + + Args: + sid: The Socket.IO session id. + event: The socket event name, for the log. + handler: The handler's coroutine. + """ + try: + await handler + except Exception as exc: + if self._handler_failed(sid, f"socket event {event!r}", exc): + await self.disconnect(sid) + + def _scope(self, sid: str) -> MutableMapping[str, Any]: + """Get the ASGI scope of a session's connection. + + A plain function, not a coroutine: it runs for every event. + + Args: + sid: The Socket.IO session id. + + Returns: + The ASGI scope. + + Raises: + RuntimeError: If the Socket.IO is badly initialized. + """ + if (scope := self._scopes.get(sid)) is not None: + return scope + # The connection was not seen by on_connect (e.g. the namespace was + # registered after the socket connected); ask the server. + if self.app.sio is None: + msg = "Socket.IO is not initialized." + raise RuntimeError(msg) + environ = self.app.sio.get_environ(sid, self.namespace) + if environ is None: + msg = "Socket.IO environ is not initialized." + raise RuntimeError(msg) + scope = self._scopes[sid] = environ["asgi.scope"] + return scope + + async def on_ping(self, sid: str): + """Event for testing the API endpoint. + + Args: + sid: The Socket.IO session id. + """ + await self._run_handler(sid, _PING, self.handle_ping(sid)) + + async def on_client_error(self, sid: str, data: Any = None): + """Handle errors reported by the frontend. + + Args: + sid: The Socket.IO session id. + data: The error data from the client. Defaults to None because + python-socketio dispatches a payload-less emit as + ``on_client_error(sid)``; the malformed-payload guard then + drops it without raising. + """ + await self._run_handler(sid, _CLIENT_ERROR, self.handle_client_error(sid, data)) + + +class _HeaderMiddleware: + """Echo the websocket subprotocol on accept, which engineio does not.""" + + def __init__(self, app: ASGIApp): + """Initialize the middleware. + + Args: + app: The ASGI app to wrap. + """ + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send): + """Handle an ASGI connection. + + Args: + scope: The ASGI scope. + receive: The ASGI receive callable. + send: The ASGI send callable. + + Returns: + The result of the wrapped app. + """ + original_send = send + + async def modified_send(message: Message): + if message["type"] == "websocket.accept": + if scope.get("subprotocols"): + # The following *does* say "subprotocol" instead of "subprotocols", intentionally. + message["subprotocol"] = scope["subprotocols"][0] + + headers = dict(message.get("headers", [])) + header_key = b"sec-websocket-protocol" + if subprotocol := headers.get(header_key): + message["headers"] = [ + *message.get("headers", []), + (header_key, subprotocol), + ] + + return await original_send(message) + + return await self.app(scope, receive, modified_send) + + +def create_socketio_app(app: App, config: Config) -> ASGIApp: + """Create the Socket.IO server for an app and return its ASGI app. + + Creates ``app.sio`` if the user did not supply their own server. + + Args: + app: The Reflex app. + config: The app configuration. + + Returns: + The ASGI app serving the Socket.IO server. + + Raises: + RuntimeError: If a custom ``sio`` server does not use asgi mode. + """ + if not app.sio: + polling = config.transport == "polling" + policy = _socket_policy() + app.sio = AsyncServer( + async_mode="asgi", + # An empty list disables CORS handling; engineio reads "*" among + # the origins as allowing any. + cors_allowed_origins=[] if polling else list(config.cors_allowed_origins), + cors_credentials=not polling, + max_http_buffer_size=policy.max_message_size, + ping_interval=policy.ping_interval, + ping_timeout=policy.ping_timeout, + json=_SOCKET_JSON_CODEC, + allow_upgrades=False, + transports=["polling" if polling else "websocket"], + # Handlers here only parse and enqueue (or emit a pong), so run + # them inline on the socket's receive loop instead of paying a + # task creation and a loop hop per incoming message. + async_handlers=False, + ) + elif getattr(app.sio, "async_mode", "") != "asgi": + msg = f"Custom `sio` must use `async_mode='asgi'`, not '{app.sio.async_mode}'." + raise RuntimeError(msg) + + # Create the socket app. Note event endpoint constant replaces the default 'socket.io' path. + return _HeaderMiddleware(EngineIOApp(app.sio, socketio_path="")) diff --git a/reflex/state.py b/reflex/state.py index 9859e0d3c2a..a5b0d628231 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -3009,7 +3009,7 @@ class StateUpdate: Substates without a registered frontend dispatch function are skipped and reported to the backend through ``client_error`` (see - ``EventNamespace.on_client_error``). + ``BaseEventNamespace.handle_client_error``). """ # The state delta. diff --git a/reflex/testing.py b/reflex/testing.py index c80133231c5..5b605eb26b6 100644 --- a/reflex/testing.py +++ b/reflex/testing.py @@ -20,16 +20,22 @@ import threading import time import types -from collections.abc import Callable, Coroutine, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Sequence from importlib.util import find_spec from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, Self, TypeVar +from granian.constants import Interfaces +from granian.errors import FatalError +from granian.log import LogLevels +from granian.net import SocketHolder # pyright: ignore[reportPrivateImportUsage] +from granian.server.embed import Server as EmbeddedGranian from reflex_base.components.memo import MEMOS from reflex_base.config import get_config, reload_config from reflex_base.environment import environment from reflex_base.registry import RegistrationContext -from reflex_base.utils.types import ASGIApp +from reflex_base.utils import console +from reflex_base.utils.types import ASGIApp, Message, Receive, Scope, Send import reflex import reflex.reflex @@ -59,9 +65,6 @@ except ImportError: has_selenium = False -if TYPE_CHECKING: - import uvicorn - # The timeout (minutes) to check for the port. DEFAULT_TIMEOUT = 15 POLL_INTERVAL = 0.25 @@ -70,23 +73,6 @@ TimeoutType = int | float | None -def _get_uvicorn(): - """Import uvicorn for an AppHarness server. - - Returns: - The imported uvicorn module. - """ - try: - import uvicorn - except ImportError as exc: - msg = ( - "AppHarness backend support requires `uvicorn`. Install it with " - "`pip install 'reflex[testing]'`." - ) - raise ImportError(msg) from exc - return uvicorn - - if platform.system() == "Windows": FRONTEND_POPEN_ARGS["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP # pyright: ignore [reportAttributeAccessIssue] FRONTEND_POPEN_ARGS["shell"] = True @@ -94,6 +80,262 @@ def _get_uvicorn(): FRONTEND_POPEN_ARGS["start_new_session"] = True +class _PreboundGranian(EmbeddedGranian): + """Embedded granian server serving a listening socket bound beforehand.""" + + def __init__(self, listener: socket.socket, *args, **kwargs) -> None: + """Create the server around a bound socket. + + Args: + listener: the listening socket; granian owns it once serving starts. + *args: positional arguments for the granian server. + **kwargs: keyword arguments for the granian server. + """ + super().__init__(*args, **kwargs) + self._listener = listener + + def _init_shared_socket(self): + """Hand the listening socket to granian instead of binding a new one.""" + fd = self._listener.detach() + # SocketHolder takes its pickled state, whose shape differs per platform. + if sys.platform == "win32": + state = (fd,) + elif sys.platform.startswith(("linux", "freebsd")): + state = (fd, False, self.backlog) + else: + state = (fd, False) + self._shd = SocketHolder(*state) + self._sfd = fd + + def release_socket(self, served: bool) -> None: + """Close the socket unless the worker serving it already did. + + Args: + served: whether a worker took the socket over. + """ + holder, self._shd = self._shd, None + # Dropping the holder closes it on Windows, where workers serve a clone. + if holder is not None and not served and sys.platform != "win32": + os.close(self._sfd) + + +class _EmbeddedServer: + """In-process granian server with a uvicorn-like control surface. + + Serves the given ASGI app object directly, so the harness shares the app + and state instances with the running server. The listening socket is + bound on construction, without SO_REUSEPORT, so the server owns its port + (OS-assigned for port 0) before it serves. + """ + + def __init__(self, app: ASGIApp, host: str = "127.0.0.1", port: int = 0) -> None: + """Bind the listening socket without serving yet. + + Args: + app: the ASGI app object to serve. + host: the address to bind to. + port: the port to bind to; 0 lets the OS pick a free one. + """ + self.app = app + self.host = host + self._listener = self._listen(host, port) + self.port: int = self._listener.getsockname()[1] + # Monkeypatchable async shutdown hook, mirroring uvicorn.Server.shutdown. + self.shutdown: Callable[..., Coroutine[Any, Any, None]] = self._noop_shutdown + self._should_exit = threading.Event() + self._serving = False + # Set once the server serves or stops. + self._settled = threading.Event() + self._loop: asyncio.AbstractEventLoop | None = None + self._server: _PreboundGranian | None = None + + @staticmethod + def _listen(host: str, port: int) -> socket.socket: + """Bind a listening socket that no other socket can share. + + Args: + host: the address to bind to. + port: the port to bind to. + + Returns: + The listening socket. + """ + sock = socket.socket() + try: + # On posix SO_REUSEADDR only allows rebinding a port in TIME_WAIT; + # on Windows it would let other sockets share the port. + sock.setsockopt( + socket.SOL_SOCKET, + getattr(socket, "SO_EXCLUSIVEADDRUSE", socket.SO_REUSEADDR), + 1, + ) + # Accepted sockets inherit it, as from the listeners granian binds: + # with Nagle, a small frame waits for the client's delayed ACK. + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + sock.bind((host, port)) + sock.listen() + except OSError: + sock.close() + raise + return sock + + @staticmethod + async def _noop_shutdown(*args, **kwargs) -> None: + """Default shutdown hook. + + Args: + *args: ignored. + **kwargs: ignored. + """ + + def getsockname(self) -> tuple[str, int]: + """The address the server is bound to. + + Returns: + The (host, port) tuple the server serves on. + """ + return (self.host, self.port) + + @property + def url(self) -> str: + """The http URL the server serves on. + + Returns: + The URL, without a trailing slash. + """ + return f"http://{self.host}:{self.port}" + + def wait_started(self, timeout: float | None = None) -> None: + """Block until the server serves, after the app's lifespan startup. + + Args: + timeout: how long to wait in seconds; None waits indefinitely. + + Raises: + TimeoutError: when the server does not serve within the timeout. + RuntimeError: when the server stopped without serving. + """ + if not self._settled.wait(timeout): + msg = f"Server on port {self.port} did not start within {timeout}s." + raise TimeoutError(msg) + if not self._serving: + msg = f"Server on port {self.port} stopped without serving; see the log above." + raise RuntimeError(msg) + + @property + def started(self) -> bool: + """Whether the server serves, like uvicorn's `Server.started`. + + Returns: + True once the app's lifespan startup completed. + """ + console.deprecate( + feature_name="AppHarness.backend.started", + reason="AppHarness.start() returns once the backend serves; " + "call `wait_started()` to wait for a server yourself", + deprecation_version="0.10.0", + removal_version="1.0", + ) + return self._serving + + @property + def should_exit(self) -> bool: + """Whether the server was asked to stop. + + Returns: + True after `should_exit` has been set. + """ + return self._should_exit.is_set() + + @should_exit.setter + def should_exit(self, value: bool) -> None: + if not value: + return + self._should_exit.set() + loop, server = self._loop, self._server + if loop is not None and server is not None: + # A closed loop means the server is already down. + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(server.stop) + + def run(self) -> None: + """Serve the app until `should_exit` is set; used as a thread target.""" + try: + asyncio.run(self._serve()) + finally: + # A no-op once granian took the socket over. + self._listener.close() + self._settled.set() + + def _mark_serving(self) -> None: + """Record that the server serves.""" + self._serving = True + self._settled.set() + + def _asgi(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: + """Call the app, tracking its lifespan startup. + + Args: + scope: the ASGI scope. + receive: the ASGI receive callable. + send: the ASGI send callable. + + Returns: + The app's ASGI awaitable. + """ + if scope["type"] == "lifespan": + return self._lifespan(scope, receive, send) + return self.app(scope, receive, send) + + async def _lifespan(self, scope: Scope, receive: Receive, send: Send) -> None: + """Run the app's lifespan, marking the server serving after startup. + + Args: + scope: the ASGI lifespan scope. + receive: the ASGI receive callable. + send: the ASGI send callable. + """ + failed = False + + async def tracking_send(message: Message) -> None: + nonlocal failed + await send(message) + if message["type"] == "lifespan.startup.complete": + self._mark_serving() + elif message["type"] == "lifespan.startup.failed": + failed = True + + try: + await self.app(scope, receive, tracking_send) + finally: + # Like granian, serve without lifespan support unless startup failed. + if not failed: + self._mark_serving() + + async def _serve(self) -> None: + server = _PreboundGranian( + self._listener, + self._asgi, + address=self.host, + port=self.port, + interface=Interfaces.ASGI, + # Keeps lifespan failures and unhandled app errors visible. + log_level=LogLevels.error, + ) + self._server = server + self._loop = asyncio.get_running_loop() + try: + # A stop requested before this point found no server to interrupt. + # A failed lifespan startup is logged by granian and reported by + # wait_started(). + if not self._should_exit.is_set(): + with contextlib.suppress(FatalError): + await server.serve() + finally: + server.release_socket(served=self._serving) + await self.shutdown() + + # borrowed from py3.11 class chdir(contextlib.AbstractContextManager): # noqa: N801 """Non thread-safe context manager to change the current working directory.""" @@ -138,7 +380,7 @@ class AppHarness: frontend_url: str | None = None frontend_output_thread: threading.Thread | None = None backend_thread: threading.Thread | None = None - backend: uvicorn.Server | None = None + backend: _EmbeddedServer | None = None _frontends: list[WebDriver] = dataclasses.field(default_factory=list) _registry_token: contextvars.Token[RegistrationContext] | None = None _base_registration_context: ClassVar[RegistrationContext] | None = None @@ -368,14 +610,9 @@ def _start_backend(self, port: int = 0): if self.app_asgi is None: msg = "App was not initialized." raise RuntimeError(msg) - uvicorn = _get_uvicorn() - self.backend = uvicorn.Server( - uvicorn.Config( - app=self.app_asgi, - host="127.0.0.1", - port=port, - ) - ) + # Granian serves it, so it advertises the frame size Granian receives. + reflex.utils.exec._cap_message_size_for_granian() + self.backend = _EmbeddedServer(self.app_asgi, port=port) self.backend.shutdown = self._get_backend_shutdown_handler() def _run_backend(context: contextvars.Context) -> None: @@ -397,9 +634,7 @@ def _start_frontend(self): with chdir(self.app_path): config = get_config() print("Polling for servers...") # for pytest diagnosis #noqa: T201 - config.api_url = "http://{}:{}".format( - *self._poll_for_servers(timeout=30).getsockname(), - ) + config.api_url = self._poll_for_servers(timeout=30).url print("Building frontend...") # for pytest diagnosis #noqa: T201 reflex.utils.build.setup_frontend(self.app_path) @@ -610,38 +845,25 @@ async def _poll_for_async( await asyncio.sleep(step) return False - def _poll_for_servers(self, timeout: TimeoutType = None) -> socket.socket: - """Poll backend server for listening sockets. + def _poll_for_servers(self, timeout: TimeoutType = None) -> _EmbeddedServer: + """Wait for the backend server to serve. Args: - timeout: how long to wait for listening socket. + timeout: how long to wait for the server. Returns: - first active listening socket on the backend + the backend server, exposing `getsockname()` for its bound address Raises: - RuntimeError: when the backend hasn't started running - TimeoutError: when server or sockets are not ready + RuntimeError: when the backend hasn't started running, or stopped + without serving + TimeoutError: when the server is not ready """ if self.backend is None: msg = "Backend is not running." raise RuntimeError(msg) - backend = self.backend - # check for servers to be initialized - if not self._poll_for( - target=lambda: getattr(backend, "servers", False), - timeout=timeout, - ): - msg = "Backend servers are not initialized." - raise TimeoutError(msg) - # check for sockets to be listening - if not self._poll_for( - target=lambda: getattr(backend.servers[0], "sockets", False), - timeout=timeout, - ): - msg = "Backend is not listening." - raise TimeoutError(msg) - return backend.servers[0].sockets[0] + self.backend.wait_started(DEFAULT_TIMEOUT if timeout is None else timeout) + return self.backend def frontend( self, @@ -851,34 +1073,18 @@ class AppHarnessProd(AppHarness): """AppHarnessProd executes a reflex app in-process for testing. In prod mode, instead of running `react-router dev` the app is exported as static - files and served via Starlette StaticFiles in a dedicated Uvicorn server. - Additionally, the backend runs in multi-worker mode. + files and served via Starlette StaticFiles on a dedicated embedded server. """ frontend_thread: threading.Thread | None = None - frontend_server: uvicorn.Server | None = None - - def _run_frontend(self): - uvicorn = _get_uvicorn() - with chdir(self.app_path): - frontend_app = reflex.utils.exec._frontend_prod_app() - self.frontend_server = uvicorn.Server( - uvicorn.Config( - app=frontend_app, - host="127.0.0.1", - port=0, - ) - ) - self.frontend_server.run() + frontend_server: _EmbeddedServer | None = None def _start_frontend(self): # Set up the frontend. with chdir(self.app_path): config = get_config() print("Polling for servers...") # for pytest diagnosis #noqa: T201 - config.api_url = "http://{}:{}".format( - *self._poll_for_servers(timeout=30).getsockname(), - ) + config.api_url = self._poll_for_servers(timeout=30).url print("Building frontend...") # for pytest diagnosis #noqa: T201 get_config().loglevel = reflex.constants.LogLevel.INFO @@ -895,64 +1101,29 @@ def _start_frontend(self): loglevel=reflex.constants.LogLevel.INFO, env=reflex.constants.Env.PROD, ) + self.frontend_server = _EmbeddedServer( + reflex.utils.exec._frontend_prod_app() + ) print("Frontend starting...") # for pytest diagnosis #noqa: T201 - self.frontend_thread = threading.Thread(target=self._run_frontend) + self.frontend_thread = threading.Thread(target=self.frontend_server.run) self.frontend_thread.start() def _wait_frontend(self): - self._poll_for( - lambda: ( - self.frontend_server is not None - and getattr(self.frontend_server, "servers", []) - and self.frontend_server.servers[0].sockets - ) - ) - if ( - self.frontend_server is None - or not self.frontend_server.servers[0].sockets - or not self.frontend_server.servers[0].sockets[0].fileno() - ): + if self.frontend_server is None: msg = "Frontend did not start" raise RuntimeError(msg) - frontend_socket = self.frontend_server.servers[0].sockets[0] + self.frontend_server.wait_started(DEFAULT_TIMEOUT) config = get_config() - self.frontend_url = "http://{}:{}".format( - *frontend_socket.getsockname() - ) + config.prepend_frontend_path("/") + self.frontend_url = self.frontend_server.url + config.prepend_frontend_path("/") config.deploy_url = self.frontend_url - def _start_backend(self): - if self.app_asgi is None: - msg = "App was not initialized." - raise RuntimeError(msg) + def _start_backend(self, port: int = 0): environment.REFLEX_SKIP_COMPILE.set(True) - uvicorn = _get_uvicorn() - self.backend = uvicorn.Server( - uvicorn.Config( - app=self.app_asgi, - host="127.0.0.1", - port=0, - workers=reflex.utils.processes.get_num_workers(), - ), - ) - self.backend.shutdown = self._get_backend_shutdown_handler() - - def _run_backend(context: contextvars.Context) -> None: - if self.backend is not None: - context.run(self.backend.run) - - print( # noqa: T201 - "Creating backend in a new thread..." - ) - self.backend_thread = threading.Thread( - target=_run_backend, args=(contextvars.copy_context(),) - ) - self.backend_thread.start() - print("Backend started.") # for pytest diagnosis #noqa: T201 + super()._start_backend(port) - def _poll_for_servers(self, timeout: TimeoutType = None) -> socket.socket: + def _poll_for_servers(self, timeout: TimeoutType = None) -> _EmbeddedServer: try: return super()._poll_for_servers(timeout) finally: diff --git a/reflex/utils/build.py b/reflex/utils/build.py index 45bd2895a71..c403e7f68bd 100644 --- a/reflex/utils/build.py +++ b/reflex/utils/build.py @@ -21,7 +21,11 @@ def set_env_json(): config = get_config() env: dict[str, object] = { **{endpoint.name: endpoint.get_url() for endpoint in constants.Endpoint}, - "TRANSPORT": config.transport, + # An engine.io transport name, as third-party code sharing the app's + # Socket.IO connection hands it to socket.io-client; SOCKETIO picks + # the protocol. + "TRANSPORT": "polling" if config.transport == "polling" else "websocket", + "SOCKETIO": config.transport != "websocket", "TEST_MODE": is_in_app_harness(), } for plugin in config.plugins: diff --git a/reflex/utils/exec.py b/reflex/utils/exec.py index 5c4a8fb9726..02bde49ad89 100644 --- a/reflex/utils/exec.py +++ b/reflex/utils/exec.py @@ -484,6 +484,29 @@ def run_frontend_prod(host: str, port: int): ) +def _uvicorn_can_serve_websockets() -> bool: + """Whether a websocket protocol library uvicorn can use is installed. + + Returns: + True if websockets or wsproto is importable. + """ + return ( + importlib.util.find_spec("websockets") is not None + or importlib.util.find_spec("wsproto") is not None + ) + + +@once +def _warn_about_uvicorn_websockets(): + """Warn when the selected uvicorn cannot serve websockets.""" + if not _uvicorn_can_serve_websockets(): + logger.warning( + "Uvicorn has no websocket protocol library installed, so the default " + "WebSocket transport will not connect. Install `reflex[uvicorn]` or " + "use Granian (REFLEX_USE_GRANIAN=1)." + ) + + @once def _warn_user_about_uvicorn(): logger.warning( @@ -498,10 +521,18 @@ def should_use_granian(): True if Granian should be used. """ if environment.REFLEX_USE_GRANIAN.is_set(): - return environment.REFLEX_USE_GRANIAN.get() + use_granian = environment.REFLEX_USE_GRANIAN.get() + if not use_granian: + # Asking for uvicorn explicitly is the likeliest way to end up + # without a websocket library, so this check cannot live on the + # auto-detect branch alone. + _warn_about_uvicorn_websockets() + return use_granian if ( importlib.util.find_spec("uvicorn") is None or importlib.util.find_spec("gunicorn") is None + # Granian is always installed and serves the event socket itself. + or not _uvicorn_can_serve_websockets() ): return True _warn_user_about_uvicorn() @@ -731,9 +762,55 @@ def run_uvicorn_backend(host: str, port: int, loglevel: LogLevel): reload=True, reload_dirs=list(map(str, get_reload_paths())), reload_delay=0.1, + **uvicorn_websocket_options(), ) +def _uvicorn_ws_max_size() -> int: + """Websocket message size limit for uvicorn. + + Never below uvicorn's 16 MiB default, so unrelated websocket endpoints + keep working; raised when the Reflex policy limit needs more. + + Returns: + The message size limit in bytes. + """ + return max(environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE.get(), 16 * 1024 * 1024) + + +def uvicorn_websocket_options() -> dict[str, Any]: + """The app's websocket policy as uvicorn settings. + + Every uvicorn launch path applies these, including the gunicorn worker + class, which is how the production server receives options gunicorn itself + does not forward. + + Returns: + The uvicorn configuration keyword arguments. + """ + return { + "ws_max_size": _uvicorn_ws_max_size(), + "ws_per_message_deflate": environment.REFLEX_SOCKET_PER_MESSAGE_DEFLATE.get(), + } + + +def _uvicorn_websocket_args() -> list[str]: + """The app's websocket policy as uvicorn command line arguments. + + Returns: + The command line arguments. + """ + options = uvicorn_websocket_options() + return [ + *("--ws-max-size", str(options["ws_max_size"])), + # A BOOLEAN-valued option, not a flag: uvicorn rejects --no-... forms. + *( + "--ws-per-message-deflate", + str(options["ws_per_message_deflate"]).lower(), + ), + ] + + HOTRELOAD_IGNORE_EXTENSIONS = ( "txt", "toml", @@ -769,6 +846,32 @@ def _granian_log_dictconfig() -> dict[str, Any] | None: return {"handlers": {"console": json_handler, "access": json_handler}} +# The largest websocket frame Granian receives; it drops larger ones, with no +# setting to raise the limit. +_GRANIAN_MAX_FRAME_SIZE = 16 * 1024 * 1024 + + +def _cap_message_size_for_granian(): + """Lower the message size limit to what Granian can receive. + + The limit is advertised to clients, which size what they send by it; a + message Granian drops is lost with the socket. Lowered in the environment, + so the workers it spawns read the capped value too. + """ + configured = environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE.get() + if configured <= _GRANIAN_MAX_FRAME_SIZE: + return + logger.warning( + f"REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE is {configured} bytes, but Granian " + f"drops websocket frames over {_GRANIAN_MAX_FRAME_SIZE} bytes; capping " + "it there. Run uvicorn (`reflex[uvicorn]`, REFLEX_USE_GRANIAN=0) for " + "larger messages." + ) + os.environ[environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE.name] = str( + _GRANIAN_MAX_FRAME_SIZE + ) + + def run_granian_backend(host: str, port: int, loglevel: LogLevel): """Run the backend in development mode using Granian. @@ -780,6 +883,7 @@ def run_granian_backend(host: str, port: int, loglevel: LogLevel): import multiprocessing logger.debug("Using Granian for backend") + _cap_message_size_for_granian() set_dev_start_method() @@ -1002,6 +1106,7 @@ def run_uvicorn_backend_prod( *("--host", host), *("--port", str(port)), *("--workers", str(_get_backend_workers())), + *_uvicorn_websocket_args(), "--factory", app_module, ] @@ -1017,7 +1122,10 @@ def run_uvicorn_backend_prod( "-m", "gunicorn", "--preload", - *("--worker-class", "uvicorn.workers.UvicornH11Worker"), + *( + "--worker-class", + "reflex.utils.uvicorn_worker.ReflexUvicornWorker", + ), *("--threads", str(_get_backend_workers())), *("--bind", f"{host}:{port}"), *env_args, @@ -1138,6 +1246,7 @@ def run_granian_backend_prod( from granian.server import Server as Granian logger.debug("Using Granian for backend") + _cap_message_size_for_granian() if (start_method := _backend_start_method()) is not None: # Set before the preload so the app's import-time multiprocessing diff --git a/reflex/utils/uvicorn_worker.py b/reflex/utils/uvicorn_worker.py new file mode 100644 index 00000000000..3c9c90b9245 --- /dev/null +++ b/reflex/utils/uvicorn_worker.py @@ -0,0 +1,22 @@ +"""Gunicorn worker carrying the Reflex websocket settings into uvicorn. + +Gunicorn takes no uvicorn options on its command line, so the production +backend passes them through a worker class instead. Imported by dotted path +from the gunicorn process; nothing else imports it, which keeps the optional +uvicorn dependency optional. +""" + +from __future__ import annotations + +from uvicorn.workers import UvicornH11Worker + +from reflex.utils.exec import uvicorn_websocket_options + + +class ReflexUvicornWorker(UvicornH11Worker): + """Uvicorn worker applying the app's websocket message policy.""" + + CONFIG_KWARGS = { + **UvicornH11Worker.CONFIG_KWARGS, + **uvicorn_websocket_options(), + } diff --git a/tests/benchmarks/support/idle_clients.py b/tests/benchmarks/support/idle_clients.py index f21321d13b8..0adff94e8ed 100644 --- a/tests/benchmarks/support/idle_clients.py +++ b/tests/benchmarks/support/idle_clients.py @@ -1,20 +1,50 @@ -"""Hold idle Socket.IO sessions from a separate process. +"""Hold idle event websocket sessions from a separate process. Memory-per-session benchmarks run the backend inside the test process, so client-side allocations would pollute resident-memory measurements. This module connects a number of sessions, primes each with one event so the backend materializes its state, reports readiness on stdout, and holds the -connections until stdin closes. +connections, answering the server's pings, until stdin closes. """ from __future__ import annotations +import contextlib import json import sys +import threading -import socketio +from reflex_base.environment import environment +from websockets.exceptions import ConnectionClosed -from tests.benchmarks.support.socket_client import _connect +from tests.benchmarks.support.socket_client import EventClient, _codec, _connect + + +def _keepalive_interval() -> float: + """Seconds between checks for pings while the sessions are held. + + A pong goes out at the first check after its ping, so checking faster than + the server's ping timeout keeps every session inside its deadline. + + Returns: + The interval. + """ + return min(1.0, environment.REFLEX_SOCKET_TIMEOUT.get().total_seconds() / 2) + + +def _keep_alive(clients: list[EventClient], stop: threading.Event) -> None: + """Answer the server's pings until stopped. + + Args: + clients: The held sessions. + stop: Set when the sessions are released. + """ + interval = _keepalive_interval() + while not stop.wait(interval): + for client in clients: + # A session the server ended stays ended; keep the others alive. + with contextlib.suppress(ConnectionClosed, ConnectionError): + client.drain(0) def main() -> int: @@ -25,22 +55,26 @@ def main() -> int: """ url, token_prefix, count, payload_json = sys.argv[1:5] payload = json.loads(payload_json) - clients: list[socketio.SimpleClient] = [] + clients: list[EventClient] = [] try: + codec = _codec(url, timeout=10) for index in range(int(count)): - client = socketio.SimpleClient(logger=False) - _connect(client, url, f"{token_prefix}-{index}", "/_event", timeout=10) - client.emit("event", payload) - response = client.receive(timeout=10) - while response and response[0] != "event": - response = client.receive(timeout=10) + client = _connect(url, f"{token_prefix}-{index}", codec, timeout=10) clients.append(client) + client.emit("event", payload) + client.receive(10, "event") print("ready", flush=True) - sys.stdin.read() + stop = threading.Event() + keeper = threading.Thread(target=_keep_alive, args=(clients, stop)) + keeper.start() + try: + sys.stdin.read() + finally: + stop.set() + keeper.join() finally: for client in clients: - if client.connected: - client.disconnect() + client.close() return 0 diff --git a/tests/benchmarks/support/socket_client.py b/tests/benchmarks/support/socket_client.py index e64c782c128..4c40fe036e1 100644 --- a/tests/benchmarks/support/socket_client.py +++ b/tests/benchmarks/support/socket_client.py @@ -1,24 +1,42 @@ -"""Protocol-aware Socket.IO client used by scheduled load tests.""" +"""Blocking client of Reflex's event websocket used by scheduled load tests. + +It speaks the plain WebSocket protocol of Reflex's default transport, and the +Socket.IO packets of a backend that refuses that endpoint (``transport="socketio"`` +apps and the python-socketio baseline server), so both sides of a comparison +pay the same client cost. The protocol is detected once per backend. +""" from __future__ import annotations import asyncio import concurrent.futures +import contextlib import dataclasses import functools import json +import threading import time from collections.abc import Callable, Mapping from typing import Any -from urllib.parse import quote -import socketio -from socketio.exceptions import TimeoutError as SocketTimeoutError +from reflex_bench.drivers.events import ( + CODECS, + HANDSHAKE, + PING, + Codec, + detect_protocol, + end_reason, +) +from websockets.exceptions import ConnectionClosed +from websockets.sync.client import ClientConnection, connect # Silence window that marks the connection quiescent while draining the extra # deltas the first event on a fresh token emits (hydrate + on_load_internal). _PRIME_DRAIN_TIMEOUT = 0.5 +_codecs: dict[str, Codec] = {} +_codecs_lock = threading.Lock() + @dataclasses.dataclass(frozen=True) class ClientLoadResult: @@ -40,62 +58,138 @@ class ReconnectResult: errors: tuple[str, ...] -def _payload_size(response: Any) -> int: - """Approximate the wire size of a received update payload. +class EventClient: + """One session on Reflex's event websocket, used from one thread at a time.""" + + def __init__(self, ws: ClientConnection, codec: Codec): + """Wrap a connected websocket. + + Args: + ws: The websocket. + codec: The protocol it speaks. + """ + self._ws = ws + self._codec = codec + + def emit(self, name: str, data: Any) -> None: + """Send a message. + + Args: + name: The message name. + data: Its payload. + """ + self._ws.send(self._codec.emit(name, data)) + + def receive(self, timeout: float, name: str | None = None) -> list[Any]: + """Wait for the next message, answering pings on the way. + + Args: + timeout: Seconds to wait; 0 takes only what already arrived. + name: The message name to wait for, skipping others; any if None. + + Returns: + The message, ``[name, *args]``. + + Raises: + ConnectionError: When the server refuses or ends the session. + """ + codec = self._codec + deadline = time.monotonic() + timeout + while True: + message = self._ws.recv(timeout=deadline - time.monotonic()) + args = codec.parse(message) + if args is None: + # A frame without a message, e.g. a binary one. + continue + if args[0] == PING: + self._ws.send(codec.pong) + elif (reason := end_reason(args, message)) is not None: + raise ConnectionError(reason) + elif name is None or args[0] == name: + return args + + def drain(self, quiet_s: float) -> None: + """Consume frames until the socket is quiet, answering pings. + + Args: + quiet_s: Seconds of silence that end the drain. + """ + with contextlib.suppress(TimeoutError): + while True: + self.receive(quiet_s) + + def close(self) -> None: + """Leave the session and close the websocket.""" + if (leave := self._codec.leave) is not None: + with contextlib.suppress(ConnectionClosed): + self._ws.send(leave) + self._ws.close() + + +def _codec(url: str, timeout: float) -> Codec: + """Find the protocol of a backend's event websocket, detecting it once per backend. Args: - response: Decoded ``[event, payload]`` message from the client. + url: Backend URL. + timeout: Maximum wait for the detection's connection. Returns: - Compact JSON byte length of the payload, excluding Socket.IO framing. + Its codec. """ - payload = response[1] if len(response) > 1 else None - return len(json.dumps(payload, separators=(",", ":"), default=str)) + # Concurrent clients wait for one detection instead of each running one. + with _codecs_lock: + if (codec := _codecs.get(url)) is None: + codec = _codecs[url] = CODECS[asyncio.run(detect_protocol(url, timeout))] + return codec -def _connect( - client: socketio.SimpleClient, - url: str, - token: str, - namespace: str, - timeout: float, -) -> None: - """Connect a client the way the Reflex frontend does. +def _connect(url: str, token: str, codec: Codec, timeout: float) -> EventClient: + """Connect a client the way the Reflex frontend does, without a boot event. Args: - client: Blocking Socket.IO client. - url: Backend Socket.IO URL. + url: Backend URL. token: Reflex client token. - namespace: Reflex Socket.IO namespace. - timeout: Maximum connection wait. + codec: The backend's protocol. + timeout: Maximum wait for the connection and its handshake. + + Returns: + The connected client. """ - separator = "&" if "?" in url else "?" - client.connect( - f"{url}{separator}token={quote(token)}", - transports=["websocket"], - socketio_path="_event", - headers={"Origin": url}, - namespace=namespace, - wait_timeout=max(1, int(timeout)), + ws = connect( + codec.url(url, token), + additional_headers={"Origin": url}, + compression=None, + proxy=None, + open_timeout=timeout, + ping_interval=None, + close_timeout=1, + max_size=None, ) + client = EventClient(ws, codec) + try: + ws.send(codec.join) + client.receive(timeout, HANDSHAKE) + except BaseException: + ws.close() + raise + return client -def _drain(client: socketio.SimpleClient, timeout: float) -> None: - """Consume buffered responses until the socket is quiet for ``timeout``. +def _payload_size(response: Any) -> int: + """Approximate the wire size of a received update payload. Args: - client: Blocking Socket.IO client. - timeout: Silence window that marks the connection quiescent. + response: Decoded ``[event, payload]`` message from the client. + + Returns: + Compact JSON byte length of the payload, excluding the protocol's framing. """ - while True: - try: - client.receive(timeout=timeout) - except SocketTimeoutError: - return + payload = response[1] if len(response) > 1 else None + return len(json.dumps(payload, separators=(",", ":"), default=str)) def _prime( - client: socketio.SimpleClient, + client: EventClient, event_name: str, payload: Mapping[str, Any], timeout: float, @@ -109,13 +203,13 @@ def _prime( events, which would otherwise report each latency pipelined two events deep. Args: - client: Blocking Socket.IO client. - event_name: Socket event name used for requests. + client: Event websocket client. + event_name: Message name used for requests. payload: Event payload emitted to trigger hydration. timeout: Maximum wait bounding the drain silence window. """ client.emit(event_name, dict(payload)) - _drain(client, min(timeout, _PRIME_DRAIN_TIMEOUT)) + client.drain(min(timeout, _PRIME_DRAIN_TIMEOUT)) async def run_socket_client( @@ -126,20 +220,18 @@ async def run_socket_client( *, event_name: str = "event", response_name: str = "event", - namespace: str = "/_event", timeout: float = 10, executor: concurrent.futures.Executor | None = None, ) -> ClientLoadResult: """Send events sequentially and measure matching socket responses. Args: - url: Backend Socket.IO URL. + url: Backend URL. token: Reflex client token. payload: Event payload emitted for each operation. events: Number of operations. - event_name: Socket event name used for requests. - response_name: Socket event name carrying state updates. - namespace: Reflex Socket.IO namespace. + event_name: Message name used for requests. + response_name: Message name carrying state updates. timeout: Maximum response wait per operation. executor: Optional executor that owns the blocking socket client. @@ -157,7 +249,6 @@ async def run_socket_client( events, event_name, response_name, - namespace, timeout, ), ) @@ -170,40 +261,36 @@ def _run_socket_client_sync( events: int, event_name: str, response_name: str, - namespace: str, timeout: float, ) -> ClientLoadResult: - """Run a load client with the websocket-client transport. + """Run a load client on a blocking websocket. Args: - url: Backend Socket.IO URL. + url: Backend URL. token: Reflex client token. payload: Event payload emitted for each operation. events: Number of operations. - event_name: Socket event name used for requests. - response_name: Socket event name carrying state updates. - namespace: Reflex Socket.IO namespace. + event_name: Message name used for requests. + response_name: Message name carrying state updates. timeout: Maximum response wait per operation. Returns: Per-operation latency and error observations. """ - client = socketio.SimpleClient(logger=False) + client: EventClient | None = None latencies: list[float] = [] payload_sizes: list[int] = [] errors: list[str] = [] try: - _connect(client, url, token, namespace, timeout) + client = _connect(url, token, _codec(url, timeout), timeout) _prime(client, event_name, payload, timeout) for _ in range(events): started = time.perf_counter_ns() client.emit(event_name, dict(payload)) try: - response = client.receive(timeout=timeout) - while response and response[0] != response_name: - response = client.receive(timeout=timeout) - except SocketTimeoutError: + response = client.receive(timeout, response_name) + except TimeoutError: errors.append("response_timeout") continue latencies.append((time.perf_counter_ns() - started) / 1_000_000) @@ -211,8 +298,8 @@ def _run_socket_client_sync( except Exception as err: errors.append(f"{type(err).__name__}: {err}") finally: - if client.connected: - client.disconnect() + if client is not None: + client.close() return ClientLoadResult( token, tuple(latencies), tuple(errors), tuple(payload_sizes) @@ -226,19 +313,17 @@ async def run_reconnect_client( *, event_name: str = "event", response_name: str = "event", - namespace: str = "/_event", timeout: float = 10, executor: concurrent.futures.Executor | None = None, ) -> ReconnectResult: """Connect, send one event, and measure connect and first-response time. Args: - url: Backend Socket.IO URL. + url: Backend URL. token: Reflex client token. payload: Event payload emitted after connecting. - event_name: Socket event name used for the request. - response_name: Socket event name carrying state updates. - namespace: Reflex Socket.IO namespace. + event_name: Message name used for the request. + response_name: Message name carrying state updates. timeout: Maximum connection and response wait. executor: Optional executor that owns the blocking socket client. @@ -255,7 +340,6 @@ async def run_reconnect_client( payload, event_name, response_name, - namespace, timeout, ), ) @@ -267,45 +351,43 @@ def _run_reconnect_client_sync( payload: Mapping[str, Any], event_name: str, response_name: str, - namespace: str, timeout: float, ) -> ReconnectResult: - """Run one reconnect client with the websocket-client transport. + """Run one reconnect client on a blocking websocket. Args: - url: Backend Socket.IO URL. + url: Backend URL. token: Reflex client token. payload: Event payload emitted after connecting. - event_name: Socket event name used for the request. - response_name: Socket event name carrying state updates. - namespace: Reflex Socket.IO namespace. + event_name: Message name used for the request. + response_name: Message name carrying state updates. timeout: Maximum connection and response wait. Returns: Connect and first-response observations. """ - client = socketio.SimpleClient(logger=False) + client: EventClient | None = None connect_ms = 0.0 first_response_ms = 0.0 errors: list[str] = [] try: + # Outside the clock: the first client of a backend runs the detection. + codec = _codec(url, timeout) started = time.perf_counter_ns() - _connect(client, url, token, namespace, timeout) + client = _connect(url, token, codec, timeout) connect_ms = (time.perf_counter_ns() - started) / 1_000_000 started = time.perf_counter_ns() client.emit(event_name, dict(payload)) - response = client.receive(timeout=timeout) - while response and response[0] != response_name: - response = client.receive(timeout=timeout) + client.receive(timeout, response_name) first_response_ms = (time.perf_counter_ns() - started) / 1_000_000 - except SocketTimeoutError: + except TimeoutError: errors.append("response_timeout") except Exception as err: errors.append(f"{type(err).__name__}: {err}") finally: - if client.connected: - client.disconnect() + if client is not None: + client.close() return ReconnectResult(token, connect_ms, first_response_ms, tuple(errors)) diff --git a/tests/benchmarks/test_event_processing.py b/tests/benchmarks/test_event_processing.py index 46ec2c943a2..53f99984be1 100644 --- a/tests/benchmarks/test_event_processing.py +++ b/tests/benchmarks/test_event_processing.py @@ -428,8 +428,8 @@ def on_event_harness(): """Set up an EventNamespace with a connected socket for benchmarking on_event. The event processor's enqueue is mocked out so the benchmark isolates the - per-event router_data preparation (which reuses the connection-scoped - data gathered once in on_connect). + per-event router_data preparation (which reuses the headers decoded once + into the connection's ASGI scope). Yields: An async callable that feeds the given number of events through @@ -478,6 +478,11 @@ async def run_events(num_events: int) -> None: }, ) + # The namespace reads the connection scope through the Socket.IO server. + # A plain function stands in for the lookup, so the measurement is not + # dominated by a mock recording a call per event. + app.sio = mock.Mock(get_environ=lambda *_args: environ) + loop = asyncio.new_event_loop() loop.run_until_complete(namespace.on_connect(sid, environ)) yield run_events, loop diff --git a/tests/benchmarks/test_event_transport.py b/tests/benchmarks/test_event_transport.py new file mode 100644 index 00000000000..3b62e07dc55 --- /dev/null +++ b/tests/benchmarks/test_event_transport.py @@ -0,0 +1,277 @@ +"""Benchmarks comparing the plain WebSocket transport with Socket.IO. + +Measures the server-side transport layer in isolation: inbound event frames +from an established connection to the (mocked) event processor, and outbound +state updates to the (mocked) wire. Both transports share BaseEventNamespace, +so the difference is the framing and dispatch layer. Socket.IO runs with +``async_handlers=False`` (inline dispatch), its cheapest configuration. +""" + +import asyncio +import json +from collections.abc import Awaitable, Callable +from types import SimpleNamespace +from typing import Any +from unittest import mock + +import pytest +import pytest_asyncio +from pytest_codspeed import BenchmarkFixture + +from reflex.event_namespace import CONNECT_MESSAGE, WebsocketEventNamespace, _Connection +from reflex.state import StateUpdate + +NUM_MESSAGES = 100 +NAMESPACE = "/_event" +TOKEN = "benchmark-token" + +_EVENT_FIELDS = { + "name": "benchmark___state.increment", + "router_data": { + "pathname": "/benchmark", + "asPath": "/benchmark?tab=2", + "query": {"tab": "2"}, + }, + "payload": {"value": 42, "label": "increment", "flag": True}, +} + +_UPDATE = StateUpdate( + delta={ + "benchmark___state": {f"var_{i}": f"value_{i}" for i in range(15)} + | {"counter": 42, "flag": True, "items": list(range(10))}, + } +) + +_DISCONNECT = object() + +_ASGI_SCOPE = { + "type": "websocket", + "headers": [(b"host", b"localhost")], + "client": ("127.0.0.1", 1234), +} + + +def _make_app(sio: Any = None) -> SimpleNamespace: + """Build a minimal app double for the event namespace. + + Returns: + The app double. + """ + enqueued: list[Any] = [] + + async def enqueue(token: str, event: Any) -> None: # noqa: RUF029 + enqueued.append((token, event)) + + return SimpleNamespace( + _state=None, + # Not started, like an app before its lifespan: no page loads to cancel. + _event_processor=None, + sio=sio, + router=lambda _path: None, + event_processor=SimpleNamespace(enqueue=enqueue), + enqueued=enqueued, + ) + + +class FakeWebSocket: + """Minimal stand-in for a starlette WebSocket.""" + + def __init__(self, frames: list[str]): + """Initialize with the inbound frames to deliver.""" + self.scope: dict[str, Any] = { + "type": "websocket", + "query_string": f"token={TOKEN}".encode(), + "subprotocols": [], + "headers": [(b"host", b"localhost")], + "client": ("127.0.0.1", 1234), + } + self.headers: dict[str, str] = {} + self.sent: list[str] = [] + self._incoming = [*frames, _DISCONNECT] + self._pos = 0 + + async def accept(self, subprotocol: str | None = None): + """Accept the connection.""" + + async def send_text(self, text: str): + """Record an outgoing frame.""" + self.sent.append(text) + + async def receive(self) -> dict[str, Any]: + """Return the next queued frame as an ASGI message. + + Returns: + The ASGI websocket message. + """ + item = self._incoming[self._pos] + self._pos += 1 + if item is _DISCONNECT: + return {"type": "websocket.disconnect", "code": 1000} + return {"type": "websocket.receive", "text": item} + + async def close(self, code: int = 1000, reason: str | None = None): + """Close the connection.""" + + +@pytest_asyncio.fixture +async def websocket_inbound(): # noqa: RUF029 - async so it runs on the benchmark loop + """Runner delivering NUM_MESSAGES event frames over the plain transport. + + Yields: + An async callable running one full connection lifecycle. + """ + with mock.patch("reflex.utils.prerequisites.check_redis_used", return_value=False): + app = _make_app() + namespace = WebsocketEventNamespace(NAMESPACE, app) # pyright: ignore[reportArgumentType] + frames = [ + json.dumps([CONNECT_MESSAGE]), + *[json.dumps(["event", _EVENT_FIELDS])] * NUM_MESSAGES, + ] + + async def run() -> None: + websocket = FakeWebSocket(frames) + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + assert len(app.enqueued) >= NUM_MESSAGES + + yield run + + +def _make_socketio_transport() -> tuple[Any, Any, SimpleNamespace]: + """Build a Socket.IO server and namespace over an app double. + + The server runs with inline dispatch and a discarding writer, so a + benchmark measures framing and dispatch rather than the network. + + Returns: + The (server, namespace, app double) triple. + """ + pytest.importorskip("socketio") + from socketio import AsyncServer + + from reflex.socketio_namespace import _SOCKET_JSON_CODEC, EventNamespace + + sio = AsyncServer(async_mode="asgi", async_handlers=False, json=_SOCKET_JSON_CODEC) + with mock.patch("reflex.utils.prerequisites.check_redis_used", return_value=False): + app = _make_app(sio=sio) + namespace = EventNamespace(NAMESPACE, app) # pyright: ignore[reportArgumentType] + sio.register_namespace(namespace) + + async def eio_send(_eio_sid: str, _data: str) -> None: + pass + + sio.eio.send = eio_send + return sio, namespace, app + + +@pytest_asyncio.fixture +async def socketio_inbound(): # noqa: RUF029 - async so it runs on the benchmark loop + """Runner delivering NUM_MESSAGES event packets over Socket.IO. + + Returns: + An async callable running one full connection lifecycle. + """ + sio, _namespace, app = _make_socketio_transport() + event_packet = "2" + NAMESPACE + "," + json.dumps(["event", _EVENT_FIELDS]) + counter = 0 + + async def run() -> None: + nonlocal counter + counter += 1 + eio_sid = f"eio-{counter}" + await sio._handle_eio_connect( + eio_sid, + {"QUERY_STRING": f"token={TOKEN}-{counter}", "asgi.scope": _ASGI_SCOPE}, + ) + await sio._handle_eio_message(eio_sid, "0" + NAMESPACE + ",") + for _ in range(NUM_MESSAGES): + await sio._handle_eio_message(eio_sid, event_packet) + await sio._handle_eio_message(eio_sid, "1" + NAMESPACE + ",") + # Let the disconnect cleanup task run. + for _ in range(3): + await asyncio.sleep(0) + assert len(app.enqueued) >= NUM_MESSAGES + + return run + + +@pytest_asyncio.fixture +async def websocket_outbound(): + """Runner emitting NUM_MESSAGES state updates over the plain transport. + + Yields: + An async callable emitting the updates. + """ + with mock.patch("reflex.utils.prerequisites.check_redis_used", return_value=False): + app = _make_app() + namespace = WebsocketEventNamespace(NAMESPACE, app) # pyright: ignore[reportArgumentType] + websocket = FakeWebSocket([]) + connection = _Connection(websocket) # pyright: ignore[reportArgumentType] + namespace._connections["sid-1"] = connection + await namespace.link_token_to_sid("sid-1", TOKEN) + + async def run() -> None: + for _ in range(NUM_MESSAGES): + await namespace.emit_update(_UPDATE, TOKEN) + + yield run + connection.stop() + + +@pytest_asyncio.fixture +async def socketio_outbound(): + """Runner emitting NUM_MESSAGES state updates over Socket.IO. + + Returns: + An async callable emitting the updates. + """ + sio, namespace, _app = _make_socketio_transport() + + async def run() -> None: + for _ in range(NUM_MESSAGES): + await namespace.emit_update(_UPDATE, TOKEN) + + # Connect a socket.io session and link the token to its sid. + eio_sid = "eio-emit" + await sio._handle_eio_connect( + eio_sid, {"QUERY_STRING": f"token={TOKEN}", "asgi.scope": _ASGI_SCOPE} + ) + await sio._handle_eio_message(eio_sid, "0" + NAMESPACE + ",") + assert TOKEN in namespace.token_to_sid + + return run + + +def _benchmark_runner( + benchmark: BenchmarkFixture, runner: Callable[[], Awaitable[None]] +): + """Benchmark one full run of an async transport runner. + + Args: + benchmark: The benchmark fixture. + runner: The async callable to measure. + """ + loop = asyncio.get_event_loop() + + @benchmark + def _(): + loop.run_until_complete(runner()) + + +def test_transport_inbound_websocket(websocket_inbound, benchmark: BenchmarkFixture): + """Benchmark inbound event handling on the plain WebSocket transport.""" + _benchmark_runner(benchmark, websocket_inbound) + + +def test_transport_inbound_socketio(socketio_inbound, benchmark: BenchmarkFixture): + """Benchmark inbound event handling on the Socket.IO transport.""" + _benchmark_runner(benchmark, socketio_inbound) + + +def test_transport_outbound_websocket(websocket_outbound, benchmark: BenchmarkFixture): + """Benchmark emitting state updates on the plain WebSocket transport.""" + _benchmark_runner(benchmark, websocket_outbound) + + +def test_transport_outbound_socketio(socketio_outbound, benchmark: BenchmarkFixture): + """Benchmark emitting state updates on the Socket.IO transport.""" + _benchmark_runner(benchmark, socketio_outbound) diff --git a/tests/benchmarks/test_socket_headers.py b/tests/benchmarks/test_socket_headers.py index 3c5b56d6fe4..0f694ac88ac 100644 --- a/tests/benchmarks/test_socket_headers.py +++ b/tests/benchmarks/test_socket_headers.py @@ -1,14 +1,14 @@ -"""Benchmarks for per-event ASGI header decoding.""" +"""Benchmarks for ASGI header decoding.""" import pytest from pytest_codspeed import BenchmarkFixture -from reflex.app import _decode_asgi_headers +from reflex.event_namespace import _decode_asgi_headers @pytest.mark.parametrize("count", [0, 4, 16, 64]) def test_decode_event_headers(count: int, benchmark: BenchmarkFixture): - """Benchmark the header decoding used by ``EventNamespace.on_event``. + """Benchmark the header decoding behind a connection's router data. Args: count: Number of headers. diff --git a/tests/integration/test_client_storage.py b/tests/integration/test_client_storage.py index 1fadc888105..65e6af12bd9 100644 --- a/tests/integration/test_client_storage.py +++ b/tests/integration/test_client_storage.py @@ -524,7 +524,7 @@ def set_sub_sub(var: str, value: str): l1s = driver.find_element(By.ID, "l1s") s1s = driver.find_element(By.ID, "s1s") - assert c1.text == "c1 value" + AppHarness.expect(lambda: c1.text == "c1 value") assert c2.text == "c2 value" assert c3.text == "" # cookie expired so value removed from state assert c4.text == "c4 value" @@ -573,7 +573,7 @@ def set_sub_sub(var: str, value: str): l1s = driver.find_element(By.ID, "l1s") s1s = driver.find_element(By.ID, "s1s") - assert c1.text == "c1 value" + AppHarness.expect(lambda: c1.text == "c1 value") assert c2.text == "c2 value" assert c3.text == "" # temporary cookie expired after reset state! assert c4.text == "c4 value" @@ -672,7 +672,7 @@ def set_sub_sub(var: str, value: str): l1s = driver.find_element(By.ID, "l1s") s1s = driver.find_element(By.ID, "s1s") - assert c1.text == "c1 post expire" + AppHarness.expect(lambda: c1.text == "c1 post expire") assert c2.text == "c2 value" assert c3.text == "" # temporary cookie expired after reset state! assert c4.text == "c4 value" diff --git a/tests/integration/tests_playwright/test_channels.py b/tests/integration/tests_playwright/test_channels.py new file mode 100644 index 00000000000..eb9f1e49236 --- /dev/null +++ b/tests/integration/tests_playwright/test_channels.py @@ -0,0 +1,231 @@ +"""Integration tests for channels multiplexed onto the event websocket. + +Exercises the parts only a browser can prove: the client opens its channel +over the app's own socket, binary attachments survive the round trip in both +directions, and each one lands aligned enough to read as a typed array. +""" + +from __future__ import annotations + +from collections.abc import Generator + +import pytest +from playwright.sync_api import Page, expect + +from reflex.testing import AppHarness + + +def ChannelApp(): + """App serving an echo channel driven from the browser.""" + import reflex as rx + + PROBE_SETUP = """ + if (typeof window !== "undefined") { + const probeChannel = getChannel("probe"); + window.__probe = { + received: [], errors: [], connected: false, connects: 0, disconnects: 0, + }; + probeChannel.on("connect", () => { + window.__probe.connected = true; + window.__probe.connects += 1; + }); + probeChannel.on("disconnect", () => { + window.__probe.connected = false; + // Counted, not just flagged: the reconnect can restore the flag + // between two polls, but a count cannot be missed. + window.__probe.disconnects += 1; + }); + // Drop the socket the way the transport's own watchdog does on a dead + // connection, so the test exercises the supported reconnect path. + window.__probe.drop = () => probeChannel._transport._dropConnection("test"); + // Close the socket and emit in the same tick: it is CLOSING, but its + // close is not reported yet, so the channel still reports connected + // and the frame queues on the transport rather than on the channel. + // (A drop detaches the channel at once, queueing on the channel.) + window.__probe.dropThenEmit = (n) => { + probeChannel._transport._ws.close(); + probeChannel.emit("push", { n }, [new Uint8Array([1, 2, 3, 4])]); + }; + // Emit while down, then mutate what was passed: the queued message must + // still carry the values it was emitted with. + window.__probe.pushThenMutate = (n) => { + const meta = { n }; + const bytes = new Uint8Array([1, 2, 3, 4]); + probeChannel.emit("push", meta, [bytes]); + meta.n = -1; + bytes.fill(9); + }; + probeChannel.on("error", (error) => { + window.__probe.errors.push(error); + }); + probeChannel.on("echo", (data, buffers) => { + window.__probe.received.push({ + data, + bytes: Array.from(buffers[0]), + aligned: buffers[0].byteOffset % 8 === 0, + }); + }); + window.__probe.push = (n) => + probeChannel.emit("push", { n }, [new Uint8Array([1, 2, 3, 4])]); + window.__probe.broadcast = () => probeChannel.emit("broadcast", null); + } + """ + + class EchoChannel(rx.channels.Channel): + name = "probe" + accepts_binary = True + + async def on_open(self, session): + session.join("all") + + async def on_message(self, session, event, data, buffers): + if event == "push": + await session.send( + "echo", + {"n": data["n"], "sizes": [len(buffer) for buffer in buffers]}, + [bytes(reversed(buffers[0]))], + ) + elif event == "broadcast": + await self.send_to_room( + "all", "echo", {"n": -1, "sizes": []}, [b"\x09\x08\x07"] + ) + + class ChannelProbe(rx.Fragment): + def add_imports(self): + return {"$/utils/state": ["getChannel"]} + + def add_custom_code(self): + return [PROBE_SETUP] + + @rx.page("/") + def index(): + return rx.box(ChannelProbe.create(), rx.text("ready", id="ready")) + + app = rx.App() + app.register_channel(EchoChannel()) + + +@pytest.fixture(scope="module") +def channel_app( + tmp_path_factory: pytest.TempPathFactory, +) -> Generator[AppHarness, None, None]: + """Start the ChannelApp. + + Args: + tmp_path_factory: pytest fixture for creating temporary directories. + + Yields: + Running AppHarness instance. + """ + with AppHarness.create( + root=tmp_path_factory.mktemp("channel_app"), + app_source=ChannelApp, + ) as harness: + assert harness.app_instance is not None, "app is not running" + yield harness + + +def _connected_probe(channel_app: AppHarness, page: Page) -> None: + """Load the page and wait for the channel to open. + + Args: + channel_app: Running AppHarness instance. + page: Playwright page fixture. + """ + assert channel_app.frontend_url is not None + page.goto(channel_app.frontend_url) + expect(page.locator("#ready")).to_have_text("ready") + # A cold frontend build reloads the page once the dev server finishes + # optimizing its dependencies, which resets the probe; let that settle + # before the test starts counting on its state. + page.wait_for_load_state("networkidle") + page.wait_for_function("window.__probe?.connected === true") + + +def test_binary_message_round_trip(channel_app: AppHarness, page: Page): + """A binary attachment survives the round trip and arrives aligned. + + Args: + channel_app: Running AppHarness instance. + page: Playwright page fixture. + """ + _connected_probe(channel_app, page) + + page.evaluate("window.__probe.push(7)") + page.wait_for_function("window.__probe.received.length > 0") + + received = page.evaluate("window.__probe.received[0]") + assert received["data"] == {"n": 7, "sizes": [4]} + assert received["bytes"] == [4, 3, 2, 1] + assert received["aligned"] + assert page.evaluate("window.__probe.errors") == [] + + +def test_room_broadcast_reaches_the_client(channel_app: AppHarness, page: Page): + """A server-side room fan-out reaches a subscribed browser. + + Args: + channel_app: Running AppHarness instance. + page: Playwright page fixture. + """ + _connected_probe(channel_app, page) + + page.evaluate("window.__probe.broadcast()") + page.wait_for_function("window.__probe.received.length > 0") + + received = page.evaluate("window.__probe.received[0]") + assert received["data"] == {"n": -1, "sizes": []} + assert received["bytes"] == [9, 8, 7] + + +def test_channel_reopens_and_flushes_after_a_reconnect( + channel_app: AppHarness, page: Page +): + """A dropped socket reopens the channel and delivers what was queued. + + Args: + channel_app: Running AppHarness instance. + page: Playwright page fixture. + """ + _connected_probe(channel_app, page) + connects = page.evaluate("window.__probe.connects") + + disconnects = page.evaluate("window.__probe.disconnects") + page.evaluate("window.__probe.drop()") + page.wait_for_function(f"window.__probe.disconnects > {disconnects}") + # Emitted while the socket is down: queued, not lost, and not rewritten by + # what the caller did with the payload afterwards. + page.evaluate("window.__probe.pushThenMutate(9)") + # Reconnects are retried with backoff, which outlasts the default wait on + # a machine busy building the frontend. + page.wait_for_function(f"window.__probe.connects > {connects}", timeout=30_000) + page.wait_for_function("window.__probe.received.length > 0", timeout=30_000) + + received = page.evaluate("window.__probe.received[0]") + assert received["data"] == {"n": 9, "sizes": [4]} + assert received["bytes"] == [4, 3, 2, 1] + + +def test_frame_queued_on_the_transport_survives_the_reconnect( + channel_app: AppHarness, page: Page +): + """A frame emitted as the socket drops is delivered once it is back. + + Such a frame waits on the transport's own queue, so the channel has to be + reopened before that queue is flushed -- otherwise it reaches the backend + ahead of its `_open` and is answered with `channel_not_open`. + + Args: + channel_app: Running AppHarness instance. + page: Playwright page fixture. + """ + _connected_probe(channel_app, page) + connects = page.evaluate("window.__probe.connects") + + page.evaluate("window.__probe.dropThenEmit(5)") + page.wait_for_function(f"window.__probe.connects > {connects}", timeout=30_000) + page.wait_for_function("window.__probe.received.length > 0", timeout=30_000) + + received = page.evaluate("window.__probe.received[0]") + assert received["data"] == {"n": 5, "sizes": [4]} + assert page.evaluate("window.__probe.errors") == [] diff --git a/tests/integration/tests_playwright/test_hydration_storage.py b/tests/integration/tests_playwright/test_hydration_storage.py index 0fbeb3914dd..2e1cb694a7d 100644 --- a/tests/integration/tests_playwright/test_hydration_storage.py +++ b/tests/integration/tests_playwright/test_hydration_storage.py @@ -785,7 +785,11 @@ def proxy(route: WebSocketRoute): def test_hydration_over_server_limit_is_reported_once( hydration_storage_app: AppHarness, page: Page ): - """A hydrate over the server's announced limit is reported and not sent. + """A hydrate over the server's limit is reported once and not retried. + + Socket.IO checks it against the limit its handshake announced, before + sending it. The plain WebSocket sends it with its connect frame, which + saves the handshake round trip, and reports the server refusing it. Args: hydration_storage_app: The running app. @@ -802,7 +806,14 @@ def test_hydration_over_server_limit_is_reported_once( page.evaluate("localStorage.setItem('hydrate-local', 'x'.repeat(1_100_000))") opens.clear() page.reload() - expect(page.get_by_text("Reduce browser storage data and reload")).to_be_visible() + # Each transport names the limit: the plain WebSocket from the close reason. + expect( + page.get_by_text( + re.compile( + r"server limit \(1000000 \w+\)\. Reduce browser storage data and reload" + ) + ) + ).to_be_visible() page.wait_for_timeout(2000) assert len(opens) == 1, opens expect(page.locator("#hydrated")).to_have_text("false") diff --git a/tests/integration/tests_playwright/test_router_query.py b/tests/integration/tests_playwright/test_router_query.py index 132cd45f8e7..7816febd5d2 100644 --- a/tests/integration/tests_playwright/test_router_query.py +++ b/tests/integration/tests_playwright/test_router_query.py @@ -19,15 +19,16 @@ from __future__ import annotations +import json import re -from collections.abc import Generator +from collections.abc import Generator, Iterator +from typing import Any from urllib.parse import urlsplit import pytest from playwright.sync_api import Page, expect from reflex_base.config import get_config from reflex_base.constants.state import FIELD_MARKER -from socketio.packet import Packet from reflex.testing import AppHarness @@ -259,6 +260,30 @@ def _load(harness: AppHarness, page: Page) -> str: return base +def _state_updates(frames: list[str | bytes]) -> Iterator[dict[str, Any]]: + """Yield the state updates among frames received on the default transport. + + Its frames are JSON arrays, [event_name, payload]. + + Args: + frames: The received websocket frames. + + Yields: + The payload of each state update. + """ + for frame in frames: + if not isinstance(frame, str): + continue + message = json.loads(frame) + if ( + isinstance(message, list) + and len(message) == 2 + and message[0] == "event" + and isinstance(message[1], dict) + ): + yield message[1] + + def test_initial_connection_dispatches_both_substates( router_query_app: AppHarness, page: Page ): @@ -277,15 +302,7 @@ def test_initial_connection_dispatches_both_substates( _load(router_query_app, page) expect(page.locator("#connected-count")).to_have_value("1") - # Socket.IO event packets use the Engine.IO message and event prefixes, 42. - for frame in frames: - if not isinstance(frame, str) or not frame.startswith("42"): - continue - packet = Packet(encoded_packet=frame[1:]) - assert packet.data is not None - event, update = packet.data - if event != "event": - continue + for update in _state_updates(frames): updated_substates = sum( fields.get("load_count" + FIELD_MARKER) == 1 or fields.get("connected_count" + FIELD_MARKER) == 1 @@ -343,23 +360,17 @@ def test_request_credentials_stay_server_side(router_query_app: AppHarness, page expect(page.locator("#credentials-seen")).to_have_value("true") assert "privacy-test" not in page.evaluate("document.cookie") - received_headers = [] for frame in frames: decoded = frame.decode() if isinstance(frame, bytes) else frame assert secret not in decoded for credential in credential_headers.values(): assert credential not in decoded - if not isinstance(frame, str) or not frame.startswith("42"): - continue - packet = Packet(encoded_packet=frame[1:]) - assert packet.data is not None - event, update = packet.data - if event == "event": - received_headers.extend( - fields["rx_router_headers" + FIELD_MARKER] - for fields in update.get("delta", {}).values() - if "rx_router_headers" + FIELD_MARKER in fields - ) + received_headers = [ + fields["rx_router_headers" + FIELD_MARKER] + for update in _state_updates(frames) + for fields in update.get("delta", {}).values() + if "rx_router_headers" + FIELD_MARKER in fields + ] assert received_headers, "Hydration did not include the request headers" for headers in received_headers: assert "cookie" not in headers diff --git a/tests/integration/tests_playwright/test_stateless_app.py b/tests/integration/tests_playwright/test_stateless_app.py index 0451230a556..d7897873f10 100644 --- a/tests/integration/tests_playwright/test_stateless_app.py +++ b/tests/integration/tests_playwright/test_stateless_app.py @@ -5,8 +5,9 @@ import httpx import pytest from playwright.sync_api import Page, expect -from reflex_base.config import get_config from reflex_base.constants import Endpoint +from websockets.exceptions import InvalidStatus +from websockets.sync.client import connect from reflex.testing import AppHarness @@ -48,11 +49,13 @@ def test_statelessness(stateless_app: AppHarness, page: Page): """ assert stateless_app.frontend_url is not None assert stateless_app.backend is not None - assert stateless_app.backend.started + stateless_app.backend.wait_started(timeout=0) - config = get_config() - res = httpx.get(config.api_url + config.prepend_backend_path(str(Endpoint.EVENT))) - assert res.status_code == 404 + # Plain HTTP gets a 404 from a websocket route too: dial the handshake. + event_url = Endpoint.EVENT.get_url() + for url in (event_url, event_url.rstrip("/") + "/"): + with pytest.raises(InvalidStatus), connect(url, open_timeout=10): + pass res2 = httpx.get(Endpoint.PING.get_url()) assert res2.status_code == 200 diff --git a/tests/units/benchmarks/test_support.py b/tests/units/benchmarks/test_support.py index 7176c8672e4..6a4cb909124 100644 --- a/tests/units/benchmarks/test_support.py +++ b/tests/units/benchmarks/test_support.py @@ -2,6 +2,9 @@ import asyncio import json +import socket +import sys +import time from types import SimpleNamespace import pytest @@ -15,7 +18,13 @@ PerformanceReport, percentile, ) -from tests.benchmarks.support.socket_client import run_clients + +# The benchmark socket client speaks the wire protocols through reflex-bench, +# which is Linux-only like its own tests. Imported anywhere else, it would also +# make coverage count that whole untested package, so it is imported in tests. +linux_only = pytest.mark.skipif( + sys.platform != "linux", reason="the benchmark socket client builds on reflex-bench" +) def test_percentile_interpolates_and_validates(): @@ -140,9 +149,12 @@ def test_pipeline_trace_preserves_repeated_stage_pairs(): } +@linux_only @pytest.mark.asyncio async def test_run_clients_leaves_default_executor_usable(): """The load helper does not close the event loop's default executor.""" + from tests.benchmarks.support.socket_client import run_clients + results = await run_clients( 2, lambda index, executor: asyncio.get_running_loop().run_in_executor( @@ -154,6 +166,21 @@ async def test_run_clients_leaves_default_executor_usable(): assert await asyncio.to_thread(lambda: "still-usable") == "still-usable" +@linux_only +@pytest.mark.asyncio +async def test_reconnect_client_detects_the_protocol_within_its_timeout(): + """Detecting a hung backend's protocol is bounded by the client's timeout.""" + from tests.benchmarks.support.socket_client import run_reconnect_client + + # The listener never answers the upgrade. + with socket.create_server(("127.0.0.1", 0)) as listener: + url = f"http://127.0.0.1:{listener.getsockname()[1]}" + started = time.monotonic() + result = await run_reconnect_client(url, "token", {}, timeout=0.1) + assert result.errors == ("response_timeout",) + assert time.monotonic() - started < 5 + + class _FakeRedis: """Minimal async Redis stand-in recording lifecycle calls.""" @@ -322,3 +349,24 @@ def crash(self) -> None: BaselineSocketServer(response={}), ): pass + + +@linux_only +@pytest.mark.parametrize(("timeout", "interval"), [(None, 1.0), ("100ms", 0.05)]) +def test_idle_clients_answer_pings_within_the_ping_timeout( + monkeypatch: pytest.MonkeyPatch, timeout: str | None, interval: float +): + """Held sessions poll for pings fast enough to beat the server's deadline. + + A pong goes out at the first poll after its ping, so the gap between + client frames is the ping interval plus the poll period; it has to stay + under the interval plus the ping timeout. + """ + from tests.benchmarks.support import idle_clients + + if timeout is None: + monkeypatch.delenv("REFLEX_SOCKET_TIMEOUT", raising=False) + else: + monkeypatch.setenv("REFLEX_SOCKET_TIMEOUT", timeout) + + assert idle_clients._keepalive_interval() == pytest.approx(interval) diff --git a/tests/units/compiler/state_js.test.mjs b/tests/units/compiler/state_js.test.mjs index 688814a2fde..361794f5123 100644 --- a/tests/units/compiler/state_js.test.mjs +++ b/tests/units/compiler/state_js.test.mjs @@ -18,6 +18,9 @@ async function setup({ disabled = false, hidden = false, storedToken = "original-token", + transport = "websocket", + socketio = false, + importFails = false, clientStorage = {}, dropConnects = 0, dropReason = "transport close", @@ -70,13 +73,17 @@ async function setup({ window.setTimeout = context.setTimeout; window.clearTimeout = context.clearTimeout; - function io(url, opts) { + /** A transport double shared by the plain WebSocket and socket.io clients. */ + function fakeSocket(kind, url, opts) { const handlers = new Map(); const socket = { + kind, handlers, connected: false, + opens: 0, namespaceConnects: 0, disconnects: 0, + offs: 0, // Like socket.io-client: the URL path names the namespace. nsp: new URL(url).pathname, // Engine.IO frames of the CONNECT packets sent to the server. @@ -86,7 +93,7 @@ async function setup({ // The handshake announces the server's limit. engine: { _maxPayload: maxPayload, - transport: { name: opts.transports[0] }, + transport: { name: opts.transports?.[0] }, }, // Like socket.io-parser's Encoder.encodeAsString. encoder: { @@ -98,15 +105,16 @@ async function setup({ }, }, decoder: {}, - opens: 0, - open(callback) { - this.opens++; - this.openCallback = callback; - }, + }, + open() { + this.opens++; }, on(name, handler) { handlers.set(name, handler); }, + off() { + this.offs++; + }, connect() { // A ready transport can deliver these synchronously on namespace connect. for (const name of [ @@ -118,20 +126,42 @@ async function setup({ ]) { assert.ok(handlers.has(name), `missing handler: ${name}`); } - // Like Manager._packet: the server receives the frames the encoder returns. - const frames = this.io.encoder.encode({ - type: 0, - nsp: this.nsp, - data: this.auth, - }); - if (frames.length === 0) { + if (kind === "socketio") { + // Like Manager._packet: the server receives the frames the encoder returns. + const frames = this.io.encoder.encode({ + type: 0, + nsp: this.nsp, + data: this.auth, + }); + if (frames.length === 0) { + return; + } + this.frames.push(...frames.map((frame) => `4${frame}`)); + } else if ( + JSON.stringify(["_connect", this.auth]).length > maxPayload + ) { + // Like ReflexWebSocket when the backend refuses its connect frame + // for its size. + const error = new Error( + "Hydration payload exceeds the server limit. Reduce browser storage data and reload.", + ); + error.code = "hydration_payload_too_large"; + handlers.get("connect_error")(error); return; } - this.frames.push(...frames.map((frame) => `4${frame}`)); this.namespaceConnects++; if (this.namespaceConnects <= dropConnects) { - // The transport closes before the CONNECT is acknowledged. A - // transport error reports connect_error before the disconnect. + // The transport closes before the CONNECT is acknowledged. + if (kind === "websocket") { + // The plain WebSocket reports a close before its handshake as a + // failed attempt only. + handlers.get("connect_error")( + new Error("websocket connection failed"), + ); + return; + } + // A socket.io transport error reports connect_error before the + // disconnect. if (dropReason === "transport error") { handlers.get("connect_error")(new Error("xhr post error")); } @@ -151,11 +181,25 @@ async function setup({ sockets.push(socket); return socket; } + // Called with `new`: a constructor returning an object yields that object. + function ReflexWebSocket(url, opts) { + return fakeSocket("websocket", url, opts); + } + const io = (url, opts) => fakeSocket("socketio", url, opts); const imports = { "socket.io-client": { default: io }, + "$/utils/helpers/websocket": { + ReflexWebSocket, + disableChannels() {}, + getChannel() {}, + }, mergician: { mergician() {} }, "$/env.json": { - default: { EVENT: eventUrl, TRANSPORT: "websocket" }, + default: { + EVENT: eventUrl, + TRANSPORT: transport, + SOCKETIO: socketio, + }, }, "$/reflex.json": { default: { version: "test" } }, "universal-cookie": { default: class {} }, @@ -190,22 +234,15 @@ async function setup({ }, eventLoop: { addEvents() {}, connectErrors: [] }, }, - "$/utils/helpers/json": { parseJson: JSON.parse }, + "$/utils/helpers/json": { + parseJsonLenient() {}, + undefinedToNull: (_k, v) => (v === undefined ? null : v), + }, "$/utils/helpers/debounce": { default() {} }, "$/utils/helpers/throttle": { default() {} }, "$/utils/helpers/upload": { uploadFiles() {} }, }; - const module = new vm.SourceTextModule(source, { - context, - initializeImportMeta(meta) { - meta.hot = { - dispose(fn) { - dispose = fn; - }, - }; - }, - }); - await module.link((name) => { + const link = (name) => { assert.ok(imports[name], `unexpected import: ${name}`); return new vm.SyntheticModule( Object.keys(imports[name]), @@ -215,7 +252,29 @@ async function setup({ }, { context }, ); + }; + const module = new vm.SourceTextModule(source, { + context, + initializeImportMeta(meta) { + meta.hot = { + dispose(fn) { + dispose = fn; + }, + }; + }, + // The socket.io client is imported on demand. + async importModuleDynamically(name) { + if (importFails) { + // A chunk missing after a redeploy; browsers cache the failure. + throw new TypeError("Failed to fetch dynamically imported module"); + } + const imported = link(name); + await imported.link(() => {}); + await imported.evaluate(); + return imported; + }, }); + await module.link(link); await module.evaluate(); const socket = { current: null }; return { @@ -251,13 +310,13 @@ async function setup({ }; } -test("warm the transport without hydrating, then reuse it with every handler attached", async () => { +test("warm the transport without opening a session, then reuse it with every handler attached", async () => { const app = await setup(); app.flush(); assert.equal(app.sockets.length, 1); - assert.equal(app.sockets[0].io.opens, 1); + assert.equal(app.sockets[0].kind, "websocket"); + assert.equal(app.sockets[0].opens, 1); assert.equal(app.sockets[0].namespaceConnects, 0); - assert.equal(app.sockets[0].io.opts.autoConnect, false); assert.deepEqual(app.firstHydrates, []); await app.connect(); assert.equal(app.sockets.length, 1); @@ -293,9 +352,14 @@ const frameSize = (frame, transport) => ? [...frame].length : new TextEncoder().encode(frame).byteLength; -/** Connect an app whose first hydrate carries the payload, with the given server limit. */ +/** Connect a socket.io app whose first hydrate carries the payload, with the given server limit. */ async function connectWith(transport, firstPayload, maxPayload, options = {}) { - const app = await setup({ firstPayload, maxPayload, ...options }); + const app = await setup({ + socketio: true, + firstPayload, + maxPayload, + ...options, + }); app.flush(); await app.connect([transport]); return app; @@ -368,41 +432,48 @@ test("websocket limits count code points, polling limits count UTF-8 bytes", asy assert.match(polling.connectErrors[0].message, /bytes/); }); -test("an oversized first hydrate is reported once and not rebuilt per event", async () => { - const app = await connectWith( - "websocket", - { vars: { large: "x".repeat(1_000_001) } }, - 1_000_000, - ); - const socket = app.socket.current; - // Later events reconnect through connect(), and a visible tab directly. - await app.connect(); - socket.reconnect(); +for (const socketio of [false, true]) { + const over = socketio ? " over socket.io" : ""; - assert.deepEqual(app.firstHydrates, [true]); - assert.equal(socket.namespaceConnects, 0); - assert.equal(app.timers.size, 0); - assert.equal(app.connectErrors.length, 1); - assert.equal(app.consoleErrors.length, 1); -}); + test(`an oversized first hydrate${over} is reported once and not rebuilt per event`, async () => { + const app = await setup({ + socketio, + firstPayload: { vars: { large: "x".repeat(1_000_001) } }, + }); + app.flush(); + await app.connect(); + const socket = app.socket.current; + // Later events reconnect through connect(), and a visible tab directly. + await app.connect(); + socket.reconnect(); -test("an oversized full hydrate on reconnect is reported, not retried", async () => { - const app = await setup({ - reconnectPayload: { vars: { large: "x".repeat(1_000_001) } }, + assert.deepEqual(app.firstHydrates, [true]); + assert.equal(socket.namespaceConnects, 0); + assert.equal(app.timers.size, 0); + assert.equal(app.connectErrors.length, 1); + assert.equal(app.connectErrors[0].code, "hydration_payload_too_large"); + assert.deepEqual(app.consoleErrors, [[app.connectErrors[0].message]]); }); - app.flush(); - await app.connect(); - const socket = app.socket.current; - socket.connected = false; - socket.handlers.get("disconnect")("transport close"); - socket.reconnect(); - assert.equal(socket.namespaceConnects, 1); - assert.deepEqual(app.firstHydrates, [true, false]); - assert.equal(app.timers.size, 0); - assert.equal(app.connectErrors.length, 1); - assert.match(app.connectErrors[0].message, /exceeds/); -}); + test(`an oversized full hydrate on reconnect${over} is reported, not retried`, async () => { + const app = await setup({ + socketio, + reconnectPayload: { vars: { large: "x".repeat(1_000_001) } }, + }); + app.flush(); + await app.connect(); + const socket = app.socket.current; + socket.connected = false; + socket.handlers.get("disconnect")("transport close"); + socket.reconnect(); + + assert.equal(socket.namespaceConnects, 1); + assert.deepEqual(app.firstHydrates, [true, false]); + assert.equal(app.timers.size, 0); + assert.equal(app.connectErrors.length, 1); + assert.match(app.connectErrors[0].message, /exceeds/); + }); +} test("a new session's token is saved only once the mounted app connects", async () => { const app = await setup({ storedToken: null }); @@ -434,38 +505,80 @@ test("a warm transport carrying another session's token is replaced", async () = assert.equal(app.socket.current.io.opts.query.token, "replaced-token"); }); -test("reconnect uses the assigned token and requests a full hydrate", async () => { - const app = await setup(); - app.flush(); - await app.connect(); - app.socket.current.connected = false; - app.socket.current.reconnect(); - assert.equal(app.sockets.length, 1); - assert.equal(app.socket.current.io.opts.query.token, "assigned-token"); - assert.deepEqual(app.firstHydrates, [true, false]); - assert.equal(app.socket.current.namespaceConnects, 2); -}); +for (const socketio of [false, true]) { + test(`reconnect ${socketio ? "over socket.io " : ""}uses the assigned token and requests a full hydrate`, async () => { + const app = await setup({ socketio }); + app.flush(); + await app.connect(); + assert.equal(app.socket.current.kind, socketio ? "socketio" : "websocket"); + app.socket.current.connected = false; + app.socket.current.reconnect(); + assert.equal(app.sockets.length, 1); + assert.equal(app.socket.current.io.opts.query.token, "assigned-token"); + assert.deepEqual(app.firstHydrates, [true, false]); + assert.equal(app.socket.current.auth.event.name, "root.hydrate_and_load"); + assert.equal(app.socket.current.namespaceConnects, 2); + }); +} -test("a transport closed before the namespace connects backs off with an error", async () => { - const app = await setup({ dropConnects: Infinity }); - app.flush(); - await app.connect(); - const socket = app.socket.current; - assert.equal(socket.namespaceConnects, 1); - assert.equal(app.connectErrors.length, 1); - assert.deepEqual([...app.timers.values()], [200]); +for (const socketio of [false, true]) { + const over = socketio ? " over socket.io" : ""; - const [retry] = app.timers.keys(); - app.timers.clear(); - retry(); - assert.equal(socket.namespaceConnects, 2); - assert.equal(app.connectErrors.length, 2); - assert.deepEqual([...app.timers.values()], [400]); -}); + test(`a transport closed before the namespace connects${over} backs off with an error`, async () => { + const app = await setup({ socketio, dropConnects: Infinity }); + app.flush(); + await app.connect(); + const socket = app.socket.current; + assert.equal(socket.kind, socketio ? "socketio" : "websocket"); + assert.equal(socket.namespaceConnects, 1); + assert.equal(app.connectErrors.length, 1); + assert.deepEqual([...app.timers.values()], [200]); + + const [retry] = app.timers.keys(); + app.timers.clear(); + retry(); + assert.equal(socket.namespaceConnects, 2); + assert.equal(app.connectErrors.length, 2); + assert.deepEqual([...app.timers.values()], [400]); + }); -test("a transport error before the namespace connects retries once with backoff", async () => { - const app = await setup({ dropConnects: 2, dropReason: "transport error" }); - app.flush(); + test(`a connect${over} cancels the pending retry`, async () => { + const app = await setup({ socketio, dropConnects: 1 }); + app.flush(); + await app.connect(); + const socket = app.socket.current; + assert.equal(app.timers.size, 1); + + // Another path, such as the page becoming visible, reconnects first. + socket.reconnect(); + assert.equal(socket.connected, true); + assert.equal(app.timers.size, 0); + + socket.connected = false; + socket.handlers.get("disconnect")("io server disconnect"); + assert.equal(socket.namespaceConnects, 2); + assert.equal(app.timers.size, 0); + }); + + test(`a dropped connection${over} reconnects at once`, async () => { + const app = await setup({ socketio }); + app.flush(); + await app.connect(); + const socket = app.socket.current; + socket.connected = false; + socket.handlers.get("disconnect")("transport close"); + assert.equal(socket.namespaceConnects, 2); + assert.equal(app.connectErrors.length, 0); + assert.equal(app.timers.size, 0); + }); +} + +test("a socket.io transport error before the namespace connects retries once with backoff", async () => { + const app = await setup({ + socketio: true, + dropConnects: 2, + dropReason: "transport error", + }); await app.connect(); const socket = app.socket.current; assert.equal(socket.namespaceConnects, 1); @@ -479,46 +592,17 @@ test("a transport error before the namespace connects retries once with backoff" assert.deepEqual([...app.timers.values()], [400]); }); -test("a connect cancels the pending retry", async () => { - const app = await setup({ dropConnects: 1 }); - app.flush(); - await app.connect(); - const socket = app.socket.current; - assert.equal(app.timers.size, 1); - - // Another path, such as the page becoming visible, reconnects first. - socket.reconnect(); - assert.equal(socket.connected, true); - assert.equal(app.timers.size, 0); - - socket.connected = false; - socket.handlers.get("disconnect")("io server disconnect"); - assert.equal(socket.namespaceConnects, 2); - assert.equal(app.timers.size, 0); -}); - -test("a dropped connection reconnects at once", async () => { - const app = await setup(); - app.flush(); - await app.connect(); - const socket = app.socket.current; - socket.connected = false; - socket.handlers.get("disconnect")("transport close"); - assert.equal(socket.namespaceConnects, 2); - assert.equal(app.connectErrors.length, 0); - assert.equal(app.timers.size, 0); -}); - -for (const reason of ["error", "timeout", "pagehide", "hot reload"]) { +for (const reason of ["timeout", "pagehide", "hot reload"]) { test(`discard an unclaimed warm transport on ${reason}`, async () => { const app = await setup(); app.flush(); const warm = app.sockets[0]; - if (reason === "error") warm.io.openCallback(new Error("offline")); if (reason === "timeout") [...app.timers.keys()][0](); if (reason === "pagehide") app.window.dispatchEvent(new Event("pagehide")); if (reason === "hot reload") app.dispose(); assert.equal(warm.disconnects, 1); + // The transport's offline listener goes with it. + assert.equal(warm.offs, 1); assert.equal(app.timers.size, 0); await app.connect(); assert.equal(app.sockets.length, 2); @@ -526,14 +610,6 @@ for (const reason of ["error", "timeout", "pagehide", "hot reload"]) { }); } -test("a late warmup error cannot close the socket after React claims it", async () => { - const app = await setup(); - app.flush(); - await app.connect(); - app.socket.current.io.openCallback(new Error("late failure")); - assert.equal(app.socket.current.disconnects, 0); -}); - test("a synchronous mount does not leave a second speculative connection", async () => { const app = await setup(); await app.connect(); @@ -555,22 +631,48 @@ test("blocked session storage does not throw from speculative setup", async () = }; app.flush(); assert.equal(app.sockets.length, 0); - await assert.rejects(app.connect(), /storage blocked/); + // Reported, not rejected: the unhandled rejection handler adds an event, + // which reconnects, which would fail again without end. + await app.connect(); + assert.equal(app.sockets.length, 0); + assert.match(String(app.connectErrors[0]), /storage blocked/); + // A later attempt is not locked out by the failed one. + app.window.sessionStorage.getItem = () => "original-token"; + await app.connect(); + assert.equal(app.sockets.length, 1); }); -test("a caller selecting another transport gets its requested configuration", async () => { - const app = await setup(); - app.flush(); - await app.connect(["polling"]); - assert.equal(app.sockets[0].disconnects, 1); - assert.equal(app.socket.current.io.opts.transports[0], "polling"); +test("a socket.io client that fails to load is a connection error", async () => { + const app = await setup({ socketio: true, importFails: true }); + await app.connect(); + assert.equal(app.socket.current, null); + assert.equal(app.connectErrors.length, 1); + assert.match(String(app.connectErrors[0]), /dynamically imported module/); }); +for (const transport of ["websocket", "polling"]) { + test(`socket.io over ${transport} gets the engine.io transport and the boot event`, async () => { + // TRANSPORT is an engine.io name, handed to socket.io-client as is. + const app = await setup({ socketio: true, transport }); + app.flush(); + await app.connect([transport]); + assert.equal(app.sockets.length, 1); + assert.equal(app.socket.current.kind, "socketio"); + assert.equal(app.socket.current.io.opts.autoConnect, false); + assert.deepEqual([...app.socket.current.io.opts.transports], [transport]); + assert.equal(app.socket.current.namespaceConnects, 1); + assert.deepEqual(app.firstHydrates, [true]); + assert.equal(app.socket.current.auth.event.router_data.pathname, "/page"); + }); +} + for (const config of [ { browser: false }, { stateful: false }, { disabled: true }, { hidden: true }, + // The socket.io client is only loaded on connect. + { socketio: true }, ]) { test(`skip speculative connections for ${JSON.stringify(config)}`, async () => { const app = await setup(config); diff --git a/tests/units/compiler/test_websocket_js_template.py b/tests/units/compiler/test_websocket_js_template.py new file mode 100644 index 00000000000..b9b8acccd9d --- /dev/null +++ b/tests/units/compiler/test_websocket_js_template.py @@ -0,0 +1,26 @@ +"""Regression tests for the websocket.js frontend transport template.""" + +import shutil +import subprocess +from pathlib import Path + +import pytest + + +def test_websocket_session_lifecycle() -> None: + """Execute the transport's session open and reconnect tests against the real template.""" + node = shutil.which("node") + if node is None: + pytest.skip("Node.js is required for the frontend runtime tests") + result = subprocess.run( + [ + node, + "--experimental-vm-modules", + str(Path(__file__).with_name("websocket_js.test.mjs")), + ], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/units/compiler/websocket_js.test.mjs b/tests/units/compiler/websocket_js.test.mjs new file mode 100644 index 00000000000..12f8b2b8033 --- /dev/null +++ b/tests/units/compiler/websocket_js.test.mjs @@ -0,0 +1,590 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import test from "node:test"; +import vm from "node:vm"; + +const helpers = new URL( + "../../../packages/reflex-base/src/reflex_base/.templates/web/utils/helpers/", + import.meta.url, +); +const source = await readFile(new URL("websocket.js", helpers), "utf8"); +const jsonSource = await readFile(new URL("json.js", helpers), "utf8"); + +/** Evaluate the actual transport module against a scripted browser WebSocket. */ +async function setup({ + throwOnConstruct = false, + closeLater = false, + url = "http://localhost:8000/_event", +} = {}) { + const sockets = []; + const timers = new Map(); + const reported = []; + const listeners = new Map(); + + class WebSocket { + static CONNECTING = 0; + static OPEN = 1; + static CLOSING = 2; + static CLOSED = 3; + + constructor(url, protocols) { + if (throwOnConstruct) { + // Mixed content or a CSP violation fails the constructor itself. + throw new Error("SecurityError: insecure WebSocket from https"); + } + this.url = String(url); + this.protocols = protocols; + this.readyState = WebSocket.CONNECTING; + this.sent = []; + sockets.push(this); + } + + send(frame) { + assert.equal(this.readyState, WebSocket.OPEN); + this.sent.push(typeof frame === "string" ? JSON.parse(frame) : frame); + } + + close() { + if (closeLater) { + // A dead link: the browser reports the close only after its closing + // handshake times out, up to a minute later. + this.readyState = WebSocket.CLOSING; + return; + } + this.readyState = WebSocket.CLOSED; + this.onclose?.({ code: 1000, reason: "" }); + } + + // Test controls standing in for the network. + open() { + this.readyState = WebSocket.OPEN; + this.onopen?.(); + } + + receive(frame) { + this.onmessage({ data: JSON.stringify(frame) }); + } + + serverClose(code, reason = "") { + this.readyState = WebSocket.CLOSED; + this.onclose?.({ code, reason }); + } + } + + const context = vm.createContext({ + console, + URL, + URLSearchParams, + TextEncoder, + TextDecoder, + ArrayBuffer, + DataView, + Uint8Array, + WebSocket, + queueMicrotask, + SharedArrayBuffer, + reportError: (error) => reported.push(error), + setTimeout: (fn, ms) => { + timers.set(fn, ms); + return fn; + }, + clearTimeout: (fn) => timers.delete(fn), + addEventListener: (name, fn) => listeners.set(name, fn), + removeEventListener: (name, fn) => + listeners.get(name) === fn && listeners.delete(name), + }); + const module = new vm.SourceTextModule(source, { context }); + await module.link((specifier) => { + assert.equal(specifier, "$/utils/helpers/json"); + return new vm.SourceTextModule(jsonSource, { context }); + }); + await module.evaluate(); + const transport = new module.namespace.ReflexWebSocket(url, { + query: { token: "tok" }, + protocols: ["version"], + }); + const events = []; + for (const name of ["connect", "connect_error", "disconnect", "event"]) { + transport.on(name, (...args) => events.push([name, ...args])); + } + return { + transport, + sockets, + timers, + events, + reported, + listeners, + getChannel: module.namespace.getChannel, + decodeChannelFrame: module.namespace.decodeChannelFrame, + }; +} + +/** + * Connect the transport and complete its handshake. + * @param app The test setup. + * @param handshake Overrides for the handshake payload. + * @returns The connected scripted socket. + */ +function connectTransport(app, handshake = {}) { + app.transport.connect(); + const socket = app.sockets.at(-1); + socket.open(); + socket.receive(["_handshake", { ...HANDSHAKE[1], ...handshake }]); + return socket; +} + +const HANDSHAKE = [ + "_handshake", + { ping_interval: 25, ping_timeout: 120, protocol: 2, max_message_size: 1e6 }, +]; + +test("connect sends the connect frame with auth as soon as the socket opens", async () => { + const { transport, sockets, events } = await setup(); + assert.equal(sockets.length, 0, "the constructor must not dial"); + transport.auth = { event: { name: "boot" } }; + transport.connect(); + const [socket] = sockets; + assert.equal(new URL(socket.url).searchParams.get("token"), "tok"); + assert.deepEqual(socket.protocols, ["version"]); + socket.open(); + // No waiting for the handshake: it acknowledges this frame. + assert.deepEqual(socket.sent, [["_connect", { event: { name: "boot" } }]]); + assert.equal(transport.connected, false); + transport.emit("event", { name: "click" }); + assert.equal(socket.sent.length, 1, "events wait for the acknowledgement"); + socket.receive(HANDSHAKE); + assert.equal(transport.connected, true); + assert.deepEqual(socket.sent[1], ["event", { name: "click" }]); + assert.deepEqual(events, [["connect"]]); +}); + +test("a socket dialed ahead by open opens its session when connect claims it", async () => { + const { transport, sockets, timers } = await setup(); + transport.open(); + const [socket] = sockets; + socket.open(); + assert.deepEqual(socket.sent, []); + // The session timeout starts with the claim, not the dial. + assert.equal(timers.size, 0); + transport.auth = { event: { name: "boot" } }; + transport.connect(); + assert.equal(sockets.length, 1); + assert.deepEqual(socket.sent, [["_connect", { event: { name: "boot" } }]]); + assert.equal(timers.size, 1); +}); + +test("connect does nothing more while the session is opening", async () => { + const { transport, sockets } = await setup(); + transport.connect(); + transport.connect(); + sockets[0].open(); + transport.connect(); + assert.equal(sockets.length, 1); + assert.equal(sockets[0].sent.length, 1); +}); + +test("a closed session is reopened on a new socket with the current auth and token", async () => { + const { transport, sockets, events } = await setup(); + transport.connect(); + sockets[0].open(); + sockets[0].receive(HANDSHAKE); + sockets[0].close(); + assert.deepEqual(events.at(-1).slice(0, 2), [ + "disconnect", + "transport close", + ]); + transport.auth = { event: { name: "rehydrate" } }; + transport.io.opts.query = { token: "assigned" }; + transport.connect(); + assert.equal(sockets.length, 2); + assert.equal(new URL(sockets[1].url).searchParams.get("token"), "assigned"); + sockets[1].open(); + assert.deepEqual(sockets[1].sent, [ + ["_connect", { event: { name: "rehydrate" } }], + ]); +}); + +test("a session that is never acknowledged fails as a connect error", async () => { + const { transport, sockets, timers, events } = await setup(); + transport.connect(); + sockets[0].open(); + const [expire] = timers.keys(); + expire(); + assert.equal(sockets[0].readyState, sockets[0].constructor.CLOSED); + assert.equal(events.at(-1)[0], "connect_error"); + assert.equal(timers.size, 0); +}); + +/** + * The UTF-8 size of the connect frame carrying the transport's auth. + * @param transport The transport. + * @returns The size in bytes. + */ +const connectFrameSize = (transport) => + new TextEncoder().encode(JSON.stringify(["_connect", transport.auth])) + .byteLength; + +for (const [reason, named] of [ + ["max_message_size=100", " (100 bytes)"], + // An older backend, or a proxy with its own limit. + ["", ""], + ["Message Too Big", ""], +]) { + test(`a connect frame refused for its size with reason ${JSON.stringify(reason)} reports the hydrate`, async () => { + const { transport, sockets, events, timers } = await setup(); + transport.auth = { + event: { name: "boot", payload: { large: "é".repeat(80) } }, + }; + transport.connect(); + sockets[0].open(); + sockets[0].serverClose(1009, reason); + assert.equal(events.length, 1); + const [name, error] = events[0]; + assert.equal(name, "connect_error"); + assert.equal(error.code, "hydration_payload_too_large"); + assert.equal( + error.message, + `Hydration payload (${connectFrameSize(transport)} bytes) exceeds the ` + + `server limit${named}. Reduce browser storage data and reload.`, + ); + assert.equal(timers.size, 0); + }); +} + +test("other closes are a generic connect error before the handshake and a disconnect after it", async () => { + const app = await setup(); + app.transport.connect(); + app.sockets[0].open(); + app.sockets[0].serverClose(1008); + const [name, error] = app.events.at(-1); + assert.equal(name, "connect_error"); + assert.equal(error.message, "websocket connection failed"); + assert.equal(error.code, undefined); + + // A session's oversized message is not its hydrate. + const socket = connectTransport(app); + socket.serverClose(1009, "max_message_size=100"); + assert.deepEqual(app.events.at(-1).slice(0, 2), [ + "disconnect", + "transport close", + ]); +}); + +test("a reconnect over the limit an earlier handshake announced is refused before sending", async () => { + const app = await setup(); + const first = connectTransport(app, { max_message_size: 100 }); + first.serverClose(1001); + app.transport.auth = { event: { name: "rehydrate", large: "x".repeat(200) } }; + app.transport.connect(); + const socket = app.sockets.at(-1); + socket.open(); + assert.notEqual(socket, first); + assert.deepEqual(socket.sent, []); + assert.equal(socket.readyState, socket.constructor.CLOSED); + assert.equal(app.timers.size, 0); + assert.equal(app.events.filter(([n]) => n === "connect_error").length, 1); + const [name, error] = app.events.at(-1); + assert.equal(name, "connect_error"); + assert.equal(error.code, "hydration_payload_too_large"); + assert.equal( + error.message, + `Hydration payload (${connectFrameSize(app.transport)} bytes) exceeds ` + + "the server limit (100 bytes). Reduce browser storage data and reload.", + ); + + // One that fits is sent. + app.transport.auth = { event: { name: "rehydrate" } }; + app.transport.connect(); + app.sockets.at(-1).open(); + assert.deepEqual(app.sockets.at(-1).sent, [ + ["_connect", { event: { name: "rehydrate" } }], + ]); +}); + +test("the transport dials the event path with a trailing slash", async () => { + // Proxy rules written for Socket.IO route "/_event/*", which "/_event" misses. + const { transport, sockets } = await setup(); + transport.connect(); + assert.equal(new URL(sockets[0].url).pathname, "/_event/"); +}); + +test("a handler removing itself does not skip the next one", async () => { + const { transport } = await setup(); + const calls = []; + const once = () => { + calls.push("once"); + transport.off("custom", once); + }; + transport.on("custom", once); + transport.on("custom", () => calls.push("next")); + transport._emitLocal("custom"); + transport._emitLocal("custom"); + assert.deepEqual(calls, ["once", "next", "next"]); +}); + +test("a throwing channel handler cannot leave the transport half torn down", async () => { + const app = await setup(); + const socket = connectTransport(app); + const channel = app.getChannel("throws-on-disconnect"); + socket.receive(["_opened", null, "throws-on-disconnect"]); + channel.on("disconnect", () => { + throw new Error("component already unmounted"); + }); + app.transport.disconnect(); + assert.equal(app.transport.connected, false); + assert.equal(socket.readyState, socket.constructor.CLOSED); + assert.equal(app.events.at(-1)[0], "disconnect"); + assert.equal(app.reported.length, 1); + assert.match(app.reported[0].message, /already unmounted/); +}); + +test("the ping watchdog reports a dead link at once", async () => { + const app = await setup({ closeLater: true }); + connectTransport(app); + const channel = app.getChannel("watched"); + app.sockets[0].receive(["_opened", null, "watched"]); + const watchdog = [...app.timers.keys()].at(-1); + watchdog(); + // Not when the browser gives up on the closing handshake, a minute later. + assert.equal(app.transport.connected, false); + assert.deepEqual(app.events.at(-1).slice(0, 2), [ + "disconnect", + "ping timeout", + ]); + assert.equal(channel.connected, false); +}); + +test("going offline reports the disconnect at once", async () => { + // Localhost keeps working offline, so only a remote backend is watched. + const app = await setup({ + url: "https://example.com/_event", + closeLater: true, + }); + connectTransport(app); + app.listeners.get("offline")(); + assert.equal(app.transport.connected, false); + const [name, reason, details] = app.events.at(-1); + assert.deepEqual([name, reason], ["disconnect", "transport close"]); + assert.equal(details.description, "network connection lost"); + // Disposing of the transport releases the listener. + app.transport.off(); + assert.equal(app.listeners.size, 0); +}); + +test("a channel opened twice reports connect once", async () => { + const app = await setup(); + const socket = connectTransport(app); + const channel = app.getChannel("twice"); + let connects = 0; + channel.on("connect", () => connects++); + socket.receive(["_opened", null, "twice"]); + socket.receive(["_opened", null, "twice"]); + assert.equal(connects, 1); +}); + +test("a socket constructor that throws becomes a connect error", async () => { + const app = await setup({ throwOnConstruct: true }); + app.transport.connect(); + // Reported on a later tick, like any failed dial, once handlers exist. + assert.deepEqual(app.events, []); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(app.events.length, 1); + assert.equal(app.events[0][0], "connect_error"); + assert.match(app.events[0][1].message, /SecurityError/); + // A later attempt dials again rather than finding the transport stuck. + app.transport.connect(); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(app.events.length, 2); +}); + +test("a channel refuses names and attachments the backend cannot take", async () => { + const app = await setup(); + const socket = connectTransport(app); + const channel = app.getChannel("strict"); + socket.receive(["_opened", null, "strict"]); + const before = socket.sent.length; + for (const name of [ + undefined, + 42, + "", + "_open", + "_custom", + "connect", + "disconnect", + "error", + ]) { + assert.throws(() => channel.emit(name, null)); + } + for (const buffers of [ + new ArrayBuffer(8), + [{ size: 3 }], + [[1, 2, 3]], + ["abc"], + ]) { + assert.throws(() => channel.emit("push", null, buffers), { + name: "TypeError", + }); + } + for (const name of [undefined, "has space", "x".repeat(65)]) { + assert.throws(() => app.getChannel(name), { name: "TypeError" }); + } + // Nothing refused was sent or queued behind the message that is valid. + channel.emit("push", null); + assert.deepEqual(socket.sent.slice(before), [["push", null, "strict"]]); +}); + +test("shared buffers are sent whole", async () => { + const app = await setup(); + const socket = connectTransport(app); + const channel = app.getChannel("shared"); + socket.receive(["_opened", null, "shared"]); + channel.emit("push", null, [new SharedArrayBuffer(8)]); + const [, , , buffers] = app.decodeChannelFrame(socket.sent.at(-1)); + assert.equal(buffers[0].byteLength, 8); +}); + +test("an oversized event is dropped instead of costing the socket", async () => { + const app = await setup(); + const socket = connectTransport(app, { max_message_size: 100 }); + const before = socket.sent.length; + app.transport.emit("event", { blob: "x".repeat(200) }); + assert.equal(socket.sent.length, before); + assert.equal(app.reported.length, 1); + assert.match(app.reported[0].message, /over the 100/); + app.transport.emit("event", { small: true }); + assert.equal(socket.sent.length, before + 1); +}); + +test("an event queued before the handshake is checked against the limit it brings", async () => { + const app = await setup(); + app.transport.connect(); + const socket = app.sockets.at(-1); + socket.open(); + // Emitted before the backend has said what it accepts. + app.transport.emit("event", { name: "state.big", blob: "x".repeat(200) }); + app.transport.emit("event", { small: true }); + socket.receive(["_handshake", { ...HANDSHAKE[1], max_message_size: 100 }]); + assert.deepEqual(socket.sent.slice(1), [["event", { small: true }]]); + assert.equal(app.reported.length, 1); + assert.match( + app.reported[0].message, + /"state\.big" is \d+ bytes, over the 100/, + ); +}); + +test("a channel message queued on a closing socket is checked against the next backend's limit", async () => { + const app = await setup(); + let socket = connectTransport(app); + const channel = app.getChannel("big"); + const errors = []; + channel.on("error", (error) => errors.push(error)); + socket.receive(["_opened", null, "big"]); + // The socket is closing, but the browser has not reported the close yet. + socket.readyState = 2; + channel.emit("push", { blob: "x".repeat(200) }); + socket.readyState = 3; + socket.onclose({ code: 1006, reason: "" }); + + // The reconnect lands on a backend that accepts less. + socket = connectTransport(app, { max_message_size: 100 }); + assert.deepEqual( + socket.sent.map((frame) => frame[0]), + ["_connect", "_open"], + ); + assert.deepEqual( + errors.map((error) => error.code), + ["message_too_large"], + ); +}); + +test("a channel message over the backend's limit is refused, measured in UTF-8 bytes", async () => { + const app = await setup(); + const socket = connectTransport(app, { max_message_size: 1024 }); + const channel = app.getChannel("sized"); + socket.receive(["_opened", null, "sized"]); + const before = socket.sent.length; + assert.throws( + () => channel.emit("push", { blob: "x".repeat(5000) }), + /"push" is \d+ bytes, over the 1024 the backend accepts \(REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE\)\.$/, + ); + // 400 characters, 1200 UTF-8 bytes: only a byte-accurate check catches it. + assert.throws( + () => channel.emit("push", "\u20ac".repeat(400)), + (error) => { + const bytes = Number(/is (\d+) bytes/.exec(error.message)[1]); + return bytes > 1024 && bytes < 4 * 400; + }, + ); + channel.emit("push", { small: true }); + assert.deepEqual(socket.sent.slice(before), [ + ["push", { small: true }, "sized"], + ]); +}); + +test("a channel message emitted before any limit was known is checked when the channel opens", async () => { + const app = await setup(); + const channel = app.getChannel("early"); + const errors = []; + channel.on("error", (error) => errors.push(error)); + // No transport yet, so no limit to check against. + channel.emit("push", { blob: "y".repeat(5000) }); + channel.emit("push", { small: true }); + const socket = connectTransport(app, { max_message_size: 1024 }); + socket.receive(["_opened", null, "early"]); + assert.deepEqual(socket.sent.slice(1), [ + ["_open", null, "early"], + ["push", { small: true }, "early"], + ]); + assert.equal(errors.length, 1); + assert.equal(errors[0].code, "message_too_large"); + assert.match( + errors[0].message, + /^A queued channel message is \d+ bytes, over the 1024 the backend accepts \(REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE\); it was dropped\.$/, + ); +}); + +test("channel frames still queued are dropped on a backend without channels", async () => { + const app = await setup(); + let socket = connectTransport(app); + const channel = app.getChannel("probe"); + const errors = []; + channel.on("error", (error) => errors.push(error.code)); + socket.receive(["_opened", null, "probe"]); + // The socket is closing, but the browser has not reported the close yet, so + // these frames queue on the transport. + socket.readyState = socket.constructor.CLOSING; + app.transport.emit("event", { name: "state.event" }); + channel.emit("push", { n: 1 }); + channel.emit("frame", { n: 2 }, [new Uint8Array([1, 2, 3])]); + socket.readyState = socket.constructor.CLOSED; + socket.onclose({ code: 1006, reason: "" }); + + // The reconnect lands on a backend that closes the socket over either. + socket = connectTransport(app, { protocol: 1 }); + assert.deepEqual(socket.sent.slice(1), [["event", { name: "state.event" }]]); + assert.equal(app.transport._sendQueue.length, 0); + assert.deepEqual(errors, ["channels_unsupported"]); +}); + +test("an _opened from a transport the channel has left changes nothing", async () => { + const app = await setup(); + const old = connectTransport(app); + const channel = app.getChannel("moved"); + let connects = 0; + channel.on("connect", () => connects++); + old.close(); + channel.emit("push", { n: 1 }); + // A straggler answering the open the closed socket carried. + old.receive(["_opened", null, "moved"]); + assert.equal(channel.connected, false); + assert.equal(connects, 0); + + // The queued message is owed to the transport the channel attaches to next. + const socket = connectTransport(app); + socket.receive(["_opened", null, "moved"]); + assert.equal(connects, 1); + assert.deepEqual(socket.sent.slice(1), [ + ["_open", null, "moved"], + ["push", { n: 1 }, "moved"], + ]); +}); diff --git a/tests/units/conftest.py b/tests/units/conftest.py index e4e83140c64..090ed0b7a9a 100644 --- a/tests/units/conftest.py +++ b/tests/units/conftest.py @@ -102,6 +102,57 @@ def mock_app(app_module_mock: mock.Mock, app: App) -> App: return app +@pytest.fixture +def namespace_app(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: + """A mock app for an event namespace, with a local token manager. + + Redis is disabled so token linking cannot leak into a shared Redis. + + Args: + monkeypatch: The pytest monkeypatch fixture. + + Returns: + The mock app. + """ + monkeypatch.setattr(prerequisites, "check_redis_used", lambda: False) + app = mock.Mock() + app._state = None + app._channels = {} + app.router = mock.Mock(return_value=None) + app.event_processor.enqueue = mock.AsyncMock() + return app + + +@pytest.fixture +def linked_state(namespace_app: mock.Mock) -> mock.Mock: + """The state a namespace app loads when a connect links its token. + + Each load is recorded on ``namespace_app.state_manager.modify_state``. + + Args: + namespace_app: The mock app. + + Returns: + The state. + """ + namespace_app._state = mock.Mock() + state = mock.Mock(router_data={}) + namespace_app.state_manager.modify_state = mock.Mock( + return_value=mock.AsyncMock(__aenter__=mock.AsyncMock(return_value=state)) + ) + return state + + +@pytest.fixture +def boot_event() -> dict[str, Any]: + """The hydrate event a frontend sends with its connect. + + Returns: + The event payload. + """ + return {"name": "state.hydrate_and_load", "payload": {}, "router_data": {}} + + @pytest.fixture(scope="session") def windows_platform() -> bool: """Check if system is windows. diff --git a/tests/units/reflex_base/client_event_queue_runtime.mjs b/tests/units/reflex_base/client_event_queue_runtime.mjs index 2befa24459f..70c50617cbf 100644 --- a/tests/units/reflex_base/client_event_queue_runtime.mjs +++ b/tests/units/reflex_base/client_event_queue_runtime.mjs @@ -24,7 +24,6 @@ export async function createQueueRuntime( localStorage: { clear() {}, removeItem() {} }, sessionStorage: { clear() {}, removeItem() {} }, }, - "socket.io-client": { default: unused }, mergician: { mergician: unused }, "$/env.json": { default: {} }, "$/reflex.json": { default: {} }, @@ -51,22 +50,23 @@ export async function createQueueRuntime( eventLoop: { addEvents: unused }, }, "$/utils/helpers/debounce": { default: unused }, - "$/utils/helpers/json": { parseJson: unused }, + "$/utils/helpers/json": { + parseJsonLenient: unused, + undefinedToNull: unused, + }, "$/utils/helpers/throttle": { default: throttle ?? unused }, "$/utils/helpers/upload": { uploadFiles: uploadFiles ?? unused, }, + "$/utils/helpers/websocket": { + ReflexWebSocket: unused, + disableChannels: unused, + getChannel: unused, + }, ...dependencyOverrides, }; - // Let Node parse the unchanged module; only expose private state to tests. - const module = new SourceTextModule( - `import { window, document, localStorage, sessionStorage } from "test:browser"; -${source} -export { event_queue }; -`, - ); const linked = new Map(); - await module.link((specifier) => { + const link = (specifier) => { if (!linked.has(specifier)) { const exports = dependencies[specifier]; if (!exports) @@ -81,7 +81,29 @@ export { event_queue }; ); } return linked.get(specifier); - }); + }; + // Let Node parse the unchanged module; only expose private state to tests. + const module = new SourceTextModule( + `import { window, document, localStorage, sessionStorage } from "test:browser"; +${source} +export { event_queue }; +`, + { + // The socket.io client is imported on demand. + async importModuleDynamically(specifier) { + const imported = link(specifier); + // Newer Node versions create synthetic modules already linked. + if (imported.status === "unlinked") { + await imported.link(link); + } + if (imported.status === "linked") { + await imported.evaluate(); + } + return imported; + }, + }, + ); + await module.link(link); await module.evaluate(); return module.namespace; } diff --git a/tests/units/reflex_base/client_state_updates.mjs b/tests/units/reflex_base/client_state_updates.mjs index 4eabef315bf..fe1518b4e68 100644 --- a/tests/units/reflex_base/client_state_updates.mjs +++ b/tests/units/reflex_base/client_state_updates.mjs @@ -28,11 +28,13 @@ const element = (tag) => ({ async function createConnection( t, dispatch = {}, - { isDevMode = false, lastReload } = {}, + { isDevMode = false, lastReload, socketio = false } = {}, ) { const warn = t.mock.method(console, "warn", () => {}); const handlers = new Map(); const sent = []; + // The kind of each transport the runtime created. + const created = []; const saved = new Map(); const reloads = []; const session = new Map([["token", "test-token"]]); @@ -90,8 +92,24 @@ async function createConnection( localStorage: { setItem: (key, value) => saved.set(key, value) }, sessionStorage, }, - "socket.io-client": { default: () => socket }, - "$/env.json": { default: { EVENT: "http://backend/_event" } }, + "$/utils/helpers/websocket": { + // Called with `new`: a constructor returning an object yields that object. + ReflexWebSocket: function () { + created.push("websocket"); + return socket; + }, + disableChannels() {}, + getChannel() {}, + }, + "socket.io-client": { + default: () => { + created.push("socketio"); + return socket; + }, + }, + "$/env.json": { + default: { EVENT: "http://backend/_event", SOCKETIO: socketio }, + }, "$/utils/context-registry": { app, eventLoop: {} }, }, }); @@ -107,6 +125,7 @@ async function createConnection( navigate, params, ); + assert.deepEqual(created, [socketio ? "socketio" : "websocket"]); return { sent, saved, @@ -133,105 +152,133 @@ async function createConnection( }; } -test("an unknown substate does not block known state or events", async (t) => { - const applied = []; - const c = await createConnection(t, { - known: (value) => applied.push(value), - }); - const callback = { name: "test.callback", payload: {} }; - c.receive({ - delta: Object.freeze({ unknown: { value: "u" }, known: { value: "k" } }), - events: [callback], - }); - assert.deepEqual(applied, [{ value: "k" }]); - assert.deepEqual([...c.saved], [["known.value", "k"]]); - assert.ok(c.sent.some((event) => event.value === callback)); - assert.deepEqual( - c.reports().map((event) => event.value.substate), - ["unknown"], - ); -}); +/** + * Register a test once per transport. + * @param name The test name. + * @param fn The test body, given a `connect` that opens a connection over the + * transport. + */ +const transportTest = (name, fn) => { + for (const socketio of [false, true]) { + test(`${name} (${socketio ? "socket.io" : "websocket"})`, (t) => + fn((dispatch, options) => + createConnection(t, dispatch, { ...options, socketio }), + )); + } +}; -test("an unknown substate is reported once per connection", async (t) => { - const c = await createConnection(t); - const update = { delta: { unknown: {} } }; - c.receive(update); - c.receive(update); - assert.equal(c.reports().length, 1); - await c.reconnect(); - c.receive(update); - assert.equal(c.reports().length, 2); -}); +transportTest( + "an unknown substate does not block known state or events", + async (connect) => { + const applied = []; + const c = await connect({ + known: (value) => applied.push(value), + }); + const callback = { name: "test.callback", payload: {} }; + c.receive({ + delta: Object.freeze({ unknown: { value: "u" }, known: { value: "k" } }), + events: [callback], + }); + assert.deepEqual(applied, [{ value: "k" }]); + assert.deepEqual([...c.saved], [["known.value", "k"]]); + assert.ok(c.sent.some((event) => event.value === callback)); + assert.deepEqual( + c.reports().map((event) => event.value.substate), + ["unknown"], + ); + }, +); -test("the state schema rides in the CONNECT auth on connect and reconnect", async (t) => { - const c = await createConnection(t); - assert.equal(c.socket.auth.state_schema, FRONTEND_SCHEMA); - assert.equal(c.socket.auth.event.name, "test.hydrate"); - assert.deepEqual(c.socket.auth.event.payload.hashes, []); +transportTest( + "an unknown substate is reported once per connection", + async (connect) => { + const c = await connect(); + const update = { delta: { unknown: {} } }; + c.receive(update); + c.receive(update); + assert.equal(c.reports().length, 1); + await c.reconnect(); + c.receive(update); + assert.equal(c.reports().length, 2); + }, +); - c.reconnectSocket(); - assert.equal(c.socket.connected, true); - assert.equal(c.socket.auth.state_schema, FRONTEND_SCHEMA); - assert.equal(c.socket.auth.event.payload.hashes, undefined); -}); +transportTest( + "the state schema rides in the connect auth on connect and reconnect", + async (connect) => { + const c = await connect(); + assert.equal(c.socket.auth.state_schema, FRONTEND_SCHEMA); + assert.equal(c.socket.auth.event.name, "test.hydrate"); + assert.deepEqual(c.socket.auth.event.payload.hashes, []); -test("a mismatch in prod reloads once per cooldown, then shows a banner", async (t) => { - const c = await createConnection(t); - c.mismatch(); - assert.equal(c.reloads.length, 1); - assert.ok(Date.now() - Number(c.session.get(LAST_RELOADED_KEY)) < 1000); - assert.equal(c.banner(), undefined); + c.reconnectSocket(); + assert.equal(c.socket.connected, true); + assert.equal(c.socket.auth.state_schema, FRONTEND_SCHEMA); + assert.equal(c.socket.auth.event.payload.hashes, undefined); + }, +); - // The reload landed on the same skew (a cached bundle, or several backend - // versions behind one load balancer): no loop, a banner instead. - c.mismatch(); - c.mismatch(); - assert.equal(c.reloads.length, 1); - const banner = c.banner(); - assert.ok(banner); - assert.equal(banner.attributes.role, "alert"); - assert.match(banner.children[0], /new version/i); - assert.equal(c.warn.mock.callCount(), 2); - const button = banner.children.find((node) => node.tag === "button"); - assert.equal(button.textContent, "Reload"); - - // The button reloads and restarts the cooldown. - c.session.set(LAST_RELOADED_KEY, "1"); - button.onclick(); - assert.equal(c.reloads.length, 2); - assert.ok(Date.now() - Number(c.session.get(LAST_RELOADED_KEY)) < 1000); - - // Detection alone does not stop events: the page keeps working as it can. - await c.enqueue([{ name: "test.event", payload: { id: 1 } }]); - assert.deepEqual(c.sent.at(-1).value.payload, { id: 1 }); -}); +transportTest( + "a mismatch in prod reloads once per cooldown, then shows a banner", + async (connect) => { + const c = await connect(); + c.mismatch(); + assert.equal(c.reloads.length, 1); + assert.ok(Date.now() - Number(c.session.get(LAST_RELOADED_KEY)) < 1000); + assert.equal(c.banner(), undefined); -test("a reload inside the cooldown window is skipped, outside it is taken", async (t) => { - const recent = await createConnection( - t, - {}, - { - lastReload: Date.now() - 1000, - }, - ); - recent.mismatch(10000); - assert.equal(recent.reloads.length, 0); - assert.ok(recent.banner()); - - const stale = await createConnection( - t, - {}, - { - lastReload: Date.now() - 20000, - }, - ); - stale.mismatch(10000); - assert.equal(stale.reloads.length, 1); - assert.equal(stale.banner(), undefined); -}); + // The reload landed on the same skew (a cached bundle, or several backend + // versions behind one load balancer): no loop, a banner instead. + c.mismatch(); + c.mismatch(); + assert.equal(c.reloads.length, 1); + const banner = c.banner(); + assert.ok(banner); + assert.equal(banner.attributes.role, "alert"); + assert.match(banner.children[0], /new version/i); + assert.equal(c.warn.mock.callCount(), 2); + const button = banner.children.find((node) => node.tag === "button"); + assert.equal(button.textContent, "Reload"); + + // The button reloads and restarts the cooldown. + c.session.set(LAST_RELOADED_KEY, "1"); + button.onclick(); + assert.equal(c.reloads.length, 2); + assert.ok(Date.now() - Number(c.session.get(LAST_RELOADED_KEY)) < 1000); + + // Detection alone does not stop events: the page keeps working as it can. + await c.enqueue([{ name: "test.event", payload: { id: 1 } }]); + assert.deepEqual(c.sent.at(-1).value.payload, { id: 1 }); + }, +); + +transportTest( + "a reload inside the cooldown window is skipped, outside it is taken", + async (connect) => { + const recent = await connect( + {}, + { + lastReload: Date.now() - 1000, + }, + ); + recent.mismatch(10000); + assert.equal(recent.reloads.length, 0); + assert.ok(recent.banner()); + + const stale = await connect( + {}, + { + lastReload: Date.now() - 20000, + }, + ); + stale.mismatch(10000); + assert.equal(stale.reloads.length, 1); + assert.equal(stale.banner(), undefined); + }, +); -test("a mismatch in dev only warns", async (t) => { - const c = await createConnection(t, {}, { isDevMode: true }); +transportTest("a mismatch in dev only warns", async (connect) => { + const c = await connect({}, { isDevMode: true }); c.mismatch(); c.mismatch(); assert.equal(c.warn.mock.callCount(), 2); diff --git a/tests/units/reflex_bench/drivers/test_browser.py b/tests/units/reflex_bench/drivers/test_browser.py index 71ef592f0b7..eb9c644ece4 100644 --- a/tests/units/reflex_bench/drivers/test_browser.py +++ b/tests/units/reflex_bench/drivers/test_browser.py @@ -19,14 +19,14 @@ from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from typing import Any, TypeVar +from typing import Any, TypeVar, cast import psutil import pytest from reflex_bench import machine from reflex_bench.drivers import browser as browser_module from reflex_bench.drivers.app_process import Readiness -from reflex_bench.drivers.browser import Anchor, Browser +from reflex_bench.drivers.browser import Anchor, Browser, Tab from reflex_bench.scheduler import ABANDON_GRACE_S HTML = Path(__file__).parents[1] / "fixtures" / "html" @@ -162,6 +162,37 @@ def test_the_page_script_is_one_file_next_to_the_driver(): assert script.read_text(encoding="utf-8") == browser_module.BENCH_JS +class _Events: + """A page or DevTools session that ignores what a tab asks of it.""" + + def on(self, *args: Any) -> None: + """Ignore a listener.""" + + def send(self, *args: Any) -> None: + """Ignore a command.""" + + +def test_received_frames_count_the_messages_of_both_protocols(): + tab = Tab(cast(Any, _Events()), cast(Any, _Events())) + frames: list[str | bytes] = [ + '["event",{"delta":{}}]', + '["_ping"]', + '42/_event,["event",{"delta":{}}]', + '42["reload","/"]', + # Socket.IO's open, ack and ping, and binary frames carry no message. + '0{"sid":"a"}', + '40/_event,{"sid":"a"}', + "2", + b"\x00\x01", + ] + for frame in frames: + tab._on_frame_received(frame) + assert tab.ws_events == {"event": 2, "_ping": 1, "reload": 1} + assert tab.ws_bytes == sum( + len(frame.encode() if isinstance(frame, str) else frame) for frame in frames + ) + + @needs_chromium def test_marks_carry_both_clocks_and_map_onto_the_harness_clock( owner: ThreadPoolExecutor, browser: Browser, site: str diff --git a/tests/units/reflex_bench/drivers/test_echo_server.py b/tests/units/reflex_bench/drivers/test_echo_server.py index 239e35780cd..0e23adbdddd 100644 --- a/tests/units/reflex_bench/drivers/test_echo_server.py +++ b/tests/units/reflex_bench/drivers/test_echo_server.py @@ -1,19 +1,116 @@ """Tests for reflex_bench.drivers.echo_server. The echo server's answers are exercised by the generator's tests; these check -the process that the calibration runs it in. +its handshake against reflex's and the process that the calibration runs it in. """ from __future__ import annotations +import asyncio +import json +from typing import Any, cast +from unittest.mock import AsyncMock, Mock + import psutil import pytest +from reflex_base.environment import environment from reflex_bench.drivers import events -from reflex_bench.drivers.echo_server import EchoProcess +from reflex_bench.drivers.echo_server import EchoProcess, EchoServer from reflex_bench.drivers.events import Endpoint, EventShape, LoadPlan, run_load +from reflex_bench.suites.events import SEQ_VAR +from starlette.websockets import WebSocket +from websockets.asyncio.server import ServerConnection + +from reflex import event_namespace +from reflex.utils import prerequisites +from tests.units.reflex_bench.factories import WEBSOCKET STATE = "reflex___state____state.playground___state____bench_state" -SEQ_VAR = "last_seq_rx_state_" + + +class _Socket: + """Just enough of a starlette websocket for reflex to open a plain session.""" + + def __init__(self) -> None: + """Queue a connect frame without a boot event, then the disconnect.""" + self.scope = {"type": "websocket", "query_string": b"token=tok", "headers": []} + self.headers: dict[str, str] = {} + self.sent: list[Any] = [] + self._received = [ + {"type": "websocket.receive", "text": WEBSOCKET.join}, + {"type": "websocket.disconnect", "code": 1000}, + ] + + async def accept(self, subprotocol: str | None = None) -> None: + """Accept the connection. + + Args: + subprotocol: The chosen subprotocol. + """ + + async def send_text(self, text: str) -> None: + """Record a frame. + + Args: + text: The frame. + """ + self.sent.append(json.loads(text)) + + async def receive(self) -> dict[str, Any]: + """Take the next queued message once the writer ran, as the network lets it. + + Returns: + The ASGI message. + """ + await asyncio.sleep(0) + return self._received.pop(0) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + """Close the connection. + + Args: + code: The close code. + reason: The close reason. + """ + + +class _Sent(list[str]): + """A server connection that records what is sent on it.""" + + async def send(self, frame: str) -> None: + """Record a frame. + + Args: + frame: The frame. + """ + self.append(frame) + + +def test_the_handshake_is_copied_from_reflex(monkeypatch: pytest.MonkeyPatch): + # reflex_bench drives releases without reflex.event_namespace, so it keeps a + # copy of the handshake a default-configured app sends. + for var in ( + environment.REFLEX_SOCKET_INTERVAL, + environment.REFLEX_SOCKET_TIMEOUT, + environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE, + ): + monkeypatch.delenv(var.name, raising=False) + monkeypatch.setattr(prerequisites, "check_redis_used", lambda: False) + app = Mock(_state=None, _channels={}, router=Mock(return_value=None)) + app.event_processor.enqueue = AsyncMock() + namespace = event_namespace.WebsocketEventNamespace(events.NAMESPACE, app) + socket = _Socket() + asyncio.run(namespace.handle_websocket(cast(WebSocket, socket))) + reflex_name, reflex_settings = socket.sent[0] + sent = _Sent() + server = EchoServer(delta_key=STATE, seq_var=SEQ_VAR) + asyncio.run(server.on_connect(cast(ServerConnection, sent), "sid")) + (frame,) = sent + message = WEBSOCKET.parse(frame) + assert message is not None + name, settings = message + assert name == reflex_name == events.HANDSHAKE + assert settings == reflex_settings def test_echo_process_answers_a_load_and_stops(): diff --git a/tests/units/reflex_bench/drivers/test_events.py b/tests/units/reflex_bench/drivers/test_events.py index 2d3b59ecc18..958e320273b 100644 --- a/tests/units/reflex_bench/drivers/test_events.py +++ b/tests/units/reflex_bench/drivers/test_events.py @@ -1,30 +1,32 @@ """Tests for reflex_bench.drivers.events. The generator runs against :class:`~reflex_bench.drivers.echo_server.EchoServer` -on a thread of the test process; subclasses script delays, dropped and reordered -answers, a new token, a missing hydration, a disconnect and a stall. Most tests -drive the sessions on the calling thread, as one generator process does; the -multi-process ones go through :class:`~reflex_bench.drivers.events.LoadRunner`. +on a thread of the test process, on the plain protocol unless a test names +Socket.IO; subclasses script delays, dropped and reordered answers, a new token, +a missing hydration, a disconnect and a stall. Most tests drive the sessions on +the calling thread, as one generator process does; the multi-process ones go +through :class:`~reflex_bench.drivers.events.LoadRunner`. """ from __future__ import annotations import asyncio -import contextlib import dataclasses import json import multiprocessing +import socket import threading import time -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, cast import pytest +from reflex_base.constants import CompileVars, SocketEvent from reflex_bench.drivers import events from reflex_bench.drivers.app_process import free_ports from reflex_bench.drivers.echo_server import EchoServer from reflex_bench.drivers.events import ( - EVENT_PREFIX, + CODECS, Endpoint, EventShape, GeneratorSaturated, @@ -33,62 +35,45 @@ LoadPlan, LoadResult, LoadRunner, + WireProtocol, open_sessions, run_load, ) +from reflex_bench.suites.events import CLIENT_VAR, SEQ_VAR from websockets.asyncio.server import ServerConnection -from tests.units.reflex_bench.factories import make_load_result +from reflex import event_namespace +from tests.units.reflex_bench.factories import ( + PROTOCOLS, + SOCKETIO, + WEBSOCKET, + make_load_result, + serving, +) STATE = "reflex___state____state.playground___state____bench_state" -SEQ_VAR = "last_seq_rx_state_" SHAPE = EventShape( name=f"{STATE}.set_seq", payload=events.seq_payload, delta_key=STATE, seq_var=SEQ_VAR, ) +# The hydration events of a session, by protocol. +HYDRATION = { + "websocket": [events.HYDRATE_AND_LOAD_EVENT], + "socketio": [events.HYDRATE_EVENT, events.ON_LOAD_EVENT], +} +# The events whose answer sets is_hydrated. +HYDRATING = {events.ON_LOAD_EVENT, events.HYDRATE_AND_LOAD_EVENT} # Captured from the playground: HEAD leaves empty fields out, 0.8.23 does not. -HEAD_REPLY = f'42/_event,["event",{{"delta":{{"{STATE}":{{"{SEQ_VAR}":7}}}}}}]' +HEAD_REPLY = f'["event",{{"delta":{{"{STATE}":{{"{SEQ_VAR}":7}}}}}}]' +SOCKETIO_REPLY = f'42/_event,["event",{{"delta":{{"{STATE}":{{"{SEQ_VAR}":7}}}}}}]' OLD_REPLY = ( f'42/_event,["event",{{"delta":{{"{STATE}":{{"{SEQ_VAR}":7}}}},' '"events":[],"final":true}]' ) -@contextlib.contextmanager -def serving(server: EchoServer) -> Iterator[str]: - """Run an echo server on a thread with its own event loop. - - Args: - server: The server. - - Yields: - Its base URL, like the backend URL of a started app. - """ - loop = asyncio.new_event_loop() - ready = threading.Event() - port: list[int] = [] - done: list[asyncio.Future[None]] = [] - - async def main() -> None: - done.append(loop.create_future()) - async with server.serve("127.0.0.1", 0) as bound: - port.append(bound) - ready.set() - await done[0] - - thread = threading.Thread(target=loop.run_until_complete, args=(main(),)) - thread.start() - assert ready.wait(10) - try: - yield f"http://127.0.0.1:{port[0]}" - finally: - loop.call_soon_threadsafe(done[0].set_result, None) - thread.join(10) - loop.close() - - def plan(url: str, **overrides: Any) -> LoadPlan: values: dict[str, Any] = { "endpoint": Endpoint(url), @@ -136,12 +121,27 @@ def start(errors: list[str]) -> Awaitable[int | None]: class Scripted(EchoServer): """An echo server that records the seq events and asks a hook about each.""" - def __init__(self) -> None: - """Record nothing yet.""" - super().__init__(delta_key=STATE, seq_var=SEQ_VAR) + def __init__(self, protocol: WireProtocol = "websocket") -> None: + """Record nothing yet. + + Args: + protocol: The protocol to speak. + """ + super().__init__(delta_key=STATE, seq_var=SEQ_VAR, protocol=protocol) self.seqs: list[int] = [] self.tokens: list[str] = [] + async def end(self, ws: ServerConnection) -> None: + """End a session as reflex does: Socket.IO leaves the namespace, the plain protocol closes with a code. + + Args: + ws: The connection. + """ + if (leave := self.codec.leave) is not None: + await ws.send(leave) + else: + await ws.close(1008) + async def on_event(self, ws: ServerConnection, event: dict[str, Any]) -> None: """Record a seq event, then answer it unless the hook says otherwise. @@ -173,39 +173,155 @@ async def seq_event( return True -def test_event_url(): +def test_the_plain_protocol_is_copied_from_reflex(): + # reflex_bench drives releases without reflex.event_namespace, so it keeps a copy. + assert (events.CONNECT, events.HANDSHAKE, events.PING, events.PONG) == ( + event_namespace.CONNECT_MESSAGE, + event_namespace.HANDSHAKE_MESSAGE, + event_namespace.PING_MESSAGE, + event_namespace.PONG_MESSAGE, + ) + assert events.EVENT == SocketEvent.EVENT + assert events.BOOT_KEY == CompileVars.CONNECT_AUTH_EVENT + + +def test_event_urls(): assert ( - events.event_url("http://localhost:8000", "tok") + WEBSOCKET.url("http://localhost:8000", "tok") + == "ws://localhost:8000/_event?token=tok" + ) + assert ( + SOCKETIO.url("http://localhost:8000", "tok") == "ws://localhost:8000/_event/?EIO=4&transport=websocket&token=tok" ) - assert events.event_url("https://app.example.com/", "tok").startswith( - "wss://app.example.com/_event/?" + assert ( + WEBSOCKET.url("https://app.example.com/", "tok") + == "wss://app.example.com/_event?token=tok" + ) + assert SOCKETIO.url("https://app.example.com/base/", "tok").startswith( + "wss://app.example.com/base/_event/?" ) def test_event_frames_always_carry_the_token(): - # 0.8.23 rejects an event without a token and HEAD ignores it, so every - # frame carries it. - frame = events.event_frame(f"{STATE}.set_seq", {"seq": 7}, token="tok") + # 0.8.23 rejects an event without a token and later releases ignore it, + # so every frame carries it. + frame = WEBSOCKET.event_frame(f"{STATE}.set_seq", {"seq": 7}, token="tok") assert frame == ( - f'42/_event,["event",{{"name":"{STATE}.set_seq","payload":{{"seq":7}},' + f'["event",{{"name":"{STATE}.set_seq","payload":{{"seq":7}},' '"router_data":{"pathname":"/","asPath":"/","query":{}},"token":"tok"}]' ) - frames = events._Frames(SHAPE, "tok", "/", index=3) - assert frames.event(7) == frame - name, event = json.loads(frames.event(8)[len(EVENT_PREFIX) :]) - assert (name, event["payload"]) == ("event", {"seq": 8}) - # A shape with a client var sends the session's index along. - contended = events._Frames(CONTENDED_SHAPE, "tok", "/", index=3) - _, event = json.loads(contended.event(8)[len(EVENT_PREFIX) :]) - assert event["payload"] == {"seq": 8, "client": 3} + socketio = SOCKETIO.event_frame(f"{STATE}.set_seq", {"seq": 7}, token="tok") + assert socketio == f"42/_event,{frame}" + for codec in (WEBSOCKET, SOCKETIO): + frames = events._Frames(codec, SHAPE, "tok", "/", index=3) + assert frames.event(7) == codec.event_frame( + f"{STATE}.set_seq", {"seq": 7}, token="tok" + ) + message = codec.parse(frames.event(8)) + assert message is not None + name, event = message + assert (name, event["payload"]) == ("event", {"seq": 8}) + # A shape with a client var sends the session's index along. + contended = events._Frames(codec, CONTENDED_SHAPE, "tok", "/", index=3) + message = codec.parse(contended.event(8)) + assert message is not None + assert message[1]["payload"] == {"seq": 8, "client": 3} + + +def test_the_plain_connect_frame_carries_the_hydration_as_its_boot_event(): + frame = WEBSOCKET.connect_frame( + token="tok", pathname="/item/42", query={"item_id": "42"} + ) + assert frame == ( + '["_connect",{"event":{"name":"reflex___state____state.hydrate_and_load",' + '"payload":{},"router_data":{"pathname":"/item/42","asPath":"/item/42",' + '"query":{"item_id":"42"}},"token":"tok"}}]' + ) + assert WEBSOCKET.load_events == () + assert not WEBSOCKET.greets + # Socket.IO joins the namespace and hydrates once acknowledged, as 0.8.23 needs. + assert SOCKETIO.connect_frame(token="tok", pathname="/") == "40/_event," + assert SOCKETIO.load_events == (events.HYDRATE_EVENT, events.ON_LOAD_EVENT) + assert SOCKETIO.greets + + +def test_both_codecs_decode_to_the_plain_protocols_messages(): + assert WEBSOCKET.parse('["_ping"]') == SOCKETIO.parse("2") == [events.PING] + assert WEBSOCKET.parse('["_handshake",{"protocol":2}]') == [ + events.HANDSHAKE, + {"protocol": 2}, + ] + assert SOCKETIO.parse('40/_event,{"sid":"a"}') == [events.HANDSHAKE, '{"sid":"a"}'] + assert SOCKETIO.parse('0{"sid":"a"}') == [events.OPENED, '{"sid":"a"}'] + assert SOCKETIO.parse('44/_event,{"message":"no"}') == [ + events.REFUSED, + '{"message":"no"}', + ] + assert SOCKETIO.parse("41/_event,") == SOCKETIO.parse("1") == [events.DISCONNECTED] + for codec in (WEBSOCKET, SOCKETIO): + assert codec.parse(codec.emit("new_token", "t")) == ["new_token", "t"] + assert codec.parse(b'["event",{}]') is None + for frame in ("{}", "[1]", "[]", "nope"): + assert WEBSOCKET.parse(frame) is None + assert SOCKETIO.parse("6") is None + assert SOCKETIO.parse("42/_event,nope") is None + assert (WEBSOCKET.join, SOCKETIO.join) == ('["_connect",{}]', "40/_event,") + assert (WEBSOCKET.ping, SOCKETIO.ping) == ('["_ping"]', "2") + assert (WEBSOCKET.pong, SOCKETIO.pong) == ('["_pong"]', "3") + assert (WEBSOCKET.leave, SOCKETIO.leave) == (None, "41/_event,") + + +def test_only_a_refusal_or_a_disconnect_ends_a_session(): + refused = SOCKETIO.parse('44/_event,{"message":"no"}') + assert refused is not None + assert ( + events.end_reason(refused, "") == 'the server refused /_event: {"message":"no"}' + ) + assert events.end_reason([events.DISCONNECTED], "1") == ( + "the server disconnected the session ('1')" + ) + for message in ([events.PING], ["event", {}], ["new_token", "t"]): + assert events.end_reason(message, "") is None -def test_replies_of_both_versions_carry_the_same_delta(): - for reply in (HEAD_REPLY, OLD_REPLY): - name, update = json.loads(reply[len(EVENT_PREFIX) :]) - assert name == "event" - assert update["delta"] == {STATE: {SEQ_VAR: 7}} +@pytest.mark.parametrize( + ("reply", "codec"), + [(HEAD_REPLY, WEBSOCKET), (SOCKETIO_REPLY, SOCKETIO), (OLD_REPLY, SOCKETIO)], +) +def test_replies_of_every_version_carry_the_same_delta(reply: str, codec: Any): + message = codec.parse(reply) + assert message is not None + name, update = message + assert name == "event" + assert update["delta"] == {STATE: {SEQ_VAR: 7}} + + +@PROTOCOLS +def test_the_protocol_is_detected(protocol: WireProtocol): + # Each server refuses the websocket on the other protocol's path. + with serving(Scripted(protocol)) as url: + assert asyncio.run(events.detect_protocol(url)) == protocol + codec = asyncio.run(events.resolve_codec(Endpoint(url))) + assert codec is CODECS[protocol] + + +def test_a_named_protocol_is_not_detected(): + port = free_ports()[0] + endpoint = Endpoint(f"http://127.0.0.1:{port}", protocol="socketio") + assert asyncio.run(events.resolve_codec(endpoint)) is SOCKETIO + with pytest.raises(OSError): + asyncio.run(events.detect_protocol(endpoint.backend_url)) + + +def test_detection_gives_up_at_its_timeout(): + # The listener never answers the upgrade, as a hung backend does. + with socket.create_server(("127.0.0.1", 0)) as listener: + url = f"http://127.0.0.1:{listener.getsockname()[1]}" + started = time.monotonic() + with pytest.raises(TimeoutError): + asyncio.run(events.detect_protocol(url, 0.1)) + assert time.monotonic() - started < 5 def test_a_refused_namespace_fails_the_session(): @@ -216,7 +332,7 @@ async def on_connect(self, ws, sid): """Send a connect_error instead of the ack.""" await ws.send('44/_event,{"message":"Unable to connect"}') - with serving(Refusing()) as url: + with serving(Refusing("socketio")) as url: result = run_inline(plan(url, sessions=1)) assert result.sent == 0 (error,) = result.session_errors @@ -224,6 +340,22 @@ async def on_connect(self, ws, sid): assert "Unable to connect" in error +def test_a_session_closed_in_the_handshake_fails(): + class Closing(Scripted): + """Closes the connect with a policy violation, as reflex does without a token.""" + + async def on_connect(self, ws, sid): + """Close instead of the handshake.""" + await ws.close(1008) + + with serving(Closing()) as url: + result = run_inline(plan(url, sessions=1)) + assert result.sent == 0 + (error,) = result.session_errors + assert "could not connect: ConnectionClosedError" in error + assert "1008" in error + + def test_an_unexpected_frame_fails_the_session(): # reflex never sends binary attachments; anything the generator does not # speak ends the session with a reason instead of being ignored. @@ -269,8 +401,9 @@ def test_schedule_spreads_sessions_evenly(): assert merged == [int(k * second / 30) for k in range(30)] -def test_every_event_is_answered(): - server = Scripted() +@PROTOCOLS +def test_every_event_is_answered(protocol: WireProtocol): + server = Scripted(protocol) with serving(server) as url: result = run_inline(plan(url, sessions=3, rate=150.0)) # 150 ev/s over the 0.5 s window, as planned. @@ -428,16 +561,17 @@ def test_an_unordered_shape_takes_answers_in_any_order(): assert result.answered == result.sent -def test_a_new_token_is_adopted(): +@PROTOCOLS +def test_a_new_token_is_adopted(protocol: WireProtocol): class Renaming(Scripted): """Hands every connection a new token, as reflex does for a duplicate tab.""" async def on_connect(self, ws, sid): - """Send the new token before the namespace ack, like python-socketio.""" - await ws.send(events.emit_frame("new_token", f"fresh-{sid}")) + """Send the new token before the handshake, as reflex does.""" + await ws.send(self.codec.emit("new_token", f"fresh-{sid}")) await super().on_connect(ws, sid) - server = Renaming() + server = Renaming(protocol) with serving(server) as url: result = run_inline(plan(url, sessions=1, rate=20.0)) assert result.unanswered == 0 @@ -445,30 +579,32 @@ async def on_connect(self, ws, sid): assert all(token.startswith("fresh-") for token in server.tokens) -def test_a_session_that_never_hydrates_fails(monkeypatch: pytest.MonkeyPatch): - class NeverHydrated(Scripted): - """Ignores on_load_internal, so is_hydrated never becomes true.""" - - async def on_event(self, ws, event): - """Answer everything but on_load_internal.""" - if event["name"] != events.ON_LOAD_EVENT: - await super().on_event(ws, event) - +@PROTOCOLS +def test_a_session_that_never_hydrates_fails( + monkeypatch: pytest.MonkeyPatch, protocol: WireProtocol +): monkeypatch.setattr(events, "PRIME_TIMEOUT_S", 0.5) - with serving(NeverHydrated()) as url: + with serving(NeverHydrated(protocol)) as url: result = run_inline(plan(url, sessions=1)) assert result.sent == 0 assert len(result.session_errors) == 1 assert "not hydrated within 0.5 s" in result.session_errors[0] -def test_a_disconnect_mid_run_fails_the_session(): +@pytest.mark.parametrize( + ("protocol", "reason"), + [ + ("websocket", "the websocket closed: received 1008"), + ("socketio", "the server disconnected the session"), + ], +) +def test_a_disconnect_mid_run_fails_the_session(protocol: WireProtocol, reason: str): class Disconnecting(Scripted): - """Ends the namespace of the first connection instead of answering seq 5.""" + """Ends the session of the first connection instead of answering seq 5.""" def __init__(self) -> None: """Know no connection yet.""" - super().__init__() + super().__init__(protocol) self.first: ServerConnection | None = None async def seq_event(self, ws, event, seq): @@ -480,14 +616,14 @@ async def seq_event(self, ws, event, seq): if self.first is None: self.first = ws if ws is self.first and seq == 5: - await ws.send(events.DISCONNECT_FRAME) + await self.end(ws) return False return True with serving(Disconnecting()) as url: result = run_inline(plan(url, sessions=2, rate=40.0, warmup_s=0.0)) assert len(result.session_errors) == 1 - assert "disconnected" in result.session_errors[0] + assert reason in result.session_errors[0] # One session stopped after seq 5, which stays unanswered; the other ran on. assert result.sent == 15 assert result.unanswered == 1 @@ -567,7 +703,6 @@ def test_closed_loop_reports_service_time_only(): BOARD = "reflex___state____state.playground___state____board_state" -CLIENT_VAR = "last_client_rx_state_" LINK = LinkEvent(name=f"{BOARD}.join", payload={"token": "board-1"}, delta_key=BOARD) BOARD_SHAPE = EventShape( name=f"{BOARD}.set_seq_shared", @@ -581,9 +716,13 @@ def test_closed_loop_reports_service_time_only(): class Board(EchoServer): """Answers like the playground's shared board: a join links a connection to a token, and a seq event reaches every linked connection.""" - def __init__(self) -> None: - """Link nothing yet.""" - super().__init__(delta_key=BOARD, seq_var=SEQ_VAR) + def __init__(self, protocol: WireProtocol = "websocket") -> None: + """Link nothing yet. + + Args: + protocol: The protocol to speak. + """ + super().__init__(delta_key=BOARD, seq_var=SEQ_VAR, protocol=protocol) self.linked: dict[str, list[ServerConnection]] = {} self.names: dict[ServerConnection, list[str]] = {} self.payloads: list[dict[str, Any]] = [] @@ -598,7 +737,7 @@ async def on_event(self, ws: ServerConnection, event: dict[str, Any]) -> None: self.names.setdefault(ws, []).append(event["name"]) if event["name"] == LINK.name: self.linked.setdefault(event["payload"]["token"], []).append(ws) - await ws.send(events.emit_frame("event", {"delta": {BOARD: {SEQ_VAR: 0}}})) + await ws.send(self.codec.emit("event", {"delta": {BOARD: {SEQ_VAR: 0}}})) elif event["name"] == BOARD_SHAPE.name: self.payloads.append(event["payload"]) delta = { @@ -609,7 +748,7 @@ async def on_event(self, ws: ServerConnection, event: dict[str, Any]) -> None: } for index, linked in enumerate(self.linked_to(ws)): await self.deliver( - linked, index, events.emit_frame("event", {"delta": delta}) + linked, index, self.codec.emit("event", {"delta": delta}) ) else: await super().on_event(ws, event) @@ -643,16 +782,20 @@ def board_plan(url: str, **overrides: Any) -> LoadPlan: return plan(url, **{"shape": BOARD_SHAPE, "link": LINK, **overrides}) -def test_linked_sessions_join_after_hydrating_and_before_the_load(): - server = Board() +@PROTOCOLS +def test_linked_sessions_join_after_hydrating_and_before_the_load( + protocol: WireProtocol, +): + server = Board(protocol) with serving(server) as url: result = run_inline(board_plan(url, mode="closed", rate=None, sessions=3)) assert result.session_errors == [] assert result.answered > 0 assert len(server.linked["board-1"]) == 3 + primed = [*HYDRATION[protocol], LINK.name] for names in server.names.values(): - assert names[:3] == [events.HYDRATE_EVENT, events.ON_LOAD_EVENT, LINK.name] - assert set(names[3:]) == {BOARD_SHAPE.name} + assert names[: len(primed)] == primed + assert set(names[len(primed) :]) == {BOARD_SHAPE.name} def test_a_session_whose_join_is_not_acknowledged_fails( @@ -869,7 +1012,7 @@ def spy(method: str | None = None) -> Any: result = LoadRunner(load).run( on_window=lambda edge: windows.append((edge, time.perf_counter_ns())) ) - assert contexts == ["spawn"] + assert contexts == ["spawn", "spawn"] assert result.processes == 2 assert result.sessions == 4 # Without warmup, the server saw exactly the measured events of both processes. @@ -916,14 +1059,10 @@ def run() -> None: def test_a_generator_process_ends_when_its_parent_goes_away(): # However the harness exits, the closed pipe ends the generator's run. with serving(Scripted()) as url: - context = multiprocessing.get_context("spawn") - conn, child = context.Pipe() load = plan(url, sessions=1, rate=20.0, duration_s=60.0) - proc = context.Process( - target=events._worker, args=(load, range(1), child), daemon=True + proc, conn = events.spawn_process( + events._worker, load, range(1), name="reflex-bench-test" ) - proc.start() - child.close() assert conn.poll(30) assert conn.recv() == ("ready", []) conn.send(("go", time.perf_counter_ns())) @@ -933,27 +1072,46 @@ def test_a_generator_process_ends_when_its_parent_goes_away(): assert proc.exitcode is not None +class NeverHydrated(Scripted): + """Ignores the events that complete the hydration, so is_hydrated never becomes true.""" + + async def on_event(self, ws, event): + """Answer everything but those.""" + if event["name"] not in HYDRATING: + await super().on_event(ws, event) + + class Counting(Scripted): - """Counts the hydrations and the disconnects.""" + """Counts the hydrations and the sessions that ended; the protocol's probe is no session.""" - def __init__(self) -> None: - """Count nothing yet.""" - super().__init__() + def __init__(self, protocol: WireProtocol = "websocket") -> None: + """Count nothing yet. + + Args: + protocol: The protocol to speak. + """ + super().__init__(protocol) self.hydrated = 0 self.left = 0 + self._sessions: set[ServerConnection] = set() + + async def on_connect(self, ws, sid): + """Count a session.""" + self._sessions.add(ws) + await super().on_connect(ws, sid) async def on_event(self, ws, event): - """Count on_load_internal, then answer.""" - if event["name"] == events.ON_LOAD_EVENT: + """Count the event that completes a hydration, then answer.""" + if event["name"] in HYDRATING: self.hydrated += 1 await super().on_event(ws, event) async def handle(self, ws): - """Count a connection that ends.""" + """Count a session that ends.""" try: await super().handle(ws) finally: - self.left += 1 + self.left += ws in self._sessions def test_open_sessions_primes_and_closes_them(): @@ -978,14 +1136,6 @@ async def hold() -> list[str]: def test_open_sessions_reports_the_sessions_that_fail( monkeypatch: pytest.MonkeyPatch, ): - class NeverHydrated(Scripted): - """Ignores on_load_internal, so is_hydrated never becomes true.""" - - async def on_event(self, ws, event): - """Answer everything but on_load_internal.""" - if event["name"] != events.ON_LOAD_EVENT: - await super().on_event(ws, event) - monkeypatch.setattr(events, "PRIME_TIMEOUT_S", 0.5) with serving(NeverHydrated()) as url: loop = events.event_loop() @@ -1003,19 +1153,27 @@ async def hold() -> list[str]: class Holding(Counting): - """Pings often and records each connection's frames.""" + """Pings often and records each session's frames and close code.""" - def __init__(self) -> None: - """Record nothing yet.""" - super().__init__() + def __init__(self, protocol: WireProtocol = "websocket") -> None: + """Record nothing yet. + + Args: + protocol: The protocol to speak. + """ + super().__init__(protocol) self.ping_interval_s = 0.2 self.frames: list[list[str | bytes]] = [] + self.close_codes: list[int | None] = [] async def handle(self, ws): - """Serve a connection, recording its frames.""" + """Serve a connection, recording a session's frames and how it closed.""" frames: list[str | bytes] = [] - self.frames.append(frames) - await super().handle(cast(ServerConnection, _Recording(ws, frames))) + recording = cast(ServerConnection, _Recording(ws, frames)) + await super().handle(recording) + if recording in self._sessions: + self.frames.append(frames) + self.close_codes.append(ws.close_code) class _Recording: @@ -1071,8 +1229,10 @@ def wait_for(condition: Callable[[], bool], timeout: float = 10.0) -> bool: return True -def test_a_hold_keeps_hydrated_sessions_until_closed(): - server = Holding() +@PROTOCOLS +def test_a_hold_keeps_hydrated_sessions_until_closed(protocol: WireProtocol): + server = Holding(protocol) + codec = CODECS[protocol] with serving(server) as url: hold = events.hold_sessions(Endpoint(url), 3) try: @@ -1085,9 +1245,13 @@ def test_a_hold_keeps_hydrated_sessions_until_closed(): hold.kill() assert wait_for(lambda: server.left == 3) for frames in server.frames: - # The sessions answered the pings while held, then left the namespace. - assert events.PONG in frames - assert frames[-1] == events.DISCONNECT_FRAME + # The sessions answered the pings while held, then left. + assert codec.pong in frames + if codec.leave is not None: + assert frames[-1] == codec.leave + if codec.leave is None: + # The plain protocol leaves with a normal close. + assert server.close_codes == [1000] * 3 assert multiprocessing.active_children() == [] @@ -1104,16 +1268,14 @@ def test_a_hold_raises_when_sessions_cannot_connect(): def test_close_reports_sessions_lost_while_held(): class Dropping(Holding): - """Ends the namespace of the first hydrated connection half a second later.""" + """Closes the first hydrated connection half a second later.""" async def on_event(self, ws, event): - """Answer, then schedule the disconnect of the first connection.""" + """Answer, then schedule the end of the first connection.""" await super().on_event(ws, event) - if event["name"] == events.ON_LOAD_EVENT and self.hydrated == 1: + if event["name"] in HYDRATING and self.hydrated == 1: loop = asyncio.get_running_loop() - loop.call_later( - 0.5, lambda: loop.create_task(ws.send(events.DISCONNECT_FRAME)) - ) + loop.call_later(0.5, lambda: loop.create_task(self.end(ws))) with serving(Dropping()) as url: hold = events.hold_sessions(Endpoint(url), 2) @@ -1124,20 +1286,15 @@ async def on_event(self, ws, event): finally: hold.kill() assert len(errors) == 1 - assert "disconnected" in errors[0] + assert "the websocket closed: received 1008" in errors[0] assert multiprocessing.active_children() == [] def test_kill_ends_a_hold_that_is_still_priming(): - class NeverHydrated(Holding): - """Ignores on_load_internal, so no session finishes priming.""" + class NeverPrimed(NeverHydrated, Holding): + """A holding server where no session finishes priming.""" - async def on_event(self, ws, event): - """Answer everything but on_load_internal.""" - if event["name"] != events.ON_LOAD_EVENT: - await super().on_event(ws, event) - - with serving(NeverHydrated()) as url: + with serving(NeverPrimed()) as url: hold = events.hold_sessions(Endpoint(url), 2) killer = threading.Timer(1.0, hold.kill) killer.start() diff --git a/tests/units/reflex_bench/factories.py b/tests/units/reflex_bench/factories.py index 3a3e4c029cf..f92d97035c2 100644 --- a/tests/units/reflex_bench/factories.py +++ b/tests/units/reflex_bench/factories.py @@ -1,23 +1,74 @@ -"""Builders for reflex-bench result documents used across the tests.""" +"""Builders and helpers shared by the reflex-bench tests.""" from __future__ import annotations +import asyncio +import contextlib import logging import random import sys -from collections.abc import Mapping, Sequence +import threading +from collections.abc import Iterator, Mapping, Sequence from pathlib import Path from typing import Any +import pytest from reflex_bench.context import Context, Subject -from reflex_bench.drivers.events import LoadResult +from reflex_bench.drivers.echo_server import EchoServer +from reflex_bench.drivers.events import CODECS, LoadResult from reflex_bench.registry import Metric from reflex_bench.scheduler import Policy, finalize from reflex_bench.schema import SCHEMA_ID, BenchmarkDoc, MachineDoc, ResultDoc, RunKind +from reflex_bench.suites.events import PlaygroundStates + +import reflex as rx +from reflex.istate.shared import SharedStateBaseInternal VERSION = "sha256:" + "0" * 64 WALL = Metric(unit="s", direction="lower") BYTES = Metric(unit="B", direction="lower", assume="exact") +WEBSOCKET = CODECS["websocket"] +SOCKETIO = CODECS["socketio"] +PROTOCOLS = pytest.mark.parametrize("protocol", ["websocket", "socketio"]) +# The playground's states as the workspace's reflex names them: the shared +# state sits under reflex's internal parent of shared states. +STATES = PlaygroundStates( + bench=f"{rx.State.get_full_name()}.playground___state____bench_state", + board=f"{SharedStateBaseInternal.get_full_name()}.playground___state____board_state", +) + + +@contextlib.contextmanager +def serving(server: EchoServer) -> Iterator[str]: + """Run an echo server on a thread with its own event loop. + + Args: + server: The server. + + Yields: + Its base URL, like the backend URL of a started app. + """ + loop = asyncio.new_event_loop() + ready = threading.Event() + port: list[int] = [] + done: list[asyncio.Future[None]] = [] + + async def main() -> None: + done.append(loop.create_future()) + async with server.serve("127.0.0.1", 0) as bound: + port.append(bound) + ready.set() + await done[0] + + thread = threading.Thread(target=loop.run_until_complete, args=(main(),)) + thread.start() + assert ready.wait(10) + try: + yield f"http://127.0.0.1:{port[0]}" + finally: + loop.call_soon_threadsafe(done[0].set_result, None) + thread.join(10) + loop.close() def make_subject(reflex_version: str | None = "0.9.12") -> Subject: diff --git a/tests/units/reflex_bench/suites/test_events.py b/tests/units/reflex_bench/suites/test_events.py index 9d387487872..7f29967f162 100644 --- a/tests/units/reflex_bench/suites/test_events.py +++ b/tests/units/reflex_bench/suites/test_events.py @@ -21,17 +21,9 @@ from reflex_bench.scheduler import Planned, Policy, Scheduler, plan from reflex_bench.suites import events as suite -import reflex as rx -from reflex.istate.shared import SharedStateBaseInternal -from tests.units.reflex_bench.factories import make_load_result, make_subject +from tests.units.reflex_bench.factories import STATES, make_load_result, make_subject SHAPES = ("background", "complex", "cross", "simple") -# The playground's states as the workspace's reflex names them: the shared -# state sits under reflex's internal parent of shared states. -STATES = suite.PlaygroundStates( - bench=f"{rx.State.get_full_name()}.playground___state____bench_state", - board=f"{SharedStateBaseInternal.get_full_name()}.playground___state____board_state", -) CHEAP = [ "events.simple.capacity[manager=disk,sessions=10]", "events.simple.latency[manager=disk,sessions=10,rate=500]", diff --git a/tests/units/reflex_bench/suites/test_memory.py b/tests/units/reflex_bench/suites/test_memory.py index 28e701584b5..c31c4badde2 100644 --- a/tests/units/reflex_bench/suites/test_memory.py +++ b/tests/units/reflex_bench/suites/test_memory.py @@ -29,13 +29,13 @@ from reflex_bench.suites import memory from tests.units.reflex_bench.factories import ( + STATES, fake_proc, fake_rollup, make_context, make_load_result, make_subject, ) -from tests.units.reflex_bench.suites.test_events import STATES MIB = 1024**2 APP_PID = 1001 diff --git a/tests/units/reflex_bench/suites/test_wire.py b/tests/units/reflex_bench/suites/test_wire.py index 8aa867ac1fb..e333856b6f8 100644 --- a/tests/units/reflex_bench/suites/test_wire.py +++ b/tests/units/reflex_bench/suites/test_wire.py @@ -1,8 +1,9 @@ """Tests for reflex_bench.suites.wire. Nothing here starts reflex: the sessions run against scripted transports (or -the echo server on a loopback socket), and the benchmarks run through the -scheduler with the app and the websocket replaced by fakes. +the echo server on a loopback socket), on the plain protocol unless a test +names Socket.IO, and the benchmarks run through the scheduler with the app and +the websocket replaced by fakes. """ from __future__ import annotations @@ -13,7 +14,7 @@ import contextlib import dataclasses import json -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Iterator, Sequence from pathlib import Path from typing import Any @@ -22,33 +23,51 @@ from reflex_bench.drivers.app_process import CliResult from reflex_bench.drivers.echo_server import EchoServer from reflex_bench.drivers.events import ( - CONNECT_FRAME, - DISCONNECT_FRAME, + CODECS, + CONNECT, + DISCONNECTED, + HANDSHAKE, + HYDRATE_AND_LOAD_EVENT, HYDRATE_EVENT, HYDRATED_VAR, ON_LOAD_EVENT, - PING, PONG, ROOT_STATE, + SIO_CONNECT_FRAME, + SIO_DISCONNECT_FRAME, + SIO_PING, + SIO_PONG, + Codec, Endpoint, ProtocolError, - emit_frame, - event_frame, + WireProtocol, ) from reflex_bench.scheduler import Planned, Policy, Scheduler, plan from reflex_bench.suites import events as events_suite from reflex_bench.suites import wire from reflex_bench.suites.events import SEQ_VAR - -from tests.units.reflex_bench.drivers.test_events import serving -from tests.units.reflex_bench.factories import make_subject -from tests.units.reflex_bench.suites.test_events import STATES +from websockets.exceptions import ConnectionClosedError +from websockets.frames import Close + +from tests.units.reflex_bench.factories import ( + SOCKETIO, + STATES, + WEBSOCKET, + make_subject, + serving, +) BENCH_STATE = STATES.bench SHAPES = {name: shape_of(STATES) for name, shape_of in events_suite.SHAPES.items()} PLAYGROUND_STATE = "reflex___state____state.playground___state____playground_state" +HANDSHAKE_FRAME = ( + '["_handshake",{"ping_interval":25.0,"ping_timeout":120.0,"protocol":2,' + '"max_message_size":10000000}]' +) OPEN = '0{"sid":"abc","upgrades":[],"pingInterval":25000,"pingTimeout":20000}' ACK = '40/_event,{"sid":"abc"}' +# How reflex ends a plain session: it closes the websocket with a code. +CLOSED = ConnectionClosedError(Close(1008, ""), None) CHANGES = ( "set_scalar", "append_item", @@ -62,27 +81,34 @@ def size(*frames: str) -> int: return sum(len(frame.encode()) for frame in frames) -def delta_frame(delta: dict[str, Any], **more: Any) -> str: - return emit_frame("event", {"delta": delta, **more}) +def delta_frame(delta: dict[str, Any], codec: Codec = WEBSOCKET, **more: Any) -> str: + return codec.emit("event", {"delta": delta, **more}) + + +# A frame of the scripted server, or the error its websocket raises instead. +Reply = str | bytes | Exception class FakeTransport: """A scripted server: every sent frame queues the replies a script gives.""" def __init__( - self, script: Callable[[str], list[str]], greeting: list[str] | None = None + self, + script: Callable[[str], Sequence[Reply]], + codec: Codec = WEBSOCKET, + greeting: Sequence[Reply] = (), ) -> None: """Queue the server's greeting. Args: script: Maps a sent frame to the frames it triggers. - greeting: The frames waiting before anything is sent; the engine.io - open by default. + codec: The protocol; one that greets opens with an engine.io open. + greeting: More frames waiting before anything is sent. """ self.script = script self.sent: list[str] = [] - self.inbox: collections.deque[str | bytes] = collections.deque( - greeting or [OPEN] + self.inbox: collections.deque[Reply] = collections.deque( + [OPEN, *greeting] if codec.greets else greeting ) async def send(self, message: str) -> None: @@ -106,58 +132,74 @@ async def recv(self) -> str | bytes: if not self.inbox: msg = "the session waits for a frame the script never sends" raise AssertionError(msg) - return self.inbox.popleft() + reply = self.inbox.popleft() + if isinstance(reply, Exception): + raise reply + return reply def playground_script( - *, new_token: str | None = None, replies: dict[str, list[str]] | None = None -) -> Callable[[str], list[str]]: + *, + new_token: str | None = None, + replies: dict[str, list[Reply]] | None = None, + codec: Codec = WEBSOCKET, +) -> Callable[[str], list[Reply]]: """Script the playground's answers: a hydration in two deltas, then echoes. Args: - new_token: A token handed out before the namespace ack. + new_token: A token handed out before the handshake. replies: The frames answering an event, by handler name; other events get one echo of their sequence number, and ``on_load_internal`` the ``is_hydrated`` delta, with the router's page away from ``/``. + codec: The protocol, which decodes the session's frames too. Returns: The script. """ - def script(message: str) -> list[str]: - if message == CONNECT_FRAME: - handed = [emit_frame("new_token", new_token)] if new_token else [] - return [*handed, ACK] - if message in {PONG, DISCONNECT_FRAME}: - return [] - event = json.loads(message.partition(",")[2])[1] + def answer(event: dict[str, Any]) -> list[Reply]: name = event["name"] if replies and name in replies: return replies[name] + snapshot = delta_frame( + { + ROOT_STATE: {HYDRATED_VAR: False, "router_rx_state_": {"a": 1}}, + BENCH_STATE: {SEQ_VAR: 0, "part_a_rx_state_": 0}, + PLAYGROUND_STATE: {"count_rx_state_": 0, "items_rx_state_": ["alpha"]}, + }, + codec, + ) if name == HYDRATE_EVENT: - return [ - delta_frame({ - ROOT_STATE: {HYDRATED_VAR: False, "router_rx_state_": {"a": 1}}, - BENCH_STATE: {SEQ_VAR: 0, "part_a_rx_state_": 0}, - PLAYGROUND_STATE: { - "count_rx_state_": 0, - "items_rx_state_": ["alpha"], - }, - }) - ] + return [snapshot] + if name == HYDRATE_AND_LOAD_EVENT: + return [snapshot, navigation_reply(event["router_data"], codec)] if name == ON_LOAD_EVENT: - return [navigation_reply(event["router_data"])] + return [navigation_reply(event["router_data"], codec)] state, _, _handler = name.rpartition(".") - return [delta_frame({state: {SEQ_VAR: event["payload"]["seq"]}})] + return [delta_frame({state: {SEQ_VAR: event["payload"]["seq"]}}, codec)] + + def script(message: str) -> list[Reply]: + args = codec.parse(message) + if args is None or args[0] in {PONG, DISCONNECTED}: + # Socket.IO's pong decodes to nothing. + return [] + handed = [codec.emit("new_token", new_token)] if new_token else [] + if args[0] == CONNECT: + return [*handed, HANDSHAKE_FRAME, *answer(args[1]["event"])] + if args[0] == HANDSHAKE: + # Socket.IO's namespace join. + return [*handed, ACK] + return answer(args[1]) return script -def navigation_reply(router_data: dict[str, Any]) -> str: +def navigation_reply(router_data: dict[str, Any], codec: Codec = WEBSOCKET) -> str: """Build the delta that completes an ``on_load_internal``. Args: router_data: The event's router data. + codec: The protocol. Returns: ``is_hydrated`` set; away from ``/`` the root state's router too. @@ -167,7 +209,20 @@ def navigation_reply(router_data: dict[str, Any]) -> str: root["router_rx_state_"] = { "page": {"path": router_data["pathname"], "params": router_data["query"]} } - return delta_frame({ROOT_STATE: root}) + return delta_frame({ROOT_STATE: root}, codec) + + +def answers(script: Callable[[str], Sequence[Reply]], message: str) -> list[str]: + """Take the text frames a script answers a frame with. + + Args: + script: The script. + message: The frame sent. + + Returns: + Its text frames. + """ + return [frame for frame in script(message) if isinstance(frame, str)] def run(work: Any) -> Any: @@ -184,16 +239,21 @@ def test_hydration_counts_every_frame_until_hydrated(): ws = FakeTransport(playground_script(new_token="tok-2")) session = wire.WireSession(ws, "tok-1") hydration = run(session.hydrate()) - hydrate = event_frame(HYDRATE_EVENT, {}, token="tok-2", pathname="/") - on_load = event_frame(ON_LOAD_EVENT, {}, token="tok-2", pathname="/") - # The new token is adopted before the hydration events go out. - assert ws.sent == [CONNECT_FRAME, hydrate, on_load] - assert hydration.sent_bytes == size(CONNECT_FRAME, hydrate, on_load) - assert hydration.sent_frames == 3 - first, second = playground_script()(hydrate)[0], playground_script()(on_load)[0] - token_frame = emit_frame("new_token", "tok-2") - assert hydration.received_bytes == size(OPEN, token_frame, ACK, first, second) - assert hydration.received_frames == 5 + # The connect frame carries the hydration as its boot event. + connect = WEBSOCKET.connect_frame(token="tok-1", pathname="/") + assert ws.sent == [connect] + assert '"name":"reflex___state____state.hydrate_and_load"' in connect + assert hydration.sent_bytes == size(connect) + assert hydration.sent_frames == 1 + token_frame, handshake, first, second = answers( + playground_script(new_token="tok-2"), connect + ) + assert (token_frame, handshake) == ( + WEBSOCKET.emit("new_token", "tok-2"), + HANDSHAKE_FRAME, + ) + assert hydration.received_bytes == size(token_frame, handshake, first, second) + assert hydration.received_frames == 4 assert hydration.largest_frame_bytes == size(first) assert hydration.delta_bytes == { ROOT_STATE: size(json.dumps({HYDRATED_VAR: False, "router_rx_state_": {"a": 1}}, separators=(",", ":"))) @@ -201,44 +261,87 @@ def test_hydration_counts_every_frame_until_hydrated(): BENCH_STATE: size('{"last_seq_rx_state_":0,"part_a_rx_state_":0}'), PLAYGROUND_STATE: size('{"count_rx_state_":0,"items_rx_state_":["alpha"]}'), } # fmt: skip + # Later events carry the new token. assert session.token == "tok-2" -def test_a_ping_during_hydration_is_answered_and_not_counted(): - ws = FakeTransport(playground_script(), greeting=[OPEN, PING]) - hydration = run(wire.WireSession(ws, "tok").hydrate()) - assert ws.sent[:2] == [CONNECT_FRAME, PONG] - # Keepalives are timing, not payload: neither the ping nor the pong counts. +def test_a_socketio_hydration_joins_then_sends_the_hydration_events(): + script = playground_script(new_token="tok-2", codec=SOCKETIO) + ws = FakeTransport(script, SOCKETIO) + session = wire.WireSession(ws, "tok-1", codec=SOCKETIO) + hydration = run(session.hydrate()) + hydrate = SOCKETIO.event_frame(HYDRATE_EVENT, {}, token="tok-2", pathname="/") + on_load = SOCKETIO.event_frame(ON_LOAD_EVENT, {}, token="tok-2", pathname="/") + # The new token is adopted before the hydration events go out. + assert ws.sent == [SIO_CONNECT_FRAME, hydrate, on_load] + assert hydration.sent_bytes == size(SIO_CONNECT_FRAME, hydrate, on_load) assert hydration.sent_frames == 3 - assert hydration.received_frames == 4 - assert hydration.sent_bytes == size(CONNECT_FRAME, *ws.sent[2:]) + (first,), (second,) = answers(script, hydrate), answers(script, on_load) + token_frame = SOCKETIO.emit("new_token", "tok-2") + assert hydration.received_bytes == size(OPEN, token_frame, ACK, first, second) + assert hydration.received_frames == 5 + assert hydration.largest_frame_bytes == size(first) + assert set(hydration.delta_bytes) == {ROOT_STATE, BENCH_STATE, PLAYGROUND_STATE} + assert session.token == "tok-2" @pytest.mark.parametrize( - ("frames", "match"), + ("protocol", "ping", "pong", "sent_frames", "received_frames"), [ - (['44/_event,{"message":"nope"}'], "refused"), - ([DISCONNECT_FRAME], "disconnected"), - (["1"], "disconnected"), - (["6"], "unexpected frame"), - ([emit_frame("reload", "/")], "reload"), + ("websocket", '["_ping"]', '["_pong"]', 1, 3), + ("socketio", SIO_PING, SIO_PONG, 3, 4), + ], +) +def test_a_ping_during_hydration_is_answered_and_not_counted( + protocol: WireProtocol, ping: str, pong: str, sent_frames: int, received_frames: int +): + codec = CODECS[protocol] + ws = FakeTransport(playground_script(codec=codec), codec, greeting=[ping]) + hydration = run(wire.WireSession(ws, "tok", codec=codec).hydrate()) + assert ws.sent[1] == pong + # Keepalives are timing, not payload: neither the ping nor the pong counts. + assert hydration.sent_frames == sent_frames + assert hydration.received_frames == received_frames + assert hydration.sent_bytes == size(ws.sent[0], *ws.sent[2:]) + + +@pytest.mark.parametrize( + ("protocol", "frames", "match"), + [ + ("websocket", [CLOSED], "received 1008"), + ("websocket", ['{"not":"a message"}'], "unexpected frame"), + ("websocket", [b"\x00\x01"], "unexpected frame"), + ("websocket", [WEBSOCKET.emit("reload", "/")], "reload"), + ("socketio", ['44/_event,{"message":"nope"}'], "refused"), + ("socketio", [SIO_DISCONNECT_FRAME], "disconnected"), + ("socketio", ["1"], "disconnected"), + ("socketio", ["6"], "unexpected frame"), + ("socketio", [SOCKETIO.emit("reload", "/")], "reload"), ], ) def test_a_frame_that_ends_the_session_fails_the_hydration( - frames: list[str], match: str + protocol: WireProtocol, frames: list[Reply], match: str ): - ws = FakeTransport(lambda message: frames if message == CONNECT_FRAME else []) - with pytest.raises(ProtocolError, match=match): - run(wire.WireSession(ws, "tok").hydrate()) + codec = CODECS[protocol] + + def script(message: str) -> list[Reply]: + args = codec.parse(message) + return frames if args is not None and args[0] in {CONNECT, HANDSHAKE} else [] + + ws = FakeTransport(script, codec) + with pytest.raises((ProtocolError, ConnectionClosedError), match=match): + run(wire.WireSession(ws, "tok", codec=codec).hydrate()) @pytest.mark.parametrize( ("first", "match"), [(ACK, "did not open"), ("3", "unexpected frame")] ) -def test_a_server_that_does_not_open_fails_the_hydration(first: str, match: str): +def test_a_socketio_server_that_does_not_open_fails_the_hydration( + first: str, match: str +): ws = FakeTransport(lambda message: [], greeting=[first]) with pytest.raises(ProtocolError, match=match): - run(wire.WireSession(ws, "tok").hydrate()) + run(wire.WireSession(ws, "tok", codec=SOCKETIO).hydrate()) async def hydrated_exchange( @@ -252,7 +355,9 @@ async def hydrated_exchange( def test_an_exchange_counts_the_request_and_its_reply(): ws = FakeTransport(playground_script()) exchange = run(hydrated_exchange(ws, "simple", seq=7)) - request = event_frame(SHAPES["simple"].name, {"seq": 7}, token="tok", pathname="/") + request = WEBSOCKET.event_frame( + SHAPES["simple"].name, {"seq": 7}, token="tok", pathname="/" + ) reply = delta_frame({BENCH_STATE: {SEQ_VAR: 7}}) assert ws.sent[-1] == request assert exchange.request_bytes == size(request) @@ -306,7 +411,7 @@ def test_a_navigation_sends_on_load_from_the_new_route_until_hydrated_again(): session = wire.WireSession(ws, "tok") run(session.hydrate()) exchange = run(session.navigate("/item/42", {"item_id": "42"})) - request = event_frame( + request = WEBSOCKET.event_frame( ON_LOAD_EVENT, {}, token="tok", pathname="/item/42", query={"item_id": "42"} ) reply = navigation_reply({"pathname": "/item/42", "query": {"item_id": "42"}}) @@ -326,7 +431,7 @@ def test_a_navigation_sends_on_load_from_the_new_route_until_hydrated_again(): ) # Later events come from the new route. run(session.exchange(SHAPES["simple"], 2)) - assert ws.sent[-1] == event_frame( + assert ws.sent[-1] == WEBSOCKET.event_frame( SHAPES["simple"].name, {"seq": 2}, token="tok", @@ -354,16 +459,42 @@ def test_a_navigation_with_on_load_events_is_summed_until_hydrated(): assert not ws.inbox -def test_a_disconnect_while_waiting_for_the_reply_fails(): +@pytest.mark.parametrize( + ("protocol", "end", "error", "match"), + [ + ("websocket", CLOSED, ConnectionClosedError, "received 1008"), + ("socketio", SIO_DISCONNECT_FRAME, ProtocolError, "disconnected"), + ], +) +def test_a_disconnect_while_waiting_for_the_reply_fails( + protocol: WireProtocol, end: Reply, error: type[Exception], match: str +): + codec = CODECS[protocol] simple = SHAPES["simple"].name - ws = FakeTransport(playground_script(replies={simple: [DISCONNECT_FRAME]})) - with pytest.raises(ProtocolError, match="disconnected"): - run(hydrated_exchange(ws, "simple")) + ws = FakeTransport(playground_script(replies={simple: [end]}, codec=codec), codec) + async def exchange() -> wire.Exchange: + session = wire.WireSession(ws, "tok", codec=codec) + await session.hydrate() + return await session.exchange(SHAPES["simple"], 1) -def test_measure_connects_hydrates_and_leaves_the_namespace( - monkeypatch: pytest.MonkeyPatch, -): + with pytest.raises(error, match=match): + run(exchange()) + + +@contextlib.contextmanager +def connecting( + monkeypatch: pytest.MonkeyPatch, codec: Codec +) -> Iterator[tuple[list[FakeTransport], list[str], list[dict[str, str]]]]: + """Replace the websocket connect with scripted transports of a protocol. + + Args: + monkeypatch: Patches the connect. + codec: The protocol. + + Yields: + The transports, the URLs and the headers of every connect. + """ transports: list[FakeTransport] = [] urls: list[str] = [] headers: list[dict[str, str]] = [] @@ -372,30 +503,66 @@ def test_measure_connects_hydrates_and_leaves_the_namespace( async def connect(url: str, **kwargs: Any) -> AsyncIterator[FakeTransport]: urls.append(url) headers.append(kwargs["additional_headers"]) - transports.append(ws := FakeTransport(playground_script())) + transports.append(ws := FakeTransport(playground_script(codec=codec), codec)) yield ws monkeypatch.setattr(wire, "connect", connect) - hydration = wire.measure( - Endpoint("http://localhost:8000", "/page"), wire.WireSession.hydrate - ) + yield transports, urls, headers + + +@pytest.mark.parametrize( + ("protocol", "url", "received_frames"), + [ + ("websocket", "ws://localhost:8000/_event?token=", 3), + ("socketio", "ws://localhost:8000/_event/?EIO=4&transport=websocket&token=", 4), + ], +) +def test_measure_connects_hydrates_and_leaves( + monkeypatch: pytest.MonkeyPatch, + protocol: WireProtocol, + url: str, + received_frames: int, +): + codec = CODECS[protocol] + with connecting(monkeypatch, codec) as (transports, urls, headers): + hydration = wire.measure( + Endpoint("http://localhost:8000", "/page", protocol=protocol), + wire.WireSession.hydrate, + ) (ws,) = transports - assert urls[0].startswith( - "ws://localhost:8000/_event/?EIO=4&transport=websocket&token=" - ) + assert urls[0].startswith(url) # A browser sends the page's origin; the router's page host comes from it. assert headers == [{"Origin": "http://localhost:8000"}] - assert hydration.received_frames == 4 - assert ws.sent[-1] == DISCONNECT_FRAME - assert '"pathname":"/page"' in ws.sent[1] + assert hydration.received_frames == received_frames + assert '"pathname":"/page"' in ws.sent[hydration.sent_frames - 1] + # Only the leave follows the hydration; the plain protocol leaves by closing. + assert ws.sent[hydration.sent_frames :] == ( + [] if codec.leave is None else [codec.leave] + ) -def test_measure_against_the_echo_server_counts_real_frames(): - with serving(EchoServer(delta_key=BENCH_STATE, seq_var=SEQ_VAR)) as url: +@pytest.mark.parametrize( + ("protocol", "hydration_frames", "hydration_deltas"), + [ + ("websocket", 2, ['{"is_hydrated_rx_state_":true}']), + ( + "socketio", + 4, + ['{"is_hydrated_rx_state_":false}', '{"is_hydrated_rx_state_":true}'], + ), + ], +) +def test_measure_against_the_echo_server_detects_its_protocol_and_counts_real_frames( + protocol: WireProtocol, hydration_frames: int, hydration_deltas: list[str] +): + codec = CODECS[protocol] + echo = EchoServer(delta_key=BENCH_STATE, seq_var=SEQ_VAR, protocol=protocol) + with serving(echo) as url: async def work( session: wire.WireSession, ) -> tuple[wire.Hydration, wire.Exchange, wire.Exchange]: + assert session.codec is codec return ( await session.hydrate(), await session.exchange(SHAPES["simple"], 3), @@ -403,24 +570,20 @@ async def work( ) hydration, exchange, navigation = wire.measure(Endpoint(url), work) - reply = emit_frame("event", {"delta": {BENCH_STATE: {SEQ_VAR: 3}}, "events": []}) + reply = codec.emit("event", {"delta": {BENCH_STATE: {SEQ_VAR: 3}}, "events": []}) assert exchange.response_bytes == size(reply) assert exchange.reply == reply - hydrated = emit_frame( + hydrated = codec.emit( "event", {"delta": {ROOT_STATE: {HYDRATED_VAR: True}}, "events": []} ) assert navigation.reply == hydrated assert navigation.response_frames == 1 # The token is a uuid4, so the request has a fixed length. assert navigation.request_bytes == size( - event_frame(ON_LOAD_EVENT, {}, token="x" * 36, pathname="/counter") + codec.event_frame(ON_LOAD_EVENT, {}, token="x" * 36, pathname="/counter") ) - assert hydration.received_frames == 4 - assert hydration.delta_bytes == { - ROOT_STATE: size( - '{"is_hydrated_rx_state_":false}', '{"is_hydrated_rx_state_":true}' - ) - } + assert hydration.received_frames == hydration_frames + assert hydration.delta_bytes == {ROOT_STATE: size(*hydration_deltas)} def test_measure_times_out(monkeypatch: pytest.MonkeyPatch): @@ -434,7 +597,7 @@ async def hang(session: wire.WireSession) -> None: monkeypatch.setattr(wire, "connect", connect) monkeypatch.setattr(wire, "WIRE_TIMEOUT_S", 0.05) with pytest.raises(TimeoutError): - wire.measure(Endpoint("http://localhost:1"), hang) + wire.measure(Endpoint("http://localhost:1", protocol="websocket"), hang) ALLOWED_RX = { @@ -583,7 +746,7 @@ class Fakes: log: list[str] = dataclasses.field(default_factory=list) compiled: list[Path] = dataclasses.field(default_factory=list) envs: list[dict[str, str]] = dataclasses.field(default_factory=list) - replies: dict[str, list[str]] = dataclasses.field(default_factory=dict) + replies: dict[str, list[Reply]] = dataclasses.field(default_factory=dict) @pytest.fixture @@ -628,6 +791,9 @@ async def connect(url: str, **kwargs: Any) -> AsyncIterator[FakeTransport]: monkeypatch.setattr(wire, "playground_states", lambda ctx: STATES) monkeypatch.setattr(wire, "AppProcess", FakeApp) monkeypatch.setattr(wire, "connect", connect) + monkeypatch.setattr( + wire, "resolve_codec", lambda endpoint: asyncio.sleep(0, WEBSOCKET) + ) return state @@ -651,21 +817,19 @@ def test_hydrate_starts_a_backend_per_sample_and_stops_it(tmp_path: Path, fakes: "connect", "app stop", ] - script = playground_script() # The token is a uuid4, so the sent frames only have a fixed length. - hydrate = event_frame(HYDRATE_EVENT, {}, token="x" * 36) - on_load = event_frame(ON_LOAD_EVENT, {}, token="x" * 36) - first, second = script(hydrate)[0], script(on_load)[0] + connect = WEBSOCKET.connect_frame(token="x" * 36, pathname="/") + handshake, first, second = answers(playground_script(), connect) samples = { name: metric["samples"]["A"] for name, metric in entry["metrics"].items() } - assert samples["hydrate_received_bytes"] == [size(OPEN, ACK, first, second)] - assert samples["hydrate_frames"] == [4] - assert samples["hydrate_sent_bytes"] == [size(CONNECT_FRAME, hydrate, on_load)] + assert samples["hydrate_received_bytes"] == [size(handshake, first, second)] + assert samples["hydrate_frames"] == [3] + assert samples["hydrate_sent_bytes"] == [size(connect)] (extra,) = entry["sample_extra"] assert extra["largest_frame_bytes"] == size(first) assert set(extra["delta_bytes"]) == {ROOT_STATE, BENCH_STATE, PLAYGROUND_STATE} - assert extra["sent_frames"] == 3 + assert extra["sent_frames"] == 1 (env,) = fakes.envs assert env["REFLEX_STATE_MANAGER_MODE"] == "memory" assert env["GRANIAN_WORKERS"] == "1" @@ -688,11 +852,13 @@ def test_event_measures_one_exchange_per_shape(tmp_path: Path, fakes: Fakes): } assert samples["response_bytes"] == [size(empty, echo)] assert samples["response_frames"] == [2] - request = event_frame(background, {"seq": 1}, token="x" * 36, pathname="/") + request = WEBSOCKET.event_frame( + background, {"seq": 1}, token="x" * 36, pathname="/" + ) assert samples["request_bytes"] == [size(request)] (extra,) = entry["sample_extra"] assert extra["reply"] == echo - assert extra["hydration"]["received_frames"] == 4 + assert extra["hydration"]["received_frames"] == 3 assert fakes.log == [ "copy playground", "run_cli compile", @@ -712,7 +878,7 @@ def test_navigate_measures_the_route_change_after_hydration( samples = { name: metric["samples"]["A"] for name, metric in entry["metrics"].items() } - request = event_frame( + request = WEBSOCKET.event_frame( ON_LOAD_EVENT, {}, token="x" * 36, @@ -726,7 +892,7 @@ def test_navigate_measures_the_route_change_after_hydration( (extra,) = entry["sample_extra"] assert extra["reply"] == reply # Hydrated on "/" first: its on_load reply carries no router. - assert extra["hydration"]["received_frames"] == 4 + assert extra["hydration"]["received_frames"] == 3 assert extra["hydration"]["delta_bytes"][ROOT_STATE] == size( '{"is_hydrated_rx_state_":false,"router_rx_state_":{"a":1}}', '{"is_hydrated_rx_state_":true}', @@ -762,7 +928,9 @@ def test_delta_generates_its_own_app_and_keys_the_series_on_its_hash( assert entry["metrics"]["response_bytes"]["samples"]["A"] == [size(reply)] (extra,) = entry["sample_extra"] assert extra["request_bytes"] == size( - event_frame(f"{wire.WIRE_STATE}.append_item", {"seq": 1}, token="x" * 36) + WEBSOCKET.event_frame( + f"{wire.WIRE_STATE}.append_item", {"seq": 1}, token="x" * 36 + ) ) assert extra["response_frames"] == 1 assert extra["delta_bytes"] == { @@ -773,8 +941,8 @@ def test_delta_generates_its_own_app_and_keys_the_series_on_its_hash( def test_a_failed_sample_still_stops_the_backend(tmp_path: Path, fakes: Fakes): - fakes.replies[SHAPES["simple"].name] = [DISCONNECT_FRAME] + fakes.replies[SHAPES["simple"].name] = [CLOSED] entry = run_bench(tmp_path, "wire.event", shape="simple") assert entry["status"] == "failed" - assert "disconnected" in entry["error"] + assert "received 1008" in entry["error"] assert fakes.log[-1] == "app stop" diff --git a/tests/units/reflex_bench/suites/test_wire_app.py b/tests/units/reflex_bench/suites/test_wire_app.py index f50d26d4c52..1a253476efd 100644 --- a/tests/units/reflex_bench/suites/test_wire_app.py +++ b/tests/units/reflex_bench/suites/test_wire_app.py @@ -65,10 +65,9 @@ def test_navigate_for_real(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): pathname = wire.ROUTES[route][0] assert f'"path":"{pathname}"' in extra["reply"] assert set(extra["delta_bytes"]) == {ROOT_STATE} - # The hydration on "/" came first: the open, the namespace ack, the - # hydrate delta, on_load_internal's update of the always dirty - # rx.dynamic var of /widgets, and the delta setting is_hydrated. - assert extra["hydration"]["received_frames"] == 5 + # The hydration on "/" came first: the handshake, the snapshot of the + # connect's hydrate_and_load, and the delta setting is_hydrated. + assert extra["hydration"]["received_frames"] == 3 # The dynamic route argument is a var of the root state. assert '"item_id_rx_state_":"42"' in entries["item"]["sample_extra"][0]["reply"] diff --git a/tests/units/test_app.py b/tests/units/test_app.py index ae575628094..d11091614c2 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -16,17 +16,15 @@ import threading import unittest.mock import uuid -from collections.abc import AsyncGenerator, Generator +from collections.abc import Generator from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext as does_not_raise from importlib.util import find_spec from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, Mock import pytest -import pytest_asyncio -import reflex_base from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor @@ -64,11 +62,7 @@ from reflex.app import ( App, ComponentCallable, - EventNamespace, _ContextMiddleware, - _decode_asgi_headers, - _sio_dumps, - _sio_loads, default_overlay_component, ) from reflex.compiler.compiler import ( @@ -79,7 +73,7 @@ ) from reflex.compiler.plugins import default_page_plugins from reflex.environment import environment -from reflex.istate.data import HeaderData, RouterData, URLData +from reflex.istate.data import RouterData, URLData from reflex.istate.manager.disk import StateManagerDisk from reflex.istate.manager.memory import StateManagerMemory from reflex.istate.manager.redis import StateManagerRedis @@ -97,7 +91,7 @@ from reflex.utils import exec as exec_utils from reflex.utils.token_manager import RedisTokenManager, SocketRecord -from .conftest import active_tracer, chdir, metric_points +from .conftest import chdir from .states import GenState from .states.upload import ( ChildFileUploadState, @@ -4453,98 +4447,6 @@ def test_call_marks_later_dev_backend_worker_as_hot_reload( assert compile_mock.call_args.kwargs["trigger"] == "hot_reload" -def test_decode_asgi_headers(): - """_decode_asgi_headers decodes raw ASGI header pairs into a str dict.""" - assert _decode_asgi_headers([]) == {} - assert _decode_asgi_headers([ - (b"host", b"example.com"), - (b"x-forwarded-for", b"10.0.0.1, 10.0.0.2"), - ]) == { - "host": "example.com", - "x-forwarded-for": "10.0.0.1, 10.0.0.2", - } - # Repeated cookie fields are joined with "; " (RFC 9113 section 8.2.3). - assert _decode_asgi_headers([ - (b"cookie", b"a=1"), - (b"cookie", b"b=2"), - ]) == {"cookie": "a=1; b=2"} - # Any other repeated field is joined in order with ", " (RFC 9110 section 5.3). - assert _decode_asgi_headers([ - (b"x-forwarded-for", b"10.0.0.1"), - (b"accept", b"text/html"), - (b"x-forwarded-for", b"10.0.0.2, 10.0.0.3"), - ]) == { - "x-forwarded-for": "10.0.0.1, 10.0.0.2, 10.0.0.3", - "accept": "text/html", - } - # Names and values are decoded as UTF-8. - assert _decode_asgi_headers([(b"x-name", "café".encode())]) == {"x-name": "café"} - # Invalid UTF-8 values fall back to Latin-1 instead of aborting the connect. - assert _decode_asgi_headers([(b"x-name", b"\xe9")]) == {"x-name": "\u00e9"} - - -def test_static_router_data_keeps_every_repeated_header_field( - monkeypatch: pytest.MonkeyPatch, -): - """Repeated cookie and x-forwarded-for fields all reach the router data.""" - monkeypatch.setattr("reflex.app.get_config", lambda: Mock(trusted_proxy_hops=2)) - router_data = EventNamespace("/event", Mock())._build_static_router_data( - "sid1", - { - "asgi.scope": { - "headers": [ - (b"cookie", b"sid=1"), - (b"x-forwarded-for", b"203.0.113.7"), - (b"cookie", b"other=2"), - (b"x-forwarded-for", b"10.0.0.2"), - ], - "client": ("127.0.0.1", 1234), - } - }, - ) - - headers = router_data[constants.RouteVar.HEADERS] - assert headers["cookie"] == "sid=1; other=2" - assert headers["x-forwarded-for"] == "203.0.113.7, 10.0.0.2" - # The original client hop comes first, not the proxy that added the second field. - assert router_data[constants.RouteVar.CLIENT_IP] == "203.0.113.7" - assert HeaderData.from_router_data(router_data).cookie == "sid=1; other=2" - - -@pytest.mark.parametrize( - ("trusted_proxy_hops", "forwarded_for", "peer_ip", "expected_ip"), - [ - (0, "6.6.6.6, 203.0.113.7", "10.0.0.2", "10.0.0.2"), - (1, "203.0.113.7", "10.0.0.2", "203.0.113.7"), - (2, "6.6.6.6, 203.0.113.7, 10.0.0.2", "127.0.0.1", "203.0.113.7"), - (2, "203.0.113.7", "10.0.0.2", "10.0.0.2"), - ], -) -def test_static_router_data_resolves_client_ip_from_trusted_proxy_hops( - monkeypatch: pytest.MonkeyPatch, - trusted_proxy_hops: int, - forwarded_for: str, - peer_ip: str, - expected_ip: str, -): - """Only configured rightmost forwarded hops determine the client IP.""" - monkeypatch.setattr( - "reflex.app.get_config", - lambda: Mock(trusted_proxy_hops=trusted_proxy_hops), - ) - router_data = EventNamespace("/event", Mock())._build_static_router_data( - "sid1", - { - "asgi.scope": { - "headers": [(b"x-forwarded-for", forwarded_for.encode())], - "client": (peer_ip, 1234), - } - }, - ) - - assert router_data[constants.RouteVar.CLIENT_IP] == expected_ip - - def test_call_ignores_stale_marker_without_dev_backend_reload( compilable_app: tuple[App, Path], mocker: MockerFixture, @@ -4614,531 +4516,6 @@ def register_route(self, *, add_page, **_): assert app._plugin_routes_registered -@pytest.fixture -def event_namespace() -> EventNamespace: - """An EventNamespace with a mock app and one linked client session. - - Returns: - The event namespace. - """ - namespace = EventNamespace(namespace="/_event", app=Mock()) - namespace.sid_to_token["known_sid"] = "some_token" - return namespace - - -@pytest.fixture -def frontend_errors(event_namespace: EventNamespace) -> list[str]: - """Capture exceptions routed to the app's frontend exception handler. - - Args: - event_namespace: The event namespace. - - Returns: - The captured exception messages. - """ - errors: list[str] = [] - event_namespace.app.frontend_exception_handler = lambda exc: errors.append(str(exc)) - return errors - - -@pytest.fixture -def client_error_console() -> Generator[dict[str, list[str]], None, None]: - """Capture messages logged through the app's logger, keyed by level. - - Yields: - Captured messages keyed by log level. - """ - captured: dict[str, list[str]] = {"error": [], "warn": [], "debug": []} - level_keys = { - logging.ERROR: "error", - logging.WARNING: "warn", - logging.DEBUG: "debug", - } - - class _CaptureHandler(logging.Handler): - def emit(self, record: logging.LogRecord): - key = level_keys.get(record.levelno) - if key is not None: - captured[key].append(record.getMessage()) - - app_logger = logging.getLogger("reflex.app") - handler = _CaptureHandler(level=logging.DEBUG) - previous_level = app_logger.level - app_logger.addHandler(handler) - app_logger.setLevel(logging.DEBUG) - yield captured - app_logger.removeHandler(handler) - app_logger.setLevel(previous_level) - - -@pytest.mark.asyncio -async def test_client_error_dispatch_missing_reports_actionable_error( - event_namespace: EventNamespace, frontend_errors: list[str] -): - """A dispatch_function_missing error reports the substate and remediation steps. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - """ - await event_namespace.on_client_error( - "known_sid", - { - "error_type": constants.ClientErrorType.DISPATCH_MISSING, - "message": "Cannot process state update", - "substate": "reflex___state____state.my___state____my_state", - }, - ) - assert len(frontend_errors) == 1 - message = frontend_errors[0] - assert "reflex___state____state.my___state____my_state" in message - assert "skipped" in message.lower() - assert "rebuild" in message.lower() - - -@pytest.mark.asyncio -async def test_client_error_generic_reports_type_and_message( - event_namespace: EventNamespace, frontend_errors: list[str] -): - """A generic client error reports the error type and message. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - """ - await event_namespace.on_client_error( - "known_sid", - { - "error_type": constants.ClientErrorType.STATE_UPDATE, - "message": "boom", - }, - ) - assert len(frontend_errors) == 1 - message = frontend_errors[0] - assert constants.ClientErrorType.STATE_UPDATE in message - assert "boom" in message - - -@pytest.mark.asyncio -@pytest.mark.parametrize("payload", ["not a dict", None, ["list"], 42]) -async def test_client_error_malformed_payload_is_ignored( - event_namespace: EventNamespace, - frontend_errors: list[str], - payload: Any, -): - """Non-dict payloads are dropped without raising or reporting errors. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - payload: The malformed payload to send. - """ - await event_namespace.on_client_error("known_sid", payload) - assert not frontend_errors - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sid", ["known_sid", "unknown_sid"]) -async def test_client_error_no_argument_emit_is_ignored( - event_namespace: EventNamespace, - frontend_errors: list[str], - client_error_console: dict[str, list[str]], - sid: str, -): - """A payload-less client_error emit is dropped like any malformed payload. - - python-socketio dispatches ``emit("client_error")`` without a payload as - ``on_client_error(sid)``, so a missing ``data`` argument must not raise a - TypeError before the anti-abuse guards run. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - client_error_console: Captured console messages. - sid: The Socket.IO session id to report from. - """ - await event_namespace.on_client_error(sid) - assert not frontend_errors - assert not client_error_console["error"] - assert not client_error_console["warn"] - - -@pytest.mark.asyncio -async def test_client_error_unknown_sid_does_not_report_error( - event_namespace: EventNamespace, - frontend_errors: list[str], - client_error_console: dict[str, list[str]], -): - """Errors from sockets without a linked token are not reported. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - client_error_console: Captured console messages. - """ - await event_namespace.on_client_error( - "unknown_sid", - { - "error_type": constants.ClientErrorType.STATE_UPDATE, - "message": "spam from unauthenticated socket", - }, - ) - assert not frontend_errors - assert not client_error_console["error"] - - -@pytest.mark.asyncio -async def test_client_error_values_are_sanitized_and_truncated( - event_namespace: EventNamespace, frontend_errors: list[str] -): - """Control characters are stripped and long messages truncated before reporting. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - """ - evil = "\x1b[31mINJECT\x1b[0m\nFAKE LOG LINE\t" + "A" * 5000 - await event_namespace.on_client_error( - "known_sid", - {"error_type": "custom_type", "message": evil}, - ) - assert len(frontend_errors) == 1 - message = frontend_errors[0] - assert "\x1b" not in message - assert "\n" not in message - assert "\t" not in message - assert len(message) < 700 - - -@pytest.mark.asyncio -async def test_client_error_reporting_is_rate_limited_per_sid( - event_namespace: EventNamespace, frontend_errors: list[str] -): - """A single session cannot flood the backend logs with error reports. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - """ - for _ in range(20): - await event_namespace.on_client_error( - "known_sid", - {"error_type": "custom_type", "message": "spam"}, - ) - assert len(frontend_errors) == EventNamespace._MAX_CLIENT_ERRORS_PER_SID - # Disconnecting removes the counter so the mapping cannot grow unboundedly. - task = event_namespace.on_disconnect("known_sid") - if task is not None: - await task - assert "known_sid" not in event_namespace._client_error_counts - - -@pytest.mark.asyncio -async def test_client_error_reporting_bounded_across_reconnects( - event_namespace: EventNamespace, - frontend_errors: list[str], - client_error_console: dict[str, list[str]], -): - """Reconnecting with fresh SIDs does not grant an unlimited report budget. - - Args: - event_namespace: The event namespace. - frontend_errors: Captured frontend exception handler messages. - client_error_console: Captured console messages. - """ - for reconnect in range(50): - sid = f"sid_{reconnect}" - event_namespace.sid_to_token[sid] = f"token_{reconnect}" - for _ in range(5): - await event_namespace.on_client_error( - sid, {"error_type": "custom_type", "message": "spam"} - ) - assert len(frontend_errors) == EventNamespace._MAX_CLIENT_ERRORS_PER_WINDOW - # Suppression is not silent: one warning is logged when the cap trips, so - # a flooding client cannot invisibly starve reports from other sessions. - assert ( - len([ - msg for msg in client_error_console["warn"] if "suppressing" in msg.lower() - ]) - == 1 - ) - # Once the window elapses, errors are reported again (not silenced forever). - event_namespace._client_error_window_start -= ( - EventNamespace._CLIENT_ERROR_WINDOW_SECONDS + 1 - ) - event_namespace.sid_to_token["sid_fresh"] = "token_fresh" - await event_namespace.on_client_error( - "sid_fresh", {"error_type": "custom_type", "message": "after window"} - ) - assert len(frontend_errors) == EventNamespace._MAX_CLIENT_ERRORS_PER_WINDOW + 1 - - -def test_client_error_event_name_matches_handler(): - """python-socketio dispatches events to on_ methods by naming - convention; this pins the handler to SocketEvent.CLIENT_ERROR. - """ - assert ( - f"on_{constants.SocketEvent.CLIENT_ERROR}" - == EventNamespace.on_client_error.__name__ - ) - - -def test_client_error_constants_match_frontend(): - """The socket event and error types are duplicated as literals in state.js. - - Nothing at runtime keeps the two definitions in sync, so pin them here. - """ - state_js = ( - Path(reflex_base.__file__).parent / ".templates/web/utils/state.js" - ).read_text() - assert ( - f'const CLIENT_ERROR_EVENT = "{constants.SocketEvent.CLIENT_ERROR}"' in state_js - ) - assert ( - f'const ERROR_TYPE_DISPATCH_MISSING = "{constants.ClientErrorType.DISPATCH_MISSING}"' - in state_js - ) - assert ( - f'const ERROR_TYPE_STATE_UPDATE = "{constants.ClientErrorType.STATE_UPDATE}"' - in state_js - ) - - -@pytest_asyncio.fixture -async def event_namespace_with_processor_mock() -> AsyncGenerator[EventNamespace, None]: - """An EventNamespace whose app has a mocked event processor. - - Yields: - The EventNamespace instance. - """ - app = App() - app._event_processor = Mock(enqueue=AsyncMock()) - event_namespace = EventNamespace("/event", app) - yield event_namespace - # The token manager is backed by redis when one is configured; drop the - # tokens these tests link so they do not show up in another test's - # enumeration of the shared instance. Awaited rather than run in a fresh - # loop via asyncio.run: the redis client is bound to the test's loop. - await event_namespace._token_manager.disconnect_all() - - -@pytest.mark.parametrize("owner", ["current", "replacement", "remote", "missing"]) -async def test_disconnect_only_cancels_current_local_socket( - event_namespace_with_processor_mock: EventNamespace, - owner: str, -): - """A stale disconnect must not cancel a replacement connection's page load. - - Args: - event_namespace_with_processor_mock: Namespace with a mocked processor. - owner: The socket currently owning the client token. - """ - namespace = event_namespace_with_processor_mock - manager = namespace._token_manager - manager.sid_to_token["old-sid"] = "token" - if owner != "missing": - manager.token_to_socket["token"] = SocketRecord( - instance_id="remote" if owner == "remote" else manager.instance_id, - sid="new-sid" if owner == "replacement" else "old-sid", - ) - manager.disconnect_token = AsyncMock() - cleanup = namespace.on_disconnect("old-sid") - assert cleanup is not None - await cleanup - cancel = cast(Mock, namespace.app.event_processor._on_disconnect) - if owner == "current": - cancel.assert_called_once_with("token") - else: - cancel.assert_not_called() - assert namespace.on_disconnect("unknown-sid") is None - - -def _connect_environ(token: str) -> dict[str, Any]: - return { - "QUERY_STRING": f"token={token}", - "asgi.scope": { - "headers": [ - (b"origin", b"http://localhost:3000"), - (b"user-agent", b"test-agent"), - ], - "client": ("127.0.0.1", 1234), - }, - } - - -def _client_event_payload() -> dict[str, Any]: - return { - "name": "state.hydrate", - "router_data": {"pathname": "/", "query": {}, "asPath": "/"}, - "payload": {}, - } - - -@pytest.mark.asyncio -async def test_on_event_uses_connect_time_router_data( - token: str, - event_namespace_with_processor_mock: EventNamespace, -): - """on_event merges the connection-scoped router_data gathered at connect. - - Headers, client IP, and session id are computed once in on_connect; the - per-event path must not re-read the connection environ at all. - - Args: - token: A token. - event_namespace_with_processor_mock: The event namespace fixture. - """ - event_namespace = event_namespace_with_processor_mock - await event_namespace.on_connect("sid1", _connect_environ(token)) - assert "sid1" in event_namespace._static_router_data - - # The per-event path must not re-read the connection environ. - event_namespace.app.sio = Mock( - get_environ=Mock(side_effect=AssertionError("environ must not be consulted")) - ) - await event_namespace.on_event("sid1", _client_event_payload()) - - enqueue_mock = cast(AsyncMock, event_namespace.app.event_processor.enqueue) - enqueue_mock.assert_called_once() - enqueued_token, event = enqueue_mock.call_args[0] - assert enqueued_token == token - assert event.router_data[constants.RouteVar.CLIENT_TOKEN] == token - assert event.router_data[constants.RouteVar.SESSION_ID] == "sid1" - assert event.router_data[constants.RouteVar.CLIENT_IP] == "127.0.0.1" - assert event.router_data[constants.RouteVar.HEADERS] == { - "origin": "http://localhost:3000", - "user-agent": "test-agent", - "asgi-scope-client": "127.0.0.1", - } - assert event.router_data[constants.RouteVar.PATH] == "/404" - assert event.router_data[constants.RouteVar.QUERY] == {} - - # Disconnect drops the cached connection data. - event_namespace.on_disconnect("sid1") - assert "sid1" not in event_namespace._static_router_data - - -@pytest.mark.asyncio -async def test_link_token_to_sid_records_the_connecting_identity( - token: str, - event_namespace_with_processor_mock: EventNamespace, - mocker: MockerFixture, -): - """The session var carries the token the state was loaded under. - - Duplicate-token handling hands back a fresh token, and the state is loaded - under it. Leaving `rx_router_session.client_token` empty until the first event - would let anything reading it in between -- a background task, a - shared-state link -- address the wrong state tree. - - Args: - token: A token. - event_namespace_with_processor_mock: The event namespace fixture. - mocker: pytest-mock fixture. - """ - event_namespace = event_namespace_with_processor_mock - state = Mock() - state.router_data = {} - mocker.patch.object( - event_namespace.app.state_manager, - "modify_state", - Mock(return_value=AsyncMock(__aenter__=AsyncMock(return_value=state))), - ) - - # No duplicate: the connecting token is recorded. - await event_namespace.link_token_to_sid("sid1", token) - assert state.router_data[constants.RouteVar.CLIENT_TOKEN] == token - assert state.rx_router_session.client_token == token - assert state.rx_router_session.session_id == "sid1" - - # Duplicate: the *new* token is recorded, not the one the client sent. - # The duplicate branch emits the replacement token to the client, which - # needs a server the bare namespace does not have. - event_namespace.emit = AsyncMock() # pyright: ignore[reportAttributeAccessIssue] - new_token = "a-fresh-token" - mocker.patch.object( - event_namespace._token_manager, - "link_token_to_sid", - AsyncMock(return_value=new_token), - ) - await event_namespace.link_token_to_sid("sid2", token) - assert state.router_data[constants.RouteVar.CLIENT_TOKEN] == new_token - assert state.rx_router_session.client_token == new_token - assert state.rx_router_session.session_id == "sid2" - - -@pytest.mark.asyncio -async def test_on_event_does_not_share_the_cached_headers( - token: str, - event_namespace_with_processor_mock: EventNamespace, -): - """Each event gets its own headers mapping, not the cached one. - - The headers reach `state.router_data`, a plain mutable dict, so sharing - the cached mapping would let a handler mutating it corrupt the connection - cache for every later event on the socket. - - Args: - token: A token. - event_namespace_with_processor_mock: The event namespace fixture. - """ - event_namespace = event_namespace_with_processor_mock - await event_namespace.on_connect("sid1", _connect_environ(token)) - cached_headers = event_namespace._static_router_data["sid1"][ - constants.RouteVar.HEADERS - ] - - await event_namespace.on_event("sid1", _client_event_payload()) - enqueue_mock = cast(AsyncMock, event_namespace.app.event_processor.enqueue) - _, event = enqueue_mock.call_args[0] - event_headers = event.router_data[constants.RouteVar.HEADERS] - - assert event_headers == cached_headers - assert event_headers is not cached_headers - # Mutating what the handler sees must not reach the connection cache. - event_headers["user-agent"] = "mutated" - assert cached_headers["user-agent"] == "test-agent" - - enqueue_mock.reset_mock() - await event_namespace.on_event("sid1", _client_event_payload()) - _, next_event = enqueue_mock.call_args[0] - assert ( - next_event.router_data[constants.RouteVar.HEADERS]["user-agent"] == "test-agent" - ) - - -@pytest.mark.asyncio -async def test_on_event_falls_back_to_environ_without_connect( - token: str, - event_namespace_with_processor_mock: EventNamespace, -): - """on_event computes and caches the static router_data if connect was missed. - - Args: - token: A token. - event_namespace_with_processor_mock: The event namespace fixture. - """ - event_namespace = event_namespace_with_processor_mock - await event_namespace._token_manager.link_token_to_sid(token, "sid1") - event_namespace.app.sio = Mock( - get_environ=Mock(return_value=_connect_environ(token)) - ) - - await event_namespace.on_event("sid1", _client_event_payload()) - await event_namespace.on_event("sid1", _client_event_payload()) - - # The environ is only consulted once; the result is cached for the sid. - event_namespace.app.sio.get_environ.assert_called_once() - enqueue_mock = cast(AsyncMock, event_namespace.app.event_processor.enqueue) - assert enqueue_mock.call_count == 2 - for call in enqueue_mock.call_args_list: - _, event = call[0] - assert event.router_data[constants.RouteVar.SESSION_ID] == "sid1" - assert ( - event.router_data[constants.RouteVar.HEADERS]["user-agent"] == "test-agent" - ) - - @pytest.mark.parametrize("compile_raises", [False, True]) def test_compile_releases_memo_naming_caches( mocker: MockerFixture, compile_raises: bool @@ -5174,214 +4551,6 @@ def fake_compile_app(*_args: Any, **_kwargs: Any) -> bool: assert not _hash_str_encodings -@pytest.mark.asyncio -async def test_on_connect_processes_boot_event_from_auth( - event_namespace: EventNamespace, -): - """The hydrate event carried in the socket.io CONNECT packet is processed on connect. - - As the connection's first event it records the new sid and token on the - state, so the connect does not load and save the state tree for that - first; a connect without a boot event still does. - - Args: - event_namespace: The event namespace. - """ - event_namespace._token_manager = Mock() - event_namespace._token_manager.link_token_to_sid = AsyncMock(return_value=None) - event_namespace.on_event = AsyncMock() - state = Mock(router_data={}) - modify_state = event_namespace.app.state_manager.modify_state = Mock( - return_value=AsyncMock(__aenter__=AsyncMock(return_value=state)) - ) - boot_event = {"name": "state.hydrate_and_load", "payload": {}, "router_data": {}} - - await event_namespace.on_connect( - "new_sid", {"QUERY_STRING": "token=abc"}, {"event": boot_event} - ) - event_namespace._token_manager.link_token_to_sid.assert_awaited_once_with( - "abc", "new_sid" - ) - event_namespace.on_event.assert_awaited_once_with("new_sid", boot_event) - modify_state.assert_not_called() - - # Without a boot event (or without auth at all) nothing is processed, and - # the connect records the new sid and token on the state itself. - event_namespace.on_event.reset_mock() - await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}, None) - await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}) - event_namespace.on_event.assert_not_awaited() - assert modify_state.call_count == 2 - assert state.router_data[constants.RouteVar.SESSION_ID] == "new_sid" - assert state.router_data[constants.RouteVar.CLIENT_TOKEN] == "abc" - - -@pytest.mark.asyncio -async def test_on_connect_accepts_non_utf8_header_value( - event_namespace: EventNamespace, -): - """A non-UTF-8 request header value does not abort websocket connect. - - Args: - event_namespace: The event namespace. - """ - event_namespace._token_manager = Mock() - event_namespace._token_manager.link_token_to_sid = AsyncMock(return_value=None) - event_namespace.on_event = AsyncMock() - - await event_namespace.on_connect( - "new_sid", - { - "QUERY_STRING": "token=abc", - "asgi.scope": {"headers": [(b"x-name", b"\xe9")]}, - }, - {"event": {"name": "state.hydrate_and_load", "payload": {}}}, - ) - - assert ( - event_namespace._static_router_data["new_sid"][constants.RouteVar.HEADERS][ - "x-name" - ] - == "é" - ) - - -def _schema_checking_namespace( - event_namespace: EventNamespace, backend_schema: str -) -> tuple[AsyncMock, AsyncMock]: - """Prepare an event namespace for on_connect with a known backend schema. - - Args: - event_namespace: The event namespace. - backend_schema: The frontend schema hash the backend's state tree has. - - Returns: - The mocks installed as the namespace's ``on_event`` and ``emit``. - """ - event_namespace._token_manager = Mock() - event_namespace._token_manager.link_token_to_sid = AsyncMock(return_value=None) - on_event = event_namespace.on_event = AsyncMock() - emit = event_namespace.emit = AsyncMock() - cast(Any, event_namespace.app)._state = Mock( - _to_frontend_schema=Mock(return_value=backend_schema) - ) - return on_event, emit - - -@pytest.mark.asyncio -async def test_on_connect_reports_state_schema_mismatch( - event_namespace: EventNamespace, caplog: pytest.LogCaptureFixture -): - """A frontend compiled for another state tree is told once, and still hydrates. - - Args: - event_namespace: The event namespace. - caplog: pytest fixture capturing backend log records. - """ - on_event, emit = _schema_checking_namespace(event_namespace, "backend-schema") - boot_event = {"name": "state.hydrate_and_load", "payload": {}, "router_data": {}} - - with caplog.at_level(logging.WARNING, logger="reflex"): - await event_namespace.on_connect( - "new_sid", - {"QUERY_STRING": "token=abc"}, - {"event": boot_event, "state_schema": "frontend-schema"}, - ) - - emit.assert_awaited_once_with( - "state_mismatch", - {"backend": "backend-schema", "reload_cooldown_ms": 10_000}, - to="new_sid", - ) - on_event.assert_awaited_once_with("new_sid", boot_event) - app_warnings = [ - r.getMessage() - for r in caplog.records - if r.name == "reflex.app" and r.levelno == logging.WARNING - ] - assert len(app_warnings) == 1 - assert "frontend-schema" in app_warnings[0] - assert "backend-schema" in app_warnings[0] - assert "new_sid" in app_warnings[0] - - # The client-supplied hash is sanitized and bounded before it is logged. - caplog.clear() - await event_namespace.on_connect( - "new_sid", - {"QUERY_STRING": "token=abc"}, - {"event": boot_event, "state_schema": "x\nforged line [bold]" + "y" * 500}, - ) - (message,) = [ - r.getMessage() - for r in caplog.records - if r.name == "reflex.app" and r.levelno == logging.WARNING - ] - assert "forged" in message - assert "\n" not in message - assert "\\[bold]" in message - assert "(truncated)" in message - assert len(message) < 400 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "auth", - [ - pytest.param( - {"event": {"name": "state.hydrate_and_load"}, "state_schema": "same"}, - id="matching", - ), - pytest.param({"event": {"name": "state.hydrate_and_load"}}, id="old-frontend"), - pytest.param(None, id="no-auth"), - ], -) -async def test_on_connect_sends_no_mismatch_without_one( - event_namespace: EventNamespace, auth: dict[str, Any] | None -): - """A matching or missing frontend schema sends nothing to the client. - - Args: - event_namespace: The event namespace. - auth: The CONNECT auth payload. - """ - on_event, emit = _schema_checking_namespace(event_namespace, "same") - state = Mock(router_data={}) - event_namespace.app.state_manager.modify_state = Mock( - return_value=AsyncMock(__aenter__=AsyncMock(return_value=state)) - ) - - await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}, auth) - - emit.assert_not_awaited() - assert on_event.await_count == (1 if auth else 0) - - -@pytest.mark.asyncio -async def test_on_connect_unlinks_token_when_boot_event_fails( - event_namespace: EventNamespace, -): - """A boot event that fails to process drops the sid/token link before refusing the connect. - - Args: - event_namespace: The event namespace. - """ - event_namespace._token_manager = Mock() - event_namespace._token_manager.link_token_to_sid = AsyncMock(return_value=None) - event_namespace._token_manager.disconnect_token = AsyncMock() - event_namespace._token_manager.sid_to_token = {"new_sid": "abc"} - event_namespace.on_event = AsyncMock(side_effect=ValueError("bad boot event")) - - with pytest.raises(ValueError, match="bad boot event"): - await event_namespace.on_connect( - "new_sid", {"QUERY_STRING": "token=abc"}, {"event": {"name": "x"}} - ) - event_namespace._token_manager.disconnect_token.assert_awaited_once_with( - "abc", "new_sid" - ) - # The connection-scoped router data cached for the refused sid goes too. - assert "new_sid" not in event_namespace._static_router_data - - def test_call_app_wraps_with_otel_asgi_middleware(): """The app's ASGI callable is wrapped when instrumentation installs a middleware.""" app = App() @@ -5395,54 +4564,6 @@ def test_call_app_wraps_with_otel_asgi_middleware(): assert wrapped == [api] -def test_sio_json_records_message_sizes(otel_metrics): - """Socket.IO packet serialization records UTF-8 sizes in both directions.""" - data = _sio_dumps({"a": "é"}, separators=(",", ":")) - assert data == '{"a":"é"}' - assert _sio_loads(data) == {"a": "é"} - assert _sio_loads(data.encode()) == {"a": "é"} - # ASCII payloads take the fast path that sizes without encoding a copy. - ascii_data = _sio_dumps({"b": "x"}, separators=(",", ":")) - assert _sio_loads(ascii_data) == {"b": "x"} - points = { - p.attributes[otel.ATTR_NETWORK_IO_DIRECTION]: p.sum - for p in metric_points(otel_metrics, otel.METRIC_WEBSOCKET_MESSAGE_SIZE) - } - size = len(data.encode()) - assert size == len(data) + 1 - assert points == { - "transmit": size + len(ascii_data), - "receive": 2 * size + len(ascii_data), - } - - -@pytest.mark.asyncio -async def test_on_event_uses_frontend_traceparent(otel_exporter): - """A traceparent in the event payload becomes the parent of the event span.""" - mock_app = unittest.mock.Mock() - mock_app.router.return_value = "/" - mock_app.sio.get_environ.return_value = { - "asgi.scope": {"headers": [], "client": ("127.0.0.1", 1)} - } - seen: list = [] - - async def enqueue(token, event): - await asyncio.sleep(0) - seen.append(trace.get_current_span().get_span_context()) - - mock_app.event_processor.enqueue = enqueue - ns = EventNamespace(namespace="/", app=mock_app) - ns._token_manager.sid_to_token["sid"] = "tok" - - traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - with active_tracer().start_as_current_span("websocket"): - await ns.on_event("sid", {"name": "state.h", "traceparent": traceparent}) - await ns.on_event("sid", {"name": "state.h"}) - remote, fresh = seen - assert f"{remote.trace_id:032x}" == "0af7651916cd43dd8448eb211c80319c" - assert not fresh.is_valid - - @pytest.mark.asyncio async def test_upload_handler_span_joins_the_request_trace( tmp_path: Path, @@ -5544,24 +4665,6 @@ async def send(message): # noqa: RUF029 assert handler_parent.span_id == request_context.span_id -@pytest.mark.asyncio -async def test_connect_disconnect_counts_connections(otel_metrics): - """Connect and disconnect adjust the open connection gauge.""" - mock_app = unittest.mock.Mock() - mock_app._state = None - ns = EventNamespace(namespace="/", app=mock_app) - ns.emit = unittest.mock.AsyncMock() - await ns.on_connect("sid1", {"QUERY_STRING": "token=t1"}) - await ns.on_connect("sid2", {"QUERY_STRING": "token=t2"}) - task = ns.on_disconnect("sid1") - if task is not None: - await task - (point,) = metric_points(otel_metrics, otel.METRIC_WEBSOCKET_CONNECTIONS) - assert point.value == 1 - # Release t2 so a shared token store (redis) does not leak into other tests. - await ns._token_manager.disconnect_all() - - def test_compile_installs_browser_plugin( compilable_app: tuple[App, Path], mocker: MockerFixture, @@ -5795,3 +4898,74 @@ def reader(): for future in readers: future.result() assert json.loads(marker.read_text()) == routes + + +class _ProbeChannel(rx.channels.Channel): + """A channel that ignores everything, for registration tests.""" + + name = "probe" + + async def on_message(self, session, event, data, buffers) -> None: + """Ignore inbound messages.""" + return + + +def test_register_channel_serves_the_channel(): + """A registered channel is reachable by name for the transport.""" + app = App(enable_state=True) + channel = _ProbeChannel() + + app.register_channel(channel) + + assert app._channels == {"probe": channel} + + +def test_register_channel_rejects_a_duplicate_name(): + """Two channels cannot claim the same name.""" + app = App(enable_state=True) + app.register_channel(_ProbeChannel()) + + with pytest.raises(RuntimeError, match="already registered"): + app.register_channel(_ProbeChannel()) + + +@pytest.mark.parametrize("state", [None, State]) +def test_register_channel_skips_without_the_event_websocket( + state: type[State] | None, caplog: pytest.LogCaptureFixture +): + """Without state there is no transport, whatever `_state` was passed. + + A supplied `_state` does not set one up on its own: `enable_state=False` + skips the setup that creates the event namespace and its route. A package + registering its channel from a plugin cannot know that, so the app warns + and keeps serving rather than failing to start. + """ + with RegistrationContext.get().fork(): + app = App(_state=state, enable_state=False) + assert app.event_namespace is None + + with caplog.at_level(logging.WARNING): + app.register_channel(_ProbeChannel()) + + assert app._channels == {} + assert "needs the event websocket" in caplog.text + + +def test_register_channel_skips_without_the_websocket_transport( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + """Channels are a plain-WebSocket feature; Socket.IO cannot carry them. + + The namespace's capability is patched rather than a Socket.IO app built: + that would need the optional python-socketio package, which this check + has nothing to do with. + """ + app = App(enable_state=True) + assert app.event_namespace is not None + monkeypatch.setattr(app.event_namespace, "supports_channels", False) + + with caplog.at_level(logging.WARNING): + app.register_channel(_ProbeChannel()) + + assert app._channels == {} + assert "requires the plain WebSocket transport" in caplog.text diff --git a/tests/units/test_channels.py b/tests/units/test_channels.py new file mode 100644 index 00000000000..0bb7131c2b8 --- /dev/null +++ b/tests/units/test_channels.py @@ -0,0 +1,330 @@ +"""Tests for the channel API in reflex/channels.py.""" + +from collections.abc import Sequence +from typing import Any + +import pytest + +from reflex.channels import ( + MAX_MESSAGE_BUFFERS, + RESERVED_EVENTS, + Channel, + ChannelSession, + validate_channel_name, +) + + +class CollectingChannel(Channel): + """A channel that records every message its transport was asked to send.""" + + name = "probe" + + def __init__(self): + """Initialize the channel and its recorded sends.""" + super().__init__() + self.sent: list[tuple[list[str], str, str, Any, list[bytes]]] = [] + # What the stand-in transport reports for each send. + self.accepted = True + + async def on_message( + self, session: ChannelSession, event: str, data: Any, buffers: list[bytes] + ) -> None: + """Ignore inbound messages; these tests drive the channel directly.""" + return + + async def _record( + self, + sids: Sequence[str], + channel: str, + event: str, + data: Any, + buffers: Any, + ) -> bool: + """Stand in for the transport's send callable. + + Returns: + Whether the transport accepted a recipient. + """ + self.sent.append((list(sids), channel, event, data, list(buffers))) + return self.accepted + + def session(self, sid: str, client_token: str = "tok") -> ChannelSession: + """Open a session wired to the recording sender. + + Args: + sid: The session id. + client_token: The client token. + + Returns: + The open session. + """ + return self.open_session(sid, client_token, self._record) + + +@pytest.mark.parametrize("name", ["xy", "_xy", "/_xy", "a.b-c:d/e", "x" * 64]) +def test_valid_channel_names(name: str): + """Names usable on the wire pass validation.""" + validate_channel_name(name) + + +@pytest.mark.parametrize("name", ["", "x" * 65, "has space", 'quote"', "new\nline"]) +def test_invalid_channel_names(name: str): + """Names that would need escaping on the wire are rejected.""" + with pytest.raises(ValueError, match="Invalid channel name"): + validate_channel_name(name) + + +def test_channel_rejects_an_unusable_name(): + """A channel declaring an unusable name fails at construction.""" + + class BadChannel(CollectingChannel): + name = "not a name" + + with pytest.raises(ValueError, match="Invalid channel name"): + BadChannel() + + +@pytest.mark.asyncio +async def test_session_send_reaches_the_transport(): + """A session's send is handed to the transport with its channel name.""" + channel = CollectingChannel() + session = channel.session("sid1") + + await session.send("payload", {"fig": "f1"}, [b"\x00"]) + + assert channel.sent == [(["sid1"], "probe", "payload", {"fig": "f1"}, [b"\x00"])] + + +@pytest.mark.asyncio +async def test_room_fan_out_reaches_members_only(): + """A room's members are handed to the transport together, and nobody else. + + Together, because the frame is then serialized once however many members + the room has, rather than once per member. + """ + channel = CollectingChannel() + first = channel.session("sid1") + second = channel.session("sid2") + channel.session("sid3") + first.join("fig:1") + second.join("fig:1") + + await channel.send_to_room("fig:1", "push", {"n": 1}) + + assert len(channel.sent) == 1 + assert sorted(channel.sent[0][0]) == ["sid1", "sid2"] + + +@pytest.mark.asyncio +async def test_leaving_a_room_stops_delivery(): + """A session that left a room no longer receives its messages.""" + channel = CollectingChannel() + session = channel.session("sid1") + session.join("fig:1") + session.leave("fig:1") + + await channel.send_to_room("fig:1", "push") + + assert channel.sent == [] + # The empty room is not kept around. + assert channel._rooms == {} + + +@pytest.mark.asyncio +async def test_sending_to_an_unknown_room_is_a_noop(): + """Fan-out to a room nobody joined does nothing.""" + channel = CollectingChannel() + + await channel.send_to_room("fig:missing", "push") + + assert channel.sent == [] + + +@pytest.mark.asyncio +async def test_send_to_token_reports_delivery(): + """Token-addressed sends reach that client's sessions, and report misses.""" + channel = CollectingChannel() + channel.session("sid1", client_token="tok1") + channel.session("sid2", client_token="tok2") + + assert await channel.send_to_token("tok1", "push", {"n": 1}) is True + assert await channel.send_to_token("missing", "push") is False + assert [sent[0] for sent in channel.sent] == [["sid1"]] + + +@pytest.mark.asyncio +async def test_send_to_token_reports_a_send_the_transport_dropped(): + """A token send reports a miss when the transport accepted no recipient. + + The transport skips a session whose token moved to another socket, which + the channel keeps tracking until the connection's teardown forgets it. + """ + channel = CollectingChannel() + channel.session("sid1", client_token="tok1") + channel.accepted = False + + assert await channel.send_to_token("tok1", "push") is False + assert [sent[0] for sent in channel.sent] == [["sid1"]] + + +@pytest.mark.asyncio +async def test_forgetting_a_session_clears_rooms_and_tokens(): + """A forgotten session leaves no room membership or token entry behind.""" + channel = CollectingChannel() + session = channel.session("sid1", client_token="tok1") + session.join("fig:1") + + channel.forget_session(session) + + assert channel._rooms == {} + assert channel._sessions == {} + assert await channel.send_to_token("tok1", "push") is False + await channel.send_to_room("fig:1", "push") + assert channel.sent == [] + + +def test_sessions_of_one_token_are_tracked_together(): + """Two tabs sharing a token both stay reachable until each is forgotten.""" + channel = CollectingChannel() + first = channel.session("sid1", client_token="tok1") + second = channel.session("sid2", client_token="tok1") + + channel.forget_session(first) + + assert channel._sessions == {"tok1": {second}} + + +def test_forgotten_session_reports_itself_closed(): + """A session whose client went away reports it, for handlers mid-await.""" + channel = CollectingChannel() + session = channel.session("sid1") + assert session.open is True + + channel.forget_session(session) + + assert session.open is False + + +@pytest.mark.asyncio +async def test_send_rejects_more_attachments_than_a_frame_holds(): + """An oversized message raises instead of building an unsendable frame.""" + channel = CollectingChannel() + session = channel.session("sid1") + buffers = [b"\x00"] * (MAX_MESSAGE_BUFFERS + 1) + + with pytest.raises(ValueError, match="attachments"): + await session.send("push", None, buffers) + + assert channel.sent == [] + + +@pytest.mark.asyncio +async def test_room_fan_out_rejects_oversized_messages_before_delivering(): + """A fan-out that cannot be framed reaches nobody, rather than some.""" + channel = CollectingChannel() + channel.session("sid1").join("all") + channel.session("sid2").join("all") + + with pytest.raises(ValueError, match="attachments"): + await channel.send_to_room( + "all", "push", None, [b""] * (MAX_MESSAGE_BUFFERS + 1) + ) + + assert channel.sent == [] + + +@pytest.mark.asyncio +async def test_send_accepts_the_full_attachment_budget(): + """The limit is inclusive, so a channel can use all of it.""" + channel = CollectingChannel() + session = channel.session("sid1") + + await session.send("push", None, [b""] * MAX_MESSAGE_BUFFERS) + + assert len(channel.sent) == 1 + + +@pytest.mark.parametrize("name", [42, b"bytes", None]) +def test_channel_with_an_unusable_name_reports_it(name: Any): + """A non-string name fails with the explicit error, not a TypeError.""" + + class BadChannel(CollectingChannel): + pass + + BadChannel.name = name # pyright: ignore[reportAttributeAccessIssue] + + with pytest.raises(ValueError, match="Invalid channel name"): + BadChannel() + + +def test_channel_without_a_name_reports_it(): + """A channel that never declared a name fails the same way.""" + + class NamelessChannel(Channel): + async def on_message( + self, session: ChannelSession, event: str, data: Any, buffers: list[bytes] + ) -> None: + """Ignore inbound messages.""" + return + + with pytest.raises(ValueError, match="Invalid channel name"): + NamelessChannel() + + +@pytest.mark.asyncio +async def test_join_after_close_leaves_no_member_behind(): + """A handler resuming after the socket went away cannot rejoin a room. + + The transport drops a closed session's rooms once; a later join would put + it back where nothing removes it again, and every fan-out would keep + addressing the dead connection. + """ + channel = CollectingChannel() + session = channel.session("sid1") + channel.forget_session(session) + + session.join("room") + + assert channel._rooms == {} + await channel.send_to_room("room", "tick", {"n": 1}) + assert channel.sent == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event", + [*sorted(RESERVED_EVENTS), "_open", "_opened", "_close", "_error", "_future"], +) +async def test_send_rejects_reserved_message_names(event: str): + """A message may not impersonate the handle's lifecycle or the protocol. + + A relayed ``_error`` or ``_opened`` would reach every other member as a + lifecycle event, so the protocol's underscore names are reserved too. + """ + channel = CollectingChannel() + session = channel.session("sid1") + + with pytest.raises(ValueError, match="reserved"): + await session.send(event, {"anything": True}) + + assert channel.sent == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("event", "buffers"), + [("_error", ()), ("push", [b""] * (MAX_MESSAGE_BUFFERS + 1)), (42, ())], +) +async def test_fan_out_validates_even_when_nobody_listens(event: Any, buffers: Any): + """A message that could never be sent fails whether or not anyone is in the room. + + Otherwise the mistake only surfaces once a client happens to join. + """ + channel = CollectingChannel() + + with pytest.raises(ValueError, match="Channel message"): + await channel.send_to_room("empty", event, None, buffers) + with pytest.raises(ValueError, match="Channel message"): + await channel.send_to_token("nobody", event, None, buffers) + + assert channel.sent == [] diff --git a/tests/units/test_config.py b/tests/units/test_config.py index 4ec4bc3d68c..1643402656b 100644 --- a/tests/units/test_config.py +++ b/tests/units/test_config.py @@ -196,6 +196,33 @@ def test_update_from_env_cors( ] +@pytest.mark.parametrize( + ("origins", "expected"), + [ + ("https://myapp.com", ["https://myapp.com"]), + ("https://a.com, https://b.dev", ["https://a.com", "https://b.dev"]), + ], +) +def test_cors_allowed_origins_string_lists_origins( + base_config_values: dict[str, Any], origins: str, expected: list[str] +): + """A plain string lists origins as the environment variable does. + + Left a string, membership is a substring test: https://myapp.co would pass + an origin check for https://myapp.com. + + Args: + base_config_values: Config values. + origins: The configured origins. + expected: The origins they list. + """ + config = rx.Config(**base_config_values, cors_allowed_origins=origins) + + assert list(config.cors_allowed_origins) == expected + assert "https://myapp.co" not in config.cors_allowed_origins + assert "https://b.de" not in config.cors_allowed_origins + + def test_update_from_env_frontend_compression_formats( base_config_values: dict[str, Any], monkeypatch: pytest.MonkeyPatch, diff --git a/tests/units/test_event_namespace.py b/tests/units/test_event_namespace.py new file mode 100644 index 00000000000..c2d48f02b19 --- /dev/null +++ b/tests/units/test_event_namespace.py @@ -0,0 +1,2342 @@ +"""Tests for the plain WebSocket event transport in reflex/event_namespace.py.""" + +import array +import asyncio +import base64 +import json +import logging +import os +import re +import shutil +import subprocess +import sys +from collections.abc import Awaitable, Callable, Iterable +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, Mock + +import pytest +from opentelemetry import trace +from reflex_base import constants, otel +from reflex_base.config import get_config +from reflex_base.registry import RegistrationContext +from starlette.routing import WebSocketRoute + +from reflex import event_namespace +from reflex.app import App +from reflex.channels import ( + MAX_MESSAGE_BUFFERS, + RESERVED_EVENTS, + Channel, + ChannelSession, +) +from reflex.event_namespace import ( + CHANNEL_ERROR_MESSAGE, + CLOSE_MESSAGE, + CONNECT_MESSAGE, + HANDSHAKE_MESSAGE, + MAX_MESSAGE_SIZE_REASON, + OPEN_MESSAGE, + OPENED_MESSAGE, + PING_MESSAGE, + PONG_MESSAGE, + PROTOCOL_VERSION, + WebsocketEventNamespace, + _Connection, + _decode_asgi_headers, + decode_channel_frame, + encode_channel_frame, +) +from reflex.istate.data import HeaderData +from reflex.utils.token_manager import SocketRecord + +from .conftest import active_tracer, metric_points + +_DISCONNECT = object() + +# A connect frame without a boot event. +CONNECT_FRAME = [CONNECT_MESSAGE] + +WEB_TEMPLATE_ROOT = ( + Path(__file__).parents[2] / "packages/reflex-base/src/reflex_base/.templates/web" +) +WEBSOCKET_JS_TEMPLATE = WEB_TEMPLATE_ROOT / "utils/helpers/websocket.js" + + +class FakeWebSocket: + """Minimal stand-in for a starlette WebSocket.""" + + def __init__( + self, + query_string: bytes = b"token=tok1", + origin: str | None = None, + subprotocols: list[str] | None = None, + ): + """Initialize the fake websocket.""" + self.scope: dict[str, Any] = { + "type": "websocket", + "query_string": query_string, + "subprotocols": subprotocols or [], + "headers": [(b"host", b"localhost")], + "client": ("127.0.0.1", 1234), + } + self.headers = {"origin": origin} if origin is not None else {} + self.sent: list[Any] = [] + # The text frames as written, before parsing. + self.raw_sent: list[str] = [] + self.accepted_subprotocol: str | None = None + self.accepted = False + self.close_code: int | None = None + self.close_reason: str | None = None + # Starlette refuses a second close; set to model one already sent. + self.refuse_close = False + self._incoming: asyncio.Queue = asyncio.Queue() + + async def accept(self, subprotocol: str | None = None): + """Record the accept call.""" + self.accepted = True + self.accepted_subprotocol = subprotocol + + async def send_text(self, text: str): + """Record an outgoing frame.""" + self.raw_sent.append(text) + self.sent.append(json.loads(text)) + + async def send_bytes(self, data: bytes): + """Record an outgoing binary frame.""" + self.sent.append(data) + + async def close(self, code: int = 1000, reason: str | None = None): + """Record the close call. + + Raises: + RuntimeError: If the connection was already closed. + """ + if self.refuse_close: + message = 'Cannot call "send" once a close message has been sent.' + raise RuntimeError(message) + self.close_code = code + self.close_reason = reason + + async def receive(self) -> dict[str, Any]: + """Return the next queued frame as an ASGI message. + + Returns: + The ASGI websocket message. + """ + # A real receive suspends on the network, which lets the writer run. + await asyncio.sleep(0) + item = await self._incoming.get() + if item is _DISCONNECT: + return {"type": "websocket.disconnect", "code": 1000} + if isinstance(item, bytes): + return {"type": "websocket.receive", "bytes": item} + return {"type": "websocket.receive", "text": item} + + def feed(self, *frames: Any, connect: bool = True, disconnect: bool = True): + """Queue incoming frames (lists are JSON-encoded) and a disconnect. + + The session's connect frame goes first unless ``connect`` is False. + """ + for frame in (CONNECT_FRAME, *frames) if connect else frames: + self._incoming.put_nowait( + frame if isinstance(frame, (str, bytes)) else json.dumps(frame) + ) + if disconnect: + self._incoming.put_nowait(_DISCONNECT) + + +@pytest.fixture +def namespace(namespace_app: Mock) -> WebsocketEventNamespace: + """A namespace serving the mock app. + + Returns: + The namespace. + """ + return WebsocketEventNamespace("/_event", namespace_app) + + +async def _drain_tasks(): + """Let pending disconnect-cleanup tasks run to completion.""" + for _ in range(3): + await asyncio.sleep(0) + + +async def serve( + namespace: WebsocketEventNamespace, + *frames: Any, + connect: bool = True, + websocket: FakeWebSocket | None = None, +) -> FakeWebSocket: + """Serve one session that sends the given frames and hangs up, then clean up. + + Args: + namespace: The namespace serving it. + *frames: The frames the client sends, after its connect frame unless + ``connect`` is False. + connect: Whether the client opens with a connect frame. + websocket: The fake websocket to serve; a new one if None. + + Returns: + The websocket, recording what the server sent. + """ + websocket = websocket or FakeWebSocket() + websocket.feed(*frames, connect=connect) + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + return websocket + + +@pytest.mark.asyncio +async def test_handshake_and_token_link(namespace: WebsocketEventNamespace): + """The server links the token from the query, then acknowledges with the handshake.""" + websocket = FakeWebSocket(subprotocols=["0.0.1"]) + websocket.feed() + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + + assert websocket.accepted + assert websocket.accepted_subprotocol == "0.0.1" + assert websocket.sent[0][0] == HANDSHAKE_MESSAGE + assert set(websocket.sent[0][1]) == { + "ping_interval", + "ping_timeout", + "protocol", + "max_message_size", + } + assert websocket.sent[0][1]["protocol"] == PROTOCOL_VERSION + await _drain_tasks() + # The session was linked and unlinked again on disconnect. + assert "tok1" not in namespace.token_to_sid + + +@pytest.mark.asyncio +async def test_boot_event_is_the_sessions_first_event( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + linked_state: Mock, + boot_event: dict[str, Any], +): + """The boot event in the connect frame is processed right after the handshake. + + As the session's first event it records the new sid and token on the + state, so the connect does not load and save the state tree for that; a + connect without a boot event still does. + """ + modify_state = namespace_app.state_manager.modify_state + click = {"name": "state.on_click", "payload": {}, "router_data": {}} + websocket = await serve( + namespace, + [CONNECT_MESSAGE, {"event": boot_event}], + ["event", click], + connect=False, + ) + + assert websocket.sent[0][0] == HANDSHAKE_MESSAGE + assert [ + call.args[1].name + for call in namespace_app.event_processor.enqueue.await_args_list + ] == ["state.hydrate_and_load", "state.on_click"] + modify_state.assert_not_called() + + websocket = await serve(namespace) + + modify_state.assert_called_once() + + +@pytest.mark.asyncio +async def test_state_schema_mismatch_is_reported_before_the_handshake( + namespace: WebsocketEventNamespace, namespace_app: Mock, boot_event: dict[str, Any] +): + """A frontend compiled for another state tree is told so before it hydrates. + + The session still hydrates; a frontend built for the backend's state tree + is told nothing. + """ + namespace_app._state = Mock(_to_frontend_schema=Mock(return_value="backend-schema")) + websocket = await serve( + namespace, + [CONNECT_MESSAGE, {"event": boot_event, "state_schema": "frontend-schema"}], + connect=False, + ) + + assert websocket.sent[0] == [ + "state_mismatch", + {"backend": "backend-schema", "reload_cooldown_ms": 10_000}, + ] + assert websocket.sent[1][0] == HANDSHAKE_MESSAGE + namespace_app.event_processor.enqueue.assert_awaited_once() + + websocket = await serve( + namespace, + [CONNECT_MESSAGE, {"event": boot_event, "state_schema": "backend-schema"}], + connect=False, + ) + + assert [frame[0] for frame in websocket.sent] == [HANDSHAKE_MESSAGE] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("frames", "close_code"), + [ + ([["ping"]], 1002), + (["not json"], 1002), + ([[42]], 1002), + ([b"\x00"], 1002), + # The client went away before opening a session. + ([], None), + ], +) +async def test_session_opens_only_with_a_connect_frame( + namespace: WebsocketEventNamespace, + frames: list[Any], + close_code: int | None, +): + """A socket whose first frame is no connect frame never links a token.""" + websocket = await serve(namespace, *frames, connect=False) + + assert websocket.close_code == close_code + assert websocket.sent == [] + assert namespace.token_to_sid == {} + assert namespace._connections == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("connect", [True, False]) +async def test_oversized_integer_closes_with_protocol_error( + namespace: WebsocketEventNamespace, connect: bool +): + """An integer literal past the parser's digit limit is a malformed frame. + + The JSON decoder raises a plain ValueError for it, not a JSONDecodeError; + escaping the handler, it would log an ASGI traceback per connection. + """ + frame = "[" + json.dumps(CONNECT_MESSAGE) + ", " + "1" * 4301 + "]" + websocket = FakeWebSocket() + if connect: + websocket.feed(frame, connect=False) + else: + websocket.feed('["ping", ' + "1" * 4301 + "]") + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert websocket.close_code == 1002 + + +@pytest.mark.asyncio +async def test_frames_are_serialized_compactly(namespace: WebsocketEventNamespace): + """Frames carry no separator whitespace, like the Socket.IO packets they replace.""" + websocket = await serve(namespace, ["ping"], [OPEN_MESSAGE, None, "nope"]) + + assert len(websocket.raw_sent) == 3 + for text in websocket.raw_sent: + assert text == json.dumps(json.loads(text), separators=(",", ":")) + + +@pytest.mark.asyncio +async def test_socket_without_a_connect_frame_times_out( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """A socket that never opens its session ends within the heartbeat window.""" + monkeypatch.setenv("REFLEX_SOCKET_INTERVAL", "10ms") + monkeypatch.setenv("REFLEX_SOCKET_TIMEOUT", "10ms") + websocket = FakeWebSocket() + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + + assert websocket.close_code == 1008 + assert websocket.sent == [] + + +@pytest.mark.asyncio +async def test_oversize_connect_frame_closes_connection( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """The size limit applies to the connect frame too, and the close names it. + + The client has no handshake yet to learn the limit from, so it reads it + from the close reason to report its oversized hydrate. + """ + monkeypatch.setenv("REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE", "10") + websocket = FakeWebSocket() + websocket.feed() + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + + assert websocket.close_code == 1009 + assert websocket.close_reason == f"{MAX_MESSAGE_SIZE_REASON}10" + assert websocket.sent == [] + + +@pytest.mark.asyncio +async def test_undeserializable_boot_event_closes_connection( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """A boot event that fails deserialization ends the session it opened. + + The handshake queued for it is dropped with the session, so the client + sees a failed connect and retries with backoff. + """ + websocket = await serve( + namespace, [CONNECT_MESSAGE, {"event": "not an event"}], ["ping"], connect=False + ) + + assert websocket.close_code == 1002 + assert websocket.sent == [] + namespace_app.event_processor.enqueue.assert_not_awaited() + assert "tok1" not in namespace.token_to_sid + + +@pytest.mark.asyncio +async def test_boot_event_handler_error_ends_the_session( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + caplog: pytest.LogCaptureFixture, + boot_event: dict[str, Any], +): + """A server-side failure on the boot event is logged and ends the session. + + Kept open, the session would leave the client waiting for a hydrate that + never comes; closed, the client reconnects and retries, as on Socket.IO. + """ + namespace_app.event_processor.enqueue.side_effect = RuntimeError("server bug") + with caplog.at_level(logging.ERROR, logger="reflex.event_namespace"): + websocket = await serve( + namespace, [CONNECT_MESSAGE, {"event": boot_event}], ["ping"], connect=False + ) + + assert websocket.close_code == 1011 + assert ["ping", "pong"] not in websocket.sent + assert "tok1" not in namespace.token_to_sid + assert any( + "Error handling the boot event" in record.getMessage() + for record in caplog.records + ) + + +class ReceiveBoundWebSocket(FakeWebSocket): + """Models granian: a close sent while a receive is pending never completes. + + Only the task that receives can close such a socket, once its receive is + no longer pending. + """ + + def __init__(self, **kwargs: Any): + """Initialize the fake websocket.""" + super().__init__(**kwargs) + self.receiving = False + + async def receive(self) -> dict[str, Any]: + """Return the next queued frame, marking the receive as pending. + + Returns: + The ASGI websocket message. + """ + self.receiving = True + try: + return await super().receive() + finally: + self.receiving = False + + async def close(self, code: int = 1000, reason: str | None = None): + """Record the close, unless a receive is pending.""" + if self.receiving: + await asyncio.Event().wait() + await super().close(code, reason) + + +class StalledWebSocket(FakeWebSocket): + """A peer that stopped reading: once its buffers fill, every write blocks.""" + + async def send_text(self, text: str): + """Block forever, like a write to a full socket buffer.""" + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_a_silent_client_is_closed_by_its_own_session( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """The liveness timeout ends the session even where only its task can close. + + A close sent from the heartbeat task while the session waits in receive() + never completes on granian, so the session and its token would leak. + """ + monkeypatch.setenv("REFLEX_SOCKET_INTERVAL", "20ms") + monkeypatch.setenv("REFLEX_SOCKET_TIMEOUT", "20ms") + websocket = ReceiveBoundWebSocket() + websocket.feed(disconnect=False) + + await asyncio.wait_for(namespace.handle_websocket(websocket), 2) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert websocket.close_code == 1001 + assert namespace.token_to_sid == {} + + +@pytest.mark.asyncio +async def test_a_client_that_stops_reading_does_not_block_senders( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """Emitting to a stalled client returns at once, and its backlog is bounded. + + The emit used to await the socket write, so every coroutine sending to the + client -- a shared state's fan-out under its lock, the lost-and-found + loop -- froze with it. + """ + monkeypatch.setattr(event_namespace, "_MAX_SEND_BACKLOG", 1000) + linked = asyncio.Event() + handle_connect = namespace.handle_connect + + async def connect_and_signal(*args: Any, **kwargs: Any) -> None: + await handle_connect(*args, **kwargs) + linked.set() + + monkeypatch.setattr(namespace, "handle_connect", connect_and_signal) + websocket = StalledWebSocket() + websocket.feed(disconnect=False) + session = asyncio.create_task(namespace.handle_websocket(websocket)) # pyright: ignore[reportArgumentType] + await asyncio.wait_for(linked.wait(), 1) + sid = namespace.token_to_sid["tok1"] + + for _ in range(20): + await asyncio.wait_for(namespace.emit("event", {"x": "y" * 100}, to=sid), 1) + + # Over the backlog bound, the session is dropped rather than buffering on. + await asyncio.wait_for(session, 2) + await _drain_tasks() + assert websocket.close_code == 1008 + assert namespace.token_to_sid == {} + + +@pytest.mark.asyncio +async def test_a_stalled_client_still_times_out( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """The heartbeat's own ping cannot block the liveness check that closes it.""" + monkeypatch.setenv("REFLEX_SOCKET_INTERVAL", "20ms") + monkeypatch.setenv("REFLEX_SOCKET_TIMEOUT", "20ms") + websocket = StalledWebSocket() + websocket.feed(disconnect=False) + + await asyncio.wait_for(namespace.handle_websocket(websocket), 2) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert websocket.close_code == 1001 + + +@pytest.mark.asyncio +async def test_event_is_enqueued( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """An incoming event frame reaches the app's event processor.""" + await serve( + namespace, + [ + "event", + { + "token": "tok1", + "name": "state.on_click", + "payload": {}, + "router_data": {}, + }, + ], + ) + + namespace_app.event_processor.enqueue.assert_awaited_once() + token, event = namespace_app.event_processor.enqueue.await_args.args + assert token == "tok1" + assert event.name == "state.on_click" + assert event.router_data["headers"]["host"] == "localhost" + assert event.router_data["ip"] == "127.0.0.1" + + +@pytest.mark.asyncio +async def test_ping_pong(namespace: WebsocketEventNamespace): + """An application-level ping event gets a pong reply.""" + websocket = await serve(namespace, ["ping"], [PONG_MESSAGE]) + + assert ["ping", "pong"] in websocket.sent + + +@pytest.mark.asyncio +async def test_client_error_reaches_exception_handler( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """A client_error frame is routed to the frontend exception handler.""" + errors: list[str] = [] + namespace_app.frontend_exception_handler = lambda exc: errors.append(str(exc)) + await serve( + namespace, ["client_error", {"error_type": "boom", "message": "it broke"}] + ) + + assert len(errors) == 1 + assert "it broke" in errors[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", ["not json", '{"an": "object"}', "[42]"]) +async def test_malformed_frame_closes_connection( + namespace: WebsocketEventNamespace, frame: str +): + """A malformed frame closes the connection with 1002 (protocol error).""" + websocket = await serve(namespace, frame, ["ping"]) + + assert websocket.close_code == 1002 + # Nothing after the malformed frame is processed. + assert ["ping", "pong"] not in websocket.sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [ + None, + "not an event", + # A JSON-encoded event is still a string, not an event. + json.dumps({"name": "state.on_click", "payload": {}, "router_data": {}}), + 42, + {"name": 123, "payload": {}, "router_data": {}}, + {"name": "x", "payload": "nope", "router_data": {}}, + {"name": "x", "payload": {}, "router_data": []}, + {"name": "x", "payload": {}, "router_data": "nope"}, + {"name": "x", "payload": {}, "router_data": {"query": "not-a-dict"}}, + ], +) +async def test_undeserializable_event_closes_connection( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + payload: object, + caplog: pytest.LogCaptureFixture, +): + """An event frame that fails deserialization closes with 1002 and no warning.""" + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + websocket = await serve(namespace, ["event", payload], ["ping"]) + + assert websocket.close_code == 1002 + assert ["ping", "pong"] not in websocket.sent + namespace_app.event_processor.enqueue.assert_not_awaited() + # Client-controlled input must not write above debug level. + assert all(record.levelno <= logging.DEBUG for record in caplog.records) + + +@pytest.mark.asyncio +async def test_handler_error_keeps_connection( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + caplog: pytest.LogCaptureFixture, +): + """A server-side handler failure is logged and the connection survives.""" + namespace_app.event_processor.enqueue.side_effect = RuntimeError("server bug") + with caplog.at_level(logging.ERROR, logger="reflex.event_namespace"): + websocket = await serve( + namespace, + ["event", {"name": "state.on_click", "payload": {}, "router_data": {}}], + ["ping"], + ) + + assert websocket.close_code is None + assert ["ping", "pong"] in websocket.sent + assert any( + record.levelno == logging.ERROR + and "Error handling socket event" in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_tokenless_connection_rejected( + namespace: WebsocketEventNamespace, caplog: pytest.LogCaptureFixture +): + """A connection without a token closes with 1008 and logs nothing above debug. + + The version mismatch is not reported either: it is client-controlled and + the session is rejected anyway, so warning would only let anonymous + connects flood the logs. + """ + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + websocket = await serve( + namespace, + ["ping"], + websocket=FakeWebSocket(query_string=b"", subprotocols=["0.0.1"]), + ) + + assert websocket.close_code == 1008 + # A refused session is never acknowledged. + assert websocket.sent == [] + assert all(record.levelno <= logging.DEBUG for record in caplog.records) + + +@pytest.mark.asyncio +async def test_version_mismatch_warns_for_linked_session( + namespace: WebsocketEventNamespace, caplog: pytest.LogCaptureFixture +): + """A linked session with a stale frontend gets one sanitized warning.""" + with caplog.at_level(logging.WARNING, logger="reflex.event_namespace"): + await serve(namespace, websocket=FakeWebSocket(subprotocols=["0.0.1\x1b[31m"])) + + warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert len(warnings) == 1 + assert "0.0.1" in warnings[0] + assert "does not match the backend version" in warnings[0] + assert "\x1b" not in warnings[0] + + +@pytest.mark.asyncio +async def test_oversize_message_closes_connection( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """A frame over the size limit closes the connection with 1009.""" + monkeypatch.setenv("REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE", "20") + websocket = await serve(namespace, ["event", {"payload": "x" * 100}]) + + assert websocket.close_code == 1009 + + +@pytest.mark.asyncio +async def test_oversize_multibyte_message_closes_connection( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """The size limit counts bytes, so multibyte text cannot sneak past it.""" + monkeypatch.setenv("REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE", "25") + # 15 characters (under the limit) but 29 UTF-8 bytes (over it). + frame = '["x","€€€€€€€"]' + assert len(frame) <= 25 < len(frame.encode("utf-8")) + websocket = await serve(namespace, frame) + + assert websocket.close_code == 1009 + + +@pytest.mark.asyncio +async def test_multibyte_message_within_limit_is_processed( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """Multibyte frames within the byte limit pass through the exact check.""" + # 12 characters, 14 bytes: over limit/4 (triggers the exact byte count) + # but within the limit itself. + monkeypatch.setenv("REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE", "14") + websocket = await serve(namespace, '["ping","€"]') + + assert websocket.close_code is None + assert ["ping", "pong"] in websocket.sent + + +@pytest.mark.asyncio +async def test_binary_frame_closes_connection(namespace: WebsocketEventNamespace): + """A binary frame closes the connection with 1003 (unsupported data).""" + websocket = await serve(namespace, b"\x00\x01") + + assert websocket.close_code == 1003 + + +@pytest.mark.asyncio +async def test_disallowed_origin_is_rejected( + namespace: WebsocketEventNamespace, mocker +): + """A cross-origin connection is closed before being accepted.""" + from reflex_base.config import get_config + + mocker.patch.object( + get_config(), "cors_allowed_origins", ("https://allowed.example",) + ) + websocket = FakeWebSocket(origin="https://evil.example") + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + + assert not websocket.accepted + assert websocket.close_code == 1008 + + +@pytest.mark.asyncio +async def test_allowed_origin_is_accepted(namespace: WebsocketEventNamespace, mocker): + """A connection from an allowed origin is accepted.""" + from reflex_base.config import get_config + + mocker.patch.object( + get_config(), "cors_allowed_origins", ("https://allowed.example",) + ) + websocket = await serve( + namespace, websocket=FakeWebSocket(origin="https://allowed.example") + ) + + assert websocket.accepted + + +@pytest.mark.asyncio +async def test_duplicate_token_gets_new_token(namespace: WebsocketEventNamespace): + """A second tab connecting with the same token receives a new_token frame.""" + first = FakeWebSocket() + second = FakeWebSocket() + namespace._connections["sid1"] = _Connection(first) # pyright: ignore[reportArgumentType] + namespace._connections["sid2"] = _Connection(second) # pyright: ignore[reportArgumentType] + await namespace.link_token_to_sid("sid1", "tok1") + await namespace.link_token_to_sid("sid2", "tok1") + await _drain_tasks() + for connection in namespace._connections.values(): + connection.stop() + + new_token_frames = [frame for frame in second.sent if frame[0] == "new_token"] + assert len(new_token_frames) == 1 + assert new_token_frames[0][1] != "tok1" + + +@pytest.mark.asyncio +async def test_emit_to_unknown_sid_does_not_raise( + namespace: WebsocketEventNamespace, + caplog: pytest.LogCaptureFixture, +): + """Emitting to a session that went away is a silent no-op. + + A client disconnecting mid-event is routine, so nothing above DEBUG may be + logged. + """ + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + await namespace.emit("event", {"delta": {}}, to="gone") + assert all(record.levelno <= logging.DEBUG for record in caplog.records) + + +@pytest.mark.asyncio +async def test_websocket_records_connections_and_message_sizes( + namespace: WebsocketEventNamespace, otel_metrics +): + """A session adjusts the connection gauge and sizes frames in both directions.""" + await serve(namespace, ["ping"]) + + (connections,) = metric_points(otel_metrics, otel.METRIC_WEBSOCKET_CONNECTIONS) + # Connect and disconnect both happened, so the gauge is back at zero. + assert connections.value == 0 + sizes = { + p.attributes[otel.ATTR_NETWORK_IO_DIRECTION]: p.sum + for p in metric_points(otel_metrics, otel.METRIC_WEBSOCKET_MESSAGE_SIZE) + } + assert sizes == { + "receive": len(json.dumps(CONNECT_FRAME)) + len(json.dumps(["ping"])), + "transmit": len(json.dumps(["ping", "pong"], separators=(",", ":"))), + } + + +@pytest.mark.asyncio +async def test_websocket_event_uses_frontend_traceparent( + namespace: WebsocketEventNamespace, namespace_app: Mock, otel_exporter +): + """A traceparent in the event payload becomes the parent of the event span.""" + seen: list = [] + + async def enqueue(token, event): + await asyncio.sleep(0) + seen.append(trace.get_current_span().get_span_context()) + + namespace_app.event_processor.enqueue = enqueue + traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + with active_tracer().start_as_current_span("websocket"): + await serve( + namespace, + ["event", {"name": "state.h", "traceparent": traceparent}], + ["event", {"name": "state.h"}], + ) + + remote, fresh = seen + assert f"{remote.trace_id:032x}" == "0af7651916cd43dd8448eb211c80319c" + assert not fresh.is_valid + + +def test_default_transport_uses_websocket_namespace(): + """The default transport sets up the plain websocket namespace. + + The path is served with and without a trailing slash: proxy rules written + for Socket.IO route "/_event/*", and the client dials that form. + """ + app = App(enable_state=True) + assert isinstance(app.event_namespace, WebsocketEventNamespace) + assert app.sio is None + assert app._api is not None + websocket_routes = [ + route for route in app._api.router.routes if isinstance(route, WebSocketRoute) + ] + assert [route.path for route in websocket_routes] == ["/_event", "/_event/"] + + +@pytest.mark.asyncio +async def test_a_socketio_client_is_refused_before_accept( + namespace: WebsocketEventNamespace, +): + """A tab or bundle from before the plain transport cannot speak it. + + Accepted, it would wait for an engine.io handshake that never comes until + the connect timeout; refused, it fails at once. + """ + websocket = FakeWebSocket(query_string=b"EIO=4&transport=websocket&token=tok1") + websocket.feed() + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + + assert not websocket.accepted + assert websocket.close_code == 1008 + assert namespace.token_to_sid == {} + + +def test_socketio_transport_uses_socketio_namespace( + monkeypatch: pytest.MonkeyPatch, forked_registration_context: RegistrationContext +): + """transport="socketio" sets up the Socket.IO server and namespace.""" + from reflex.socketio_namespace import EventNamespace + + monkeypatch.setenv("REFLEX_TRANSPORT", "socketio") + app = App(enable_state=True) + assert isinstance(app.event_namespace, EventNamespace) + assert app.sio is not None + # Plain websocket transport under the hood. + assert app.sio.eio.transports == ["websocket"] + + +def test_polling_transport_uses_socketio_namespace( + monkeypatch: pytest.MonkeyPatch, forked_registration_context: RegistrationContext +): + """transport="polling" sets up the Socket.IO server with polling only.""" + from reflex.socketio_namespace import EventNamespace + + monkeypatch.setenv("REFLEX_TRANSPORT", "polling") + app = App(enable_state=True) + assert isinstance(app.event_namespace, EventNamespace) + assert app.sio is not None + assert app.sio.eio.transports == ["polling"] + + +def test_transport_chosen_in_a_test_does_not_outlive_it(): + """Runs after the transport tests above: the shared config is untouched. + + App() reloads the config into the active registration context, so the + tests that choose a transport through the environment build their apps + in a forked one. + """ + assert get_config().transport == "websocket" + + +def test_custom_sio_falls_back_to_the_socketio_transport( + monkeypatch: pytest.MonkeyPatch, + forked_registration_context: RegistrationContext, + mocker, +): + """A custom sio server under the default transport keeps working, deprecated. + + The default transport used to be Socket.IO over a websocket, so an app + passing its own server never had to choose one. + """ + from socketio import AsyncServer + + from reflex.socketio_namespace import EventNamespace + + deprecate = mocker.patch("reflex.app.console.deprecate") + # Recorded twice, so the variable the fallback persists is removed again. + monkeypatch.setenv("REFLEX_TRANSPORT", "websocket") + monkeypatch.delenv("REFLEX_TRANSPORT") + sio = AsyncServer(async_mode="asgi") + + app = App(sio=sio) + + assert app.sio is sio + assert isinstance(app.event_namespace, EventNamespace) + # Persisted, so the compiled frontend and the workers use it too. + assert get_config().transport == "socketio" + assert os.environ["REFLEX_TRANSPORT"] == "socketio" + deprecate.assert_called_once() + + +def test_custom_sio_with_an_explicit_websocket_transport_is_refused( + monkeypatch: pytest.MonkeyPatch, forked_registration_context: RegistrationContext +): + """Asking for the plain transport by name conflicts with a custom sio server.""" + from socketio import AsyncServer + + monkeypatch.setenv("REFLEX_TRANSPORT", "websocket") + with pytest.raises(RuntimeError, match=r"requires the Socket\.IO transport"): + App(sio=AsyncServer(async_mode="asgi")) + + +def test_custom_sio_with_socketio_transport( + monkeypatch: pytest.MonkeyPatch, forked_registration_context: RegistrationContext +): + """A custom sio server works with the Socket.IO transport.""" + from socketio import AsyncServer + + monkeypatch.setenv("REFLEX_TRANSPORT", "socketio") + sio = AsyncServer(async_mode="asgi") + app = App(sio=sio) + assert app.sio is sio + + +def test_app_socketio_reexports(): + """reflex.app still names the Socket.IO classes it used to import.""" + import socketio + + import reflex.app + from reflex.socketio_namespace import EventNamespace + + assert reflex.app.EventNamespace is EventNamespace + assert reflex.app.AsyncServer is socketio.AsyncServer + assert reflex.app.AsyncNamespace is socketio.AsyncNamespace + assert reflex.app.EngineIOApp is socketio.ASGIApp + with pytest.raises(AttributeError): + _ = reflex.app.DoesNotExist + + +def test_app_socketio_reexports_without_the_extra(monkeypatch: pytest.MonkeyPatch): + """Without python-socketio the names are absent, not import errors. + + hasattr() and getattr() with a default only swallow AttributeError. + """ + import reflex.app + + monkeypatch.setitem(sys.modules, "socketio", None) + monkeypatch.delitem(sys.modules, "reflex.socketio_namespace", raising=False) + + for name in ("EventNamespace", "AsyncServer"): + assert not hasattr(reflex.app, name) + with pytest.raises(AttributeError, match=r"reflex\[socketio\]"): + getattr(reflex.app, name) + + +@pytest.mark.parametrize( + ("transport", "name"), + [("socketio", "EventNamespace"), ("websocket", "WebsocketEventNamespace")], +) +def test_a_replaced_event_namespace_class_is_used( + monkeypatch: pytest.MonkeyPatch, + forked_registration_context: RegistrationContext, + transport: str, + name: str, +): + """An app can substitute its own namespace class for either transport.""" + import reflex.app + + class AuthorizingNamespace(getattr(reflex.app, name)): + """A namespace an app customized.""" + + monkeypatch.setattr(reflex.app, name, AuthorizingNamespace, raising=False) + monkeypatch.setenv("REFLEX_TRANSPORT", transport) + + assert isinstance(App(enable_state=True).event_namespace, AuthorizingNamespace) + + +@pytest.mark.parametrize( + ("override", "deprecated"), + [(None, False), ("handle_event", False), ("on_event", True)], +) +def test_overriding_on_event_is_deprecated( + monkeypatch: pytest.MonkeyPatch, + forked_registration_context: RegistrationContext, + mocker, + override: str | None, + deprecated: bool, +): + """An app overriding on_event is pointed to handle_event, which both transports run.""" + import reflex.app + from reflex.socketio_namespace import EventNamespace + + class TenantNamespace(EventNamespace): + async def handle_event(self, sid, data, asgi_scope): + await super().handle_event(sid, data, asgi_scope) + + class AuditingNamespace(EventNamespace): + async def on_event(self, sid, data): + await super().on_event(sid, data) + + namespace_classes = { + None: EventNamespace, + "handle_event": TenantNamespace, + "on_event": AuditingNamespace, + } + deprecate = mocker.patch("reflex.app.console.deprecate") + monkeypatch.setattr( + reflex.app, "EventNamespace", namespace_classes[override], raising=False + ) + monkeypatch.setenv("REFLEX_TRANSPORT", "socketio") + + App(enable_state=True) + + assert [call.kwargs["feature_name"] for call in deprecate.call_args_list] == ( + ["Overriding EventNamespace.on_event"] if deprecated else [] + ) + + +@pytest.mark.asyncio +async def test_a_handle_event_override_sees_every_event( + namespace_app: Mock, boot_event: dict[str, Any] +): + """A handle_event override sees the hydrate on connect and every later event.""" + seen: list[str] = [] + + class TenantNamespace(WebsocketEventNamespace): + async def handle_event(self, sid, data, asgi_scope): + seen.append(data["name"]) + await super().handle_event(sid, data, asgi_scope) + + click = {"name": "state.on_click", "payload": {}, "router_data": {}} + await serve( + TenantNamespace("/_event", namespace_app), + [CONNECT_MESSAGE, {"event": boot_event}], + ["event", click], + connect=False, + ) + + assert seen == ["state.hydrate_and_load", "state.on_click"] + + +@pytest.mark.asyncio +async def test_emit_without_a_recipient_reaches_every_session( + namespace: WebsocketEventNamespace, +): + """Like Socket.IO's, an emit addressed to nobody in particular broadcasts.""" + first, second = FakeWebSocket(), FakeWebSocket() + namespace._connections["sid1"] = _Connection(first) # pyright: ignore[reportArgumentType] + namespace._connections["sid2"] = _Connection(second) # pyright: ignore[reportArgumentType] + + await namespace.emit("custom", {"n": 1}) + await _drain_tasks() + for connection in namespace._connections.values(): + connection.stop() + + assert first.sent == second.sent == [["custom", {"n": 1}]] + + +@pytest.mark.asyncio +async def test_a_lone_surrogate_is_sent_as_an_escape( + namespace: WebsocketEventNamespace, +): + """A frame stays valid UTF-8 when state holds an unpaired surrogate. + + A non-UTF-8 file name decodes to one; written raw, the frame either fails + to send or arrives empty, and the update is silently lost. + """ + websocket = FakeWebSocket() + namespace._connections["sid1"] = _Connection(websocket) # pyright: ignore[reportArgumentType] + + await namespace.emit("event", {"file": "bad\udcffname"}, to="sid1") + await _drain_tasks() + namespace._connections["sid1"].stop() + + (text,) = websocket.raw_sent + text.encode() + assert json.loads(text) == ["event", {"file": "bad\udcffname"}] + + +def test_protocol_message_names_match_the_client(): + """The client speaks the same protocol message names as the server. + + Both ends declare these independently, and a rename on one side is + invisible until a browser fails to connect: the client would never + answer a heartbeat, so every session would be dropped on ping timeout. + """ + declarations = dict( + re.findall( + r'^const (\w+_MESSAGE) = "([^"]+)";$', + WEBSOCKET_JS_TEMPLATE.read_text(), + re.MULTILINE, + ) + ) + + assert declarations == { + "CONNECT_MESSAGE": CONNECT_MESSAGE, + "HANDSHAKE_MESSAGE": HANDSHAKE_MESSAGE, + "PING_MESSAGE": PING_MESSAGE, + "PONG_MESSAGE": PONG_MESSAGE, + "OPEN_MESSAGE": OPEN_MESSAGE, + "OPENED_MESSAGE": OPENED_MESSAGE, + "CHANNEL_ERROR_MESSAGE": CHANNEL_ERROR_MESSAGE, + } + # CLOSE_MESSAGE is missing on purpose: a browser handle is shared by every + # component that asks for the channel, so it never closes one by itself and + # the backend frees the session on disconnect. + assert "CLOSE_MESSAGE" not in declarations + + +class RecordingChannel(Channel): + """A channel that records what the transport hands it.""" + + name = "probe" + + def __init__(self, accepts_binary: bool = False): + """Initialize the channel, recording opens, messages and closes.""" + type(self).accepts_binary = accepts_binary + super().__init__() + self.opened: list[ChannelSession] = [] + self.closed: list[ChannelSession] = [] + self.messages: list[tuple[str, Any, list[bytes]]] = [] + + async def on_open(self, session: ChannelSession) -> None: + """Record the opened session.""" + self.opened.append(session) + + async def on_message( + self, session: ChannelSession, event: str, data: Any, buffers: list[bytes] + ) -> None: + """Record the message and answer an echo.""" + self.messages.append((event, data, buffers)) + await session.send("echo", data, buffers) + + async def on_close(self, session: ChannelSession) -> None: + """Record the closed session.""" + self.closed.append(session) + + +@pytest.fixture +def probe_channel(namespace_app: Mock) -> RecordingChannel: + """A recording channel the mock app serves as "probe". + + Returns: + The channel. + """ + channel = RecordingChannel() + namespace_app._channels = {"probe": channel} + return channel + + +def channel_frames(websocket: FakeWebSocket, channel: str = "probe") -> list[Any]: + """Text frames the server sent on a channel. + + Args: + websocket: The fake websocket. + channel: The channel name. + + Returns: + The matching frames. + """ + return [ + frame + for frame in websocket.sent + if isinstance(frame, list) and len(frame) > 2 and frame[2] == channel + ] + + +def test_channel_frame_round_trip(): + """A binary frame decodes to the values it was built from.""" + buffers = [b"\x01\x02\x03", b"", bytes(range(16))] + frame = encode_channel_frame("payload", {"fig": "f1"}, "probe", buffers) + assert decode_channel_frame(frame) == ("payload", {"fig": "f1"}, "probe", buffers) + + +def test_channel_frame_aligns_attachments(): + """Every attachment starts on an 8-byte boundary, whatever the header size.""" + for name_length in range(1, 24): + frame = encode_channel_frame( + "e", {"pad": "x" * name_length}, "probe", [b"\x01" * 3, b"\x02" * 5] + ) + offsets = [] + offset = 4 + int.from_bytes(frame[:4], "little") + for length in (3, 5): + offset += -offset % 8 + offsets.append(offset) + offset += length + assert all(candidate % 8 == 0 for candidate in offsets) + assert frame[offsets[0] : offsets[0] + 3] == b"\x01" * 3 + assert frame[offsets[1] : offsets[1] + 5] == b"\x02" * 5 + + +@pytest.mark.parametrize( + "frame", + [ + b"", + b"\x02\x00", + (255).to_bytes(4, "little") + b'["e",null,"c",[]]', + (17).to_bytes(4, "little") + b'["e",null,"c",42]', + (18).to_bytes(4, "little") + b'["e",null,"c",[1]]', + (21).to_bytes(4, "little") + b'["e",null,"c",[-1]]xx', + (17).to_bytes(4, "little") + b'{"not": "a list"}', + (11).to_bytes(4, "little") + b"not json at all", + ], +) +def test_malformed_channel_frame_is_rejected(frame: bytes): + """A frame that does not follow the binary format raises.""" + with pytest.raises(ValueError): + decode_channel_frame(frame) + + +def test_channel_frame_rejects_too_many_attachments(): + """A frame declaring more attachments than the cap raises.""" + header = json.dumps(["e", None, "c", [0] * 65]).encode() + with pytest.raises(ValueError, match="more than"): + decode_channel_frame(len(header).to_bytes(4, "little") + header) + + +@pytest.mark.asyncio +async def test_channel_open_and_message( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + probe_channel: RecordingChannel, +): + """Opening a channel answers _opened and routes messages to the channel.""" + websocket = await serve( + namespace, [OPEN_MESSAGE, None, "probe"], ["sub", {"fig": "f1"}, "probe"] + ) + + assert channel_frames(websocket)[0] == [OPENED_MESSAGE, None, "probe"] + assert [session.client_token for session in probe_channel.opened] == ["tok1"] + assert probe_channel.messages == [("sub", {"fig": "f1"}, [])] + assert ["echo", {"fig": "f1"}, "probe"] in websocket.sent + # The disconnect closed the session again. + assert probe_channel.closed == probe_channel.opened + + +@pytest.mark.asyncio +async def test_channel_unknown_name_reports_error( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """Opening a channel the backend does not serve answers _error.""" + namespace_app._channels = {} + websocket = await serve(namespace, [OPEN_MESSAGE, None, "nope"]) + + error = channel_frames(websocket, "nope")[0] + assert error[0] == CHANNEL_ERROR_MESSAGE + assert error[1]["code"] == "unknown_channel" + + +@pytest.mark.asyncio +async def test_channel_message_before_open_reports_error( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + probe_channel: RecordingChannel, +): + """A message for an unopened channel answers _error and is not dispatched.""" + websocket = await serve(namespace, ["sub", {}, "probe"]) + + assert channel_frames(websocket)[0][1]["code"] == "channel_not_open" + assert probe_channel.messages == [] + + +@pytest.mark.asyncio +async def test_channel_close_ends_the_session( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + probe_channel: RecordingChannel, +): + """A _close frame runs on_close and stops dispatching to the channel.""" + websocket = await serve( + namespace, + [OPEN_MESSAGE, None, "probe"], + [CLOSE_MESSAGE, None, "probe"], + ["sub", {}, "probe"], + ) + + assert len(probe_channel.closed) == 1 + assert probe_channel.messages == [] + assert channel_frames(websocket)[-1][1]["code"] == "channel_not_open" + + +@pytest.mark.asyncio +async def test_channel_binary_message_round_trip( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """A binary frame reaches a channel that accepts binary, and echoes back.""" + channel = RecordingChannel(accepts_binary=True) + namespace_app._channels = {"probe": channel} + websocket = await serve( + namespace, + [OPEN_MESSAGE, None, "probe"], + encode_channel_frame("push", {"n": 2}, "probe", [b"\x00\x01", b"\x02"]), + ) + + assert channel.messages == [("push", {"n": 2}, [b"\x00\x01", b"\x02"])] + echoed = [frame for frame in websocket.sent if isinstance(frame, bytes)] + assert decode_channel_frame(echoed[0]) == ( + "echo", + {"n": 2}, + "probe", + [b"\x00\x01", b"\x02"], + ) + + +@pytest.mark.asyncio +async def test_channel_binary_rejected_when_not_accepted( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + probe_channel: RecordingChannel, +): + """Binary attachments to a text-only channel answer _error.""" + websocket = await serve( + namespace, + [OPEN_MESSAGE, None, "probe"], + encode_channel_frame("push", None, "probe", [b"\x00"]), + ) + + assert channel_frames(websocket)[-1][1]["code"] == "binary_not_accepted" + assert probe_channel.messages == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("binary", [False, True]) +@pytest.mark.parametrize("event", [*sorted(RESERVED_EVENTS), "_opened", "_error"]) +async def test_channel_reserved_inbound_name_is_rejected( + namespace: WebsocketEventNamespace, namespace_app: Mock, event: str, binary: bool +): + """A client may not send a message under a reserved name, on either frame kind. + + A handler that relays what it receives would answer such a message with a + send the channel API refuses, so the name never reaches it. + """ + channel = RecordingChannel(accepts_binary=True) + namespace_app._channels = {"probe": channel} + websocket = await serve( + namespace, + [OPEN_MESSAGE, None, "probe"], + encode_channel_frame(event, {"fig": "f1"}, "probe", [b"\x00"]) + if binary + else [event, {"fig": "f1"}, "probe"], + ) + + assert channel.messages == [] + assert channel_frames(websocket)[-1][1]["code"] == "reserved_event" + assert websocket.close_code is None + + +@pytest.mark.asyncio +async def test_malformed_binary_frame_closes_connection( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + probe_channel: RecordingChannel, +): + """A binary frame that is not a channel frame closes the connection.""" + websocket = await serve(namespace, b"\xff\xff\xff\xff not a frame") + + assert websocket.close_code == 1002 + + +@pytest.mark.asyncio +async def test_oversize_binary_frame_closes_connection( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + monkeypatch: pytest.MonkeyPatch, + probe_channel: RecordingChannel, +): + """A binary frame over the size limit closes the connection with 1009.""" + monkeypatch.setenv("REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE", "16") + websocket = await serve( + namespace, encode_channel_frame("push", None, "probe", [b"\x00" * 64]) + ) + + assert websocket.close_code == 1009 + + +@pytest.mark.asyncio +async def test_channel_handler_error_keeps_connection( + namespace: WebsocketEventNamespace, namespace_app: Mock, caplog +): + """A channel handler raising is logged without dropping the connection.""" + + class BoomChannel(Channel): + name = "boom" + + async def on_message(self, session, event, data, buffers) -> None: + msg = "handler failed" + raise RuntimeError(msg) + + namespace_app._channels = {"boom": BoomChannel()} + websocket = FakeWebSocket() + with caplog.at_level(logging.ERROR): + websocket.feed([OPEN_MESSAGE, None, "boom"], ["go", None, "boom"], ["ping"]) + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert websocket.close_code is None + assert ["ping", "pong"] in websocket.sent + assert "handler failed" in caplog.text + + +@pytest.mark.asyncio +async def test_channel_frame_with_non_string_name_closes_connection( + namespace: WebsocketEventNamespace, +): + """A frame whose channel element is not a string closes the connection.""" + websocket = await serve(namespace, ["sub", None, 42]) + + assert websocket.close_code == 1002 + + +NODE = shutil.which("node") or "" + +# The client template is copied next to each script, so it is imported by a +# relative specifier. Its "$/" specifiers are a vite alias node cannot resolve, +# so they become file URLs: an absolute Windows path is neither a valid ESM +# specifier nor a valid JS string literal (its separators are escapes). +CLIENT_MODULE = json.dumps("./websocket.mjs") +ALIAS_SPECIFIER_RE = re.compile(r'"\$/([^"]+)"') + + +def _resolve_alias(match: re.Match[str]) -> str: + """Rewrite one "$/" specifier to a file URL, adding the extension vite infers. + + Args: + match: The matched specifier, with the path after "$/" as group 1. + + Returns: + The file URL as a JS string literal. + """ + target = WEB_TEMPLATE_ROOT / match[1] + if not target.suffix: + target = target.with_suffix(".js") + return json.dumps(target.as_uri()) + + +def _run_client_script(tmp_path: Path, source: str, *args: str) -> Any: + """Run a script against the client template, returning the JSON it printed. + + Args: + tmp_path: Where to write the script. + source: The script source, importing from CLIENT_MODULE. + args: Command line arguments for the script. + + Returns: + The parsed JSON the script wrote to stdout. + """ + (tmp_path / "websocket.mjs").write_text( + ALIAS_SPECIFIER_RE.sub(_resolve_alias, WEBSOCKET_JS_TEMPLATE.read_text()) + ) + script = tmp_path / "client.mjs" + script.write_text(source) + result = subprocess.run( + [NODE, str(script), *args], capture_output=True, text=True, check=False + ) + + assert result.returncode == 0, result.stderr + return json.loads(result.stdout) + + +def test_header_values_that_are_not_utf8_are_kept(): + """A header may carry any octets; one that is not UTF-8 must not fail the connection. + + The connection-scoped router data is built from the headers on the first + event, so a decode error dropped every event the connection sent. UTF-8 + values still decode as before. + """ + router_data = event_namespace.build_static_router_data( + "sid1", + {"headers": [(b"x-raw", b"\xff\xfe"), (b"x-name", "José".encode())]}, + ) + + headers = router_data["headers"] + assert headers["x-raw"] == b"\xff\xfe".decode("latin-1") + assert headers["x-name"] == "José" + + +def test_binary_frame_sizes_attachments_in_bytes(): + """An attachment is measured in bytes, whatever its item size. + + ``len()`` counts items, so an array of doubles would declare an eighth of + its size and every later attachment would be read from the wrong offset. + """ + doubles = array.array("d", [1.0, 2.0]) + frame = encode_channel_frame("rows", None, "probe", [doubles, b"SECOND!!"]) + + assert decode_channel_frame(frame)[3] == [doubles.tobytes(), b"SECOND!!"] + + +@pytest.mark.skipif(not NODE, reason="Requires node to run the client codec") +def test_binary_frame_codec_matches_the_client(tmp_path: Path): + """Both ends agree on the binary frame layout, both ways. + + The layout (header length, padding, alignment) is implemented twice, and a + disagreement would only surface as an unreadable payload in a browser. + """ + buffers = [b"\x01\x02\x03", b"", bytes(range(24))] + frame = encode_channel_frame("payload", {"fig": "f1", "n": 3}, "probe", buffers) + decoded = _run_client_script( + tmp_path, + f""" +import {{ encodeChannelFrame, decodeChannelFrame }} from {CLIENT_MODULE}; + +const fromPython = Uint8Array.from(Buffer.from(process.argv[2], "base64")); +const [event, data, channel, buffers] = decodeChannelFrame(fromPython.buffer); +const encoded = encodeChannelFrame(event, data, channel, buffers); +console.log(JSON.stringify({{ + event, + data, + channel, + lengths: buffers.map((b) => b.byteLength), + // Every attachment must be readable as a typed array in place. + aligned: buffers.every((b) => b.byteOffset % 8 === 0), + encoded: Buffer.from(encoded).toString("base64"), +}})); +""", + base64.b64encode(frame).decode(), + ) + + assert decoded["event"] == "payload" + assert decoded["data"] == {"fig": "f1", "n": 3} + assert decoded["channel"] == "probe" + assert decoded["lengths"] == [len(buffer) for buffer in buffers] + assert decoded["aligned"] + # The frame the client built decodes back to the same message. + assert decode_channel_frame(base64.b64decode(decoded["encoded"])) == ( + "payload", + {"fig": "f1", "n": 3}, + "probe", + buffers, + ) + + +@pytest.mark.asyncio +async def test_repeated_channel_open_reuses_the_session( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """Opening an already-open channel answers again without a second session. + + A second session would keep the first one's room membership alive with + nothing left to close it. + """ + + class RoomChannel(RecordingChannel): + name = "probe" + + async def on_open(self, session: ChannelSession) -> None: + """Join a room so an orphaned session would be observable.""" + await super().on_open(session) + session.join("all") + + channel = RoomChannel() + namespace_app._channels = {"probe": channel} + websocket = await serve( + namespace, [OPEN_MESSAGE, None, "probe"], [OPEN_MESSAGE, None, "probe"] + ) + + assert len(channel.opened) == 1 + assert [frame[0] for frame in channel_frames(websocket)] == [ + OPENED_MESSAGE, + OPENED_MESSAGE, + ] + # The disconnect left nothing behind. + assert channel._rooms == {} + assert channel._sessions == {} + + +@pytest.mark.asyncio +async def test_interrupted_teardown_still_unlinks_the_token( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """Token cleanup starts before any teardown await can be interrupted. + + Server shutdown cancels the connection task, and an await in the teardown + path raises again inside a cancelled scope. A token left linked would make + the client's reconnect look like a duplicate tab. + """ + + class StallingChannel(RecordingChannel): + name = "probe" + + async def on_close(self, session: ChannelSession) -> None: + """Fail the way a cancelled cleanup await would.""" + raise asyncio.CancelledError + + namespace_app._channels = {"probe": StallingChannel()} + websocket = FakeWebSocket() + websocket.feed([OPEN_MESSAGE, None, "probe"]) + with pytest.raises(asyncio.CancelledError): + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert "tok1" not in namespace.token_to_sid + + +@pytest.mark.asyncio +async def test_interrupted_open_leaves_no_session_behind( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """A session interrupted inside on_open is still cleaned up on disconnect. + + on_open may already have joined rooms, so a session the transport never + recorded would stay reachable by fan-out with nothing left to close it. + """ + + class InterruptedChannel(RecordingChannel): + name = "probe" + + async def on_open(self, session: ChannelSession) -> None: + """Join a room, then fail the way a cancellation would.""" + session.join("all") + raise asyncio.CancelledError + + channel = InterruptedChannel() + namespace_app._channels = {"probe": channel} + websocket = FakeWebSocket() + websocket.feed([OPEN_MESSAGE, None, "probe"]) + with pytest.raises(asyncio.CancelledError): + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert channel._rooms == {} + assert channel._sessions == {} + + +@pytest.mark.asyncio +async def test_disconnect_marks_the_session_closed( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + probe_channel: RecordingChannel, +): + """The session a channel holds reports the disconnect to its handlers.""" + await serve(namespace, [OPEN_MESSAGE, None, "probe"]) + + assert [session.open for session in probe_channel.opened] == [False] + + +@pytest.mark.asyncio +async def test_deeply_nested_frame_closes_connection( + namespace: WebsocketEventNamespace, +): + """A frame the JSON decoder cannot recurse through closes the connection. + + Deep nesting exhausts the decoder's stack instead of failing to parse, and + a RecursionError escaping the receive loop would drop the session with a + traceback rather than a protocol close. + """ + depth = sys.getrecursionlimit() * 200 + websocket = await serve(namespace, "[" * depth + "]" * depth) + + assert websocket.close_code == 1002 + + +def test_decoding_a_too_deeply_nested_header_is_a_value_error(mocker): + """A header the decoder cannot recurse through fails like any bad frame. + + Callers handle malformed frames by catching ValueError; that the decoder + signals deep nesting with RecursionError is its own business, and the + header size cap makes how deep is too deep interpreter-specific, so the + failure is simulated rather than provoked. + """ + # Only this module's reference to json, so nothing else loses its decoder. + mocker.patch.object( + event_namespace, + "json", + SimpleNamespace(loads=Mock(side_effect=RecursionError("too deep"))), + ) + frame = encode_channel_frame("push", None, "probe", [b"\x00"]) + + with pytest.raises(ValueError, match="nested too deeply"): + decode_channel_frame(frame) + + +@pytest.mark.asyncio +async def test_handshake_advertises_the_message_limit( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """The client needs the inbound limit to refuse an oversized frame itself.""" + monkeypatch.setenv("REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE", "4096") + websocket = FakeWebSocket() + websocket.feed() + + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert websocket.sent[0][1]["max_message_size"] == 4096 + + +@pytest.mark.asyncio +async def test_a_room_broadcast_serializes_one_frame_for_every_member( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch +): + """Fan-out costs one encode and one buffer, not one per recipient.""" + channel = RecordingChannel(accepts_binary=True) + sent: list[tuple[str, Any]] = [] + encodes = 0 + real_encode = event_namespace.encode_channel_frame + + def counting_encode(*args: Any) -> bytes: + nonlocal encodes + encodes += 1 + return real_encode(*args) + + def record(sids: Iterable[str], payload: Any, label: str) -> bool: + # Stands in for queueing the frame on each client's connection. + sent.extend((sid, payload) for sid in sids) + return True + + monkeypatch.setattr(event_namespace, "encode_channel_frame", counting_encode) + monkeypatch.setattr(namespace, "_deliver", record) + for sid in ("sid1", "sid2", "sid3"): + namespace._token_manager.sid_to_token[sid] = "tok" + channel.open_session(sid, "tok", namespace._send_channel_message).join("fig:1") + + await channel.send_to_room("fig:1", "frame", {"seq": 1}, [b"\x00" * 32]) + + assert encodes == 1 + assert len(sent) == 3 + assert {to for to, _ in sent} == {"sid1", "sid2", "sid3"} + # The same buffer reached all three, rather than a copy each. + assert sent[0][1] is sent[1][1] is sent[2][1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send", + [ + pytest.param( + lambda channel, session: channel.send_to_room("all", "news"), id="room" + ), + pytest.param( + lambda channel, session: channel.send_to_token("tok1", "news"), id="token" + ), + pytest.param(lambda channel, session: session.send("news"), id="session"), + ], +) +async def test_a_session_whose_token_moved_is_closed_instead_of_sent_to( + namespace: WebsocketEventNamespace, + namespace_app: Mock, + send: Callable[[Channel, ChannelSession], Awaitable[Any]], +): + """Channel messages stop at a socket whose token another instance claimed. + + The claim drops this socket's token link but not the connection, so until + the socket sent another frame, its channel sessions kept receiving room + broadcasts, token sends and session sends. + """ + opened = asyncio.Event() + + class RoomChannel(RecordingChannel): + name = "probe" + + async def on_open(self, session: ChannelSession) -> None: + """Join a room and signal the open.""" + await super().on_open(session) + session.join("all") + opened.set() + + channel = RoomChannel() + namespace_app._channels = {"probe": channel} + websocket = FakeWebSocket() + websocket.feed([OPEN_MESSAGE, None, "probe"], disconnect=False) + serving = asyncio.create_task(namespace.handle_websocket(websocket)) # pyright: ignore[reportArgumentType] + await asyncio.wait_for(opened.wait(), 1) + # What RedisTokenManager does on adopting another instance's record. + sid = namespace.token_to_sid["tok1"] + namespace._token_manager.token_to_socket["tok1"] = SocketRecord( + instance_id="other-instance", sid="other-sid" + ) + del namespace._token_manager.sid_to_token[sid] + + await send(channel, channel.opened[0]) + await _drain_tasks() + + assert "news" not in [frame[0] for frame in channel_frames(websocket)] + await asyncio.wait_for(serving, 1) + assert websocket.close_code == 1008 + # The teardown forgot the stale session and its rooms. + assert channel.closed == channel.opened + assert channel._rooms == {} + + +@pytest.mark.asyncio +async def test_send_to_token_reports_a_socket_whose_token_moved_as_a_miss( + namespace: WebsocketEventNamespace, +): + """A token send reaching only a socket whose token moved reports no delivery. + + The move drops the socket's token link before its teardown forgets the + channel session, so the channel still finds the session; the caller has to + learn that nothing was sent to take another delivery path. + """ + channel = RecordingChannel() + live, moved = FakeWebSocket(), FakeWebSocket() + namespace._connections["live"] = _Connection(live) # pyright: ignore[reportArgumentType] + namespace._connections["moved"] = _Connection(moved) # pyright: ignore[reportArgumentType] + # The moved socket keeps its connection but has no token link. + namespace._token_manager.sid_to_token["live"] = "tok-live" + channel.open_session("live", "tok-live", namespace._send_channel_message) + channel.open_session("moved", "tok-moved", namespace._send_channel_message) + + assert await channel.send_to_token("tok-live", "news") is True + assert await channel.send_to_token("tok-moved", "news") is False + await _drain_tasks() + for connection in namespace._connections.values(): + connection.stop() + + assert channel_frames(live) == [["news", None, "probe"]] + assert channel_frames(moved) == [] + assert namespace._connections["moved"].close_code == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state", ["closing", "overflowing", "gone"]) +async def test_send_to_token_reports_a_frame_no_connection_took_as_a_miss( + namespace: WebsocketEventNamespace, monkeypatch: pytest.MonkeyPatch, state: str +): + """A token send its linked session's connection cannot take reports no delivery. + + The connection may already be closing, close over this very frame, or be + gone; the frame never reaches the client either way. + """ + channel = RecordingChannel() + # The backlog only overflows for a client that stopped reading. + websocket = StalledWebSocket() if state == "overflowing" else FakeWebSocket() + connection = _Connection(websocket) # pyright: ignore[reportArgumentType] + if state != "gone": + namespace._connections["sid1"] = connection + if state == "closing": + connection.close(1001) + elif state == "overflowing": + monkeypatch.setattr(event_namespace, "_MAX_SEND_BACKLOG", 1) + connection.send("[]") + namespace._token_manager.sid_to_token["sid1"] = "tok" + channel.open_session("sid1", "tok", namespace._send_channel_message) + + assert await channel.send_to_token("tok", "news") is False + await _drain_tasks() + connection.stop() + + assert channel_frames(websocket) == [] + + +@pytest.mark.asyncio +async def test_a_frame_that_overflows_the_backlog_never_reaches_the_client( + monkeypatch: pytest.MonkeyPatch, +): + """A frame whose send reports it unqueued stays off the wire and the backlog. + + The writer runs until teardown stops it, so a client that resumes reading + in between must not get the frame its sender was told missed it. + """ + + class PausedWebSocket(FakeWebSocket): + """A peer that stopped reading for now: writes wait until it resumes.""" + + def __init__(self): + """Initialize the fake websocket.""" + super().__init__() + self.reading = asyncio.Event() + + async def send_text(self, text: str): + """Record an outgoing frame once the peer reads again.""" + await self.reading.wait() + await super().send_text(text) + + monkeypatch.setattr(event_namespace, "_MAX_SEND_BACKLOG", 10) + websocket = PausedWebSocket() + connection = _Connection(websocket) # pyright: ignore[reportArgumentType] + assert connection.send('["a"]') is True + # The writer takes the first frame and waits on the paused peer. + await _drain_tasks() + assert connection.send('["b"]') is True + assert connection.send('["overflow"]') is False + assert connection.close_code == 1008 + assert connection._backlog == len('["b"]') + + websocket.reading.set() + await _drain_tasks() + connection.stop() + + assert websocket.sent == [["a"], ["b"]] + assert connection._backlog == 0 + + +@pytest.mark.asyncio +async def test_closing_a_connection_the_heartbeat_already_closed_is_not_an_error( + namespace: WebsocketEventNamespace, +): + """A close code decided over a frame may find the socket already gone. + + The heartbeat closes an unresponsive session from its own task, so the + receive loop's own close can land on a connection starlette has already + said goodbye on. Teardown still has to finish. + """ + websocket = FakeWebSocket() + # No channels, so a binary frame is a close-worthy protocol error. + websocket.feed(b"\x00") + websocket.refuse_close = True + + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert "tok1" not in namespace.token_to_sid + + +@pytest.mark.asyncio +async def test_handler_error_counters_do_not_outlive_their_sessions( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """Per-session error counters are dropped when the session goes away. + + A close hook runs after the disconnect bookkeeping, so a failing one logs + through the budget once the counter it uses has already been dropped -- + recreating it for a session that no longer exists. + """ + + class FailingChannel(Channel): + name = "probe" + + async def on_message( + self, session: ChannelSession, event: str, data: Any, buffers: list[bytes] + ) -> None: + """Accept anything; only the close hook matters here.""" + + async def on_close(self, session: ChannelSession) -> None: + """Fail the way a buggy close hook would. + + Args: + session: The session being closed. + + Raises: + RuntimeError: Always. + """ + message = "boom" + raise RuntimeError(message) + + namespace_app._channels = {"probe": FailingChannel()} + for index in range(3): + websocket = FakeWebSocket(query_string=f"token=tok{index}".encode()) + websocket.feed([OPEN_MESSAGE, None, "probe"]) + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert namespace._handler_error_budget.counts == {} + + +def test_client_limits_match_the_protocol(): + """The client enforces the same caps the backend closes the connection over. + + Both ends declare them independently; a client cap that drifted above the + backend's would turn a loud local error back into a dropped connection. + """ + template = WEBSOCKET_JS_TEMPLATE.read_text() + + assert f"const MAX_MESSAGE_BUFFERS = {MAX_MESSAGE_BUFFERS};" in template + assert f'const MAX_MESSAGE_SIZE_REASON = "{MAX_MESSAGE_SIZE_REASON}";' in template + reserved = re.search(r"const LIFECYCLE_EVENTS = new Set\(\[([^\]]+)\]\);", template) + assert reserved is not None + assert {name.strip().strip('"') for name in reserved.group(1).split(",")} == set( + RESERVED_EVENTS + ) + + +@pytest.mark.asyncio +async def test_large_metadata_is_bounded_only_by_the_message_limit( + namespace: WebsocketEventNamespace, namespace_app: Mock +): + """Metadata is limited by the frame size, not by a second hidden cap. + + The same metadata sent without attachments travels as a text frame, which + only the message limit applies to; a binary frame that rejected it would + close the connection over a payload the client had no way to know was too + big -- the handshake advertises one limit. Driven through the receive path, + because that is where such a frame would be turned into a close. + """ + channel = RecordingChannel(accepts_binary=True) + namespace_app._channels = {"probe": channel} + # Far past any header-shaped cap, far under the 1 MB message limit. + metadata = {"spec": "x" * (128 * 1024)} + websocket = FakeWebSocket() + websocket.feed( + [OPEN_MESSAGE, None, "probe"], + encode_channel_frame("payload", metadata, "probe", [b"\x00\x01"]), + ) + + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + assert websocket.close_code is None + assert channel.messages == [("payload", metadata, [b"\x00\x01"])] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + pytest.param(["event", {"name": "state.on_click", "payload": {}}], id="event"), + pytest.param(["push", {"x": 1}, "probe"], id="channel"), + pytest.param(None, id="channel-binary"), + ], +) +async def test_frames_from_a_session_whose_token_went_away_close_it( + namespace: WebsocketEventNamespace, namespace_app: Mock, caplog, frame: Any +): + """A session that loses its token mid-connection is closed, whatever it sends. + + The token manager drops the mapping when a token moves to another socket + or its record goes stale, while that socket stays open. Nothing it sends + can be served after that: an event has no token to enqueue under, and a + channel message would keep invoking handlers under a token that has moved + on. The payload stays out of the log as well -- at warning level, per + frame, it would be a log flood and an injection vector both. + """ + channel = RecordingChannel(accepts_binary=True) + namespace_app._channels = {"probe": channel} + if frame is None: + frame = encode_channel_frame( + "push", {"x": "\n[fake] log line"}, "probe", [b"!"] + ) + + class LosesItsToken(FakeWebSocket): + """Drops the token mapping once the channel is open and serving. + + Not on the frame after the connect frame: that one opens the channel, + and losing the token before it is dispatched would close the + connection over the open itself, leaving nothing for the frame under + test to prove. + """ + + delivered = 0 + + async def receive(self) -> dict[str, Any]: + message = await super().receive() + self.delivered += 1 + if self.delivered == 3: + namespace.sid_to_token.clear() + return message + + websocket = LosesItsToken() + websocket.feed([OPEN_MESSAGE, None, "probe"], frame) + + with caplog.at_level(logging.DEBUG): + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + # The channel was open and serving when its token went away. + assert len(channel.opened) == 1 + assert websocket.close_code == 1008 + namespace_app.event_processor.enqueue.assert_not_awaited() + assert channel.messages == [] + # Nothing the client sent reached the log, at any level. + assert "[fake] log line" not in caplog.text + assert [r for r in caplog.records if r.levelno >= logging.WARNING] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hook", ["on_open", "on_message", "on_close"]) +async def test_a_failing_channel_hook_cannot_flood_the_logs( + namespace: WebsocketEventNamespace, namespace_app: Mock, caplog, hook: str +): + """Tracebacks from any handler a client keeps breaking are budgeted. + + The connection survives a raising handler, which is what lets a client + send the frame again; without a budget each repeat writes another + traceback, whichever of the three hooks it lands in. + """ + + class BoomChannel(Channel): + name = "boom" + + async def on_message(self, session, event, data, buffers) -> None: + """Fail the way a handler meeting unexpected metadata would.""" + if hook == "on_message": + raise KeyError(data) + + async def on_open(self, session) -> None: + """Fail while setting the session up.""" + if hook == "on_open": + raise KeyError(session.sid) + + async def on_close(self, session) -> None: + """Fail while tearing the session down.""" + if hook == "on_close": + raise KeyError(session.sid) + + namespace_app._channels = {"boom": BoomChannel()} + attempts = namespace._MAX_CLIENT_ERRORS_PER_SID * 4 + # Reopening is what a client does after a failed open, and open/close is + # the only way to reach on_close repeatedly. + frames: list[Any] = [] + for _ in range(attempts): + frames.append([OPEN_MESSAGE, None, "boom"]) + if hook == "on_message": + frames.append(["go", {"n": 1}, "boom"]) + elif hook == "on_close": + frames.append([CLOSE_MESSAGE, None, "boom"]) + websocket = FakeWebSocket() + websocket.feed(*frames, ["ping"]) + + with caplog.at_level(logging.DEBUG): + await namespace.handle_websocket(websocket) # pyright: ignore[reportArgumentType] + await _drain_tasks() + + tracebacks = [r for r in caplog.records if r.levelno >= logging.ERROR] + assert len(tracebacks) == namespace._MAX_CLIENT_ERRORS_PER_SID + assert all(record.exc_info for record in tracebacks) + # The connection is still serving: a channel bug is not the client's fault. + assert websocket.close_code is None + assert ["ping", "pong"] in websocket.sent + + +def test_decode_asgi_headers(): + """_decode_asgi_headers decodes raw ASGI header pairs into a str dict.""" + assert _decode_asgi_headers([]) == {} + assert _decode_asgi_headers([ + (b"host", b"example.com"), + (b"x-forwarded-for", b"10.0.0.1, 10.0.0.2"), + ]) == { + "host": "example.com", + "x-forwarded-for": "10.0.0.1, 10.0.0.2", + } + # Repeated cookie fields are joined with "; " (RFC 9113 section 8.2.3). + assert _decode_asgi_headers([ + (b"cookie", b"a=1"), + (b"cookie", b"b=2"), + ]) == {"cookie": "a=1; b=2"} + # Any other repeated field is joined in order with ", " (RFC 9110 section 5.3). + assert _decode_asgi_headers([ + (b"x-forwarded-for", b"10.0.0.1"), + (b"accept", b"text/html"), + (b"x-forwarded-for", b"10.0.0.2, 10.0.0.3"), + ]) == { + "x-forwarded-for": "10.0.0.1, 10.0.0.2, 10.0.0.3", + "accept": "text/html", + } + # Names and values are decoded as UTF-8. + assert _decode_asgi_headers([(b"x-name", "café".encode())]) == {"x-name": "café"} + # Octets that are not UTF-8, in a name or a value, decode as Latin-1; the + # other headers still decode as UTF-8. + assert _decode_asgi_headers([ + (b"x-\xff", b"a"), + (b"x-raw", b"\xfe"), + ("x-café".encode(), "café".encode()), + ]) == {"x-ÿ": "a", "x-raw": "þ", "x-café": "café"} + # A repeated field is still joined when one of its values is not UTF-8. + assert _decode_asgi_headers([ + (b"cookie", b"a=1"), + (b"cookie", b"b=\xff"), + ]) == {"cookie": "a=1; b=ÿ"} + + +def test_static_router_data_keeps_every_repeated_header_field( + monkeypatch: pytest.MonkeyPatch, +): + """Repeated cookie and x-forwarded-for fields all reach the router data.""" + monkeypatch.setattr( + event_namespace, "get_config", lambda: Mock(trusted_proxy_hops=2) + ) + router_data = event_namespace.build_static_router_data( + "sid1", + { + "headers": [ + (b"cookie", b"sid=1"), + (b"x-forwarded-for", b"203.0.113.7"), + (b"cookie", b"other=2"), + (b"x-forwarded-for", b"10.0.0.2"), + ], + "client": ("127.0.0.1", 1234), + }, + ) + + headers = router_data[constants.RouteVar.HEADERS] + assert headers["cookie"] == "sid=1; other=2" + assert headers["x-forwarded-for"] == "203.0.113.7, 10.0.0.2" + # The original client hop comes first, not the proxy that added the second field. + assert router_data[constants.RouteVar.CLIENT_IP] == "203.0.113.7" + assert HeaderData.from_router_data(router_data).cookie == "sid=1; other=2" + + +@pytest.mark.parametrize( + ("trusted_proxy_hops", "forwarded_for", "peer_ip", "expected_ip"), + [ + (0, "6.6.6.6, 203.0.113.7", "10.0.0.2", "10.0.0.2"), + (1, "203.0.113.7", "10.0.0.2", "203.0.113.7"), + (2, "6.6.6.6, 203.0.113.7, 10.0.0.2", "127.0.0.1", "203.0.113.7"), + (2, "203.0.113.7", "10.0.0.2", "10.0.0.2"), + ], +) +def test_static_router_data_resolves_client_ip_from_trusted_proxy_hops( + monkeypatch: pytest.MonkeyPatch, + trusted_proxy_hops: int, + forwarded_for: str, + peer_ip: str, + expected_ip: str, +): + """Only configured rightmost forwarded hops determine the client IP.""" + monkeypatch.setattr( + event_namespace, + "get_config", + lambda: Mock(trusted_proxy_hops=trusted_proxy_hops), + ) + router_data = event_namespace.build_static_router_data( + "sid1", + { + "headers": [(b"x-forwarded-for", forwarded_for.encode())], + "client": (peer_ip, 1234), + }, + ) + + assert router_data[constants.RouteVar.CLIENT_IP] == expected_ip + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", ["current", "replacement", "remote", "missing"]) +async def test_disconnect_only_cancels_current_local_socket( + namespace: WebsocketEventNamespace, namespace_app: Mock, owner: str +): + """A stale disconnect must not cancel a replacement connection's page load. + + Args: + namespace: The namespace. + namespace_app: The mock app it serves. + owner: The socket currently owning the client token. + """ + manager = namespace._token_manager + manager.sid_to_token["old-sid"] = "token" + if owner != "missing": + manager.token_to_socket["token"] = SocketRecord( + instance_id="remote" if owner == "remote" else manager.instance_id, + sid="new-sid" if owner == "replacement" else "old-sid", + ) + manager.disconnect_token = AsyncMock() + cleanup = namespace.handle_disconnect("old-sid") + assert cleanup is not None + await cleanup + cancel = namespace_app._event_processor._on_disconnect + if owner == "current": + cancel.assert_called_once_with("token") + else: + cancel.assert_not_called() + assert namespace.handle_disconnect("unknown-sid") is None + + +@pytest.mark.asyncio +async def test_link_token_to_sid_records_the_connecting_identity( + namespace: WebsocketEventNamespace, + linked_state: Mock, + monkeypatch: pytest.MonkeyPatch, +): + """The session var carries the token the state was loaded under. + + Duplicate-token handling hands back a fresh token, and the state is loaded + under it. Leaving `rx_router_session.client_token` empty until the first + event would let anything reading it in between -- a background task, a + shared-state link -- address the wrong state tree. + """ + state = linked_state + + # No duplicate: the connecting token is recorded. + await namespace.link_token_to_sid("sid1", "tok1") + assert state.router_data[constants.RouteVar.CLIENT_TOKEN] == "tok1" + assert state.rx_router_session.client_token == "tok1" + assert state.rx_router_session.session_id == "sid1" + + # Duplicate: the *new* token is recorded, not the one the client sent. + monkeypatch.setattr( + namespace._token_manager, + "link_token_to_sid", + AsyncMock(return_value="a-fresh-token"), + ) + await namespace.link_token_to_sid("sid2", "tok1") + assert state.router_data[constants.RouteVar.CLIENT_TOKEN] == "a-fresh-token" + assert state.rx_router_session.client_token == "a-fresh-token" + assert state.rx_router_session.session_id == "sid2" + + +def test_client_error_constants_match_frontend(): + """The socket event and error types are duplicated as literals in the client. + + Nothing at runtime keeps the definitions in sync, so pin them here. + """ + assert ( + "const HYDRATION_PAYLOAD_TOO_LARGE = " + f'"{constants.ClientErrorType.HYDRATION_PAYLOAD_TOO_LARGE}";' + in WEBSOCKET_JS_TEMPLATE.read_text() + ) + state_js = (WEB_TEMPLATE_ROOT / "utils/state.js").read_text() + assert ( + f'const CLIENT_ERROR_EVENT = "{constants.SocketEvent.CLIENT_ERROR}"' in state_js + ) + assert ( + f'const ERROR_TYPE_DISPATCH_MISSING = "{constants.ClientErrorType.DISPATCH_MISSING}"' + in state_js + ) + assert ( + f'const ERROR_TYPE_STATE_UPDATE = "{constants.ClientErrorType.STATE_UPDATE}"' + in state_js + ) + assert ( + "const ERROR_TYPE_HYDRATION_PAYLOAD_TOO_LARGE = " + f'"{constants.ClientErrorType.HYDRATION_PAYLOAD_TOO_LARGE}"' in state_js + ) diff --git a/tests/units/test_socketio_namespace.py b/tests/units/test_socketio_namespace.py new file mode 100644 index 00000000000..d832dcad327 --- /dev/null +++ b/tests/units/test_socketio_namespace.py @@ -0,0 +1,821 @@ +"""Tests for the Socket.IO event transport in reflex/socketio_namespace.py.""" + +import asyncio +import logging +from typing import Any +from unittest.mock import AsyncMock, Mock + +import pytest +from reflex_base import constants, otel + +from reflex.socketio_namespace import EventNamespace, _sio_dumps, _sio_loads + +from .conftest import metric_points + + +@pytest.fixture +def namespace(namespace_app: Mock) -> EventNamespace: + """A namespace serving the mock app. + + Returns: + The namespace. + """ + return EventNamespace("/_event", namespace_app) + + +@pytest.mark.asyncio +async def test_connect_with_token_is_accepted(namespace: EventNamespace): + """A client that authenticates keeps its connection and its ASGI scope.""" + scope = {"type": "websocket", "headers": []} + accepted = await namespace.on_connect( + "sid1", {"QUERY_STRING": "token=tok1", "asgi.scope": scope} + ) + + assert accepted is not False + assert namespace.sid_to_token["sid1"] == "tok1" + assert namespace._scopes["sid1"] is scope + + +@pytest.mark.asyncio +async def test_connect_without_token_is_refused(namespace: EventNamespace): + """A connection that never links a token is closed, not left open. + + Nothing it sends can be served, and Socket.IO does not run the disconnect + handler for a refused connect, so the namespace has to undo its own + bookkeeping here. + """ + refused = await namespace.on_connect( + "sid1", {"QUERY_STRING": "", "asgi.scope": {"type": "websocket"}} + ) + + assert refused is False + assert "sid1" not in namespace.sid_to_token + assert namespace._scopes == {} + + +@pytest.mark.asyncio +async def test_refused_connect_is_not_counted_as_a_connection( + namespace: EventNamespace, otel_metrics +): + """A refused connection leaves the gauge where it was, not one above it.""" + await namespace.on_connect("sid1", {"QUERY_STRING": ""}) + + (connections,) = metric_points(otel_metrics, otel.METRIC_WEBSOCKET_CONNECTIONS) + assert connections.value == 0 + + +@pytest.mark.asyncio +async def test_connect_processes_boot_event_from_auth( + namespace: EventNamespace, + namespace_app: Mock, + linked_state: Mock, + boot_event: dict[str, Any], +): + """The hydrate event carried in the CONNECT packet is processed on connect. + + As the connection's first event it records the new sid and token on the + state, so the connect does not load and save the state tree for that; a + connect without a boot event still does. + """ + modify_state = namespace_app.state_manager.modify_state + environ = { + "QUERY_STRING": "token=tok1", + "asgi.scope": {"type": "websocket", "headers": []}, + } + + await namespace.on_connect("sid1", environ, {"event": boot_event}) + + token, event = namespace_app.event_processor.enqueue.await_args.args + assert (token, event.name) == ("tok1", "state.hydrate_and_load") + modify_state.assert_not_called() + + # Without a boot event (or without auth at all) nothing is processed, and + # the connect records the new sid and token on the state itself. + namespace_app.event_processor.enqueue.reset_mock() + await namespace.on_connect("sid2", {"QUERY_STRING": "token=tok2"}, None) + await namespace.on_connect("sid3", {"QUERY_STRING": "token=tok3"}) + namespace_app.event_processor.enqueue.assert_not_awaited() + assert modify_state.call_count == 2 + assert linked_state.router_data[constants.RouteVar.SESSION_ID] == "sid3" + assert linked_state.router_data[constants.RouteVar.CLIENT_TOKEN] == "tok3" + + +@pytest.mark.asyncio +async def test_connect_reports_state_schema_mismatch( + namespace: EventNamespace, + namespace_app: Mock, + caplog: pytest.LogCaptureFixture, + boot_event: dict[str, Any], +): + """A frontend compiled for another state tree is told once, and still hydrates.""" + namespace_app._state = Mock(_to_frontend_schema=Mock(return_value="backend-schema")) + emit = namespace.emit = AsyncMock() + + def warnings() -> list[str]: + return [ + r.getMessage() + for r in caplog.records + if r.name == "reflex.event_namespace" and r.levelno == logging.WARNING + ] + + with caplog.at_level(logging.WARNING, logger="reflex.event_namespace"): + await namespace.on_connect( + "sid1", + {"QUERY_STRING": "token=tok1", "asgi.scope": {"headers": []}}, + {"event": boot_event, "state_schema": "frontend-schema"}, + ) + + emit.assert_awaited_once_with( + "state_mismatch", + {"backend": "backend-schema", "reload_cooldown_ms": 10_000}, + to="sid1", + ) + token, event = namespace_app.event_processor.enqueue.await_args.args + assert (token, event.name) == ("tok1", "state.hydrate_and_load") + (message,) = warnings() + assert "frontend-schema" in message + assert "backend-schema" in message + assert "sid1" in message + + # The client-supplied hash is sanitized and bounded before it is logged. + caplog.clear() + await namespace.on_connect( + "sid2", + {"QUERY_STRING": "token=tok2", "asgi.scope": {"headers": []}}, + {"event": boot_event, "state_schema": "x\nforged line [bold]" + "y" * 500}, + ) + + (message,) = warnings() + assert "forged" in message + assert "\n" not in message + assert "\\[bold]" in message + assert "(truncated)" in message + assert len(message) < 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth", + [ + pytest.param( + { + "event": {"name": "state.hydrate_and_load", "payload": {}}, + "state_schema": "same", + }, + id="matching", + ), + pytest.param( + {"event": {"name": "state.hydrate_and_load", "payload": {}}}, + id="old-frontend", + ), + pytest.param(None, id="no-auth"), + ], +) +async def test_connect_sends_no_mismatch_without_one( + namespace: EventNamespace, + namespace_app: Mock, + linked_state: Mock, + auth: dict[str, Any] | None, +): + """A matching or missing frontend schema sends nothing to the client.""" + namespace_app._state._to_frontend_schema = Mock(return_value="same") + emit = namespace.emit = AsyncMock() + + await namespace.on_connect( + "sid1", {"QUERY_STRING": "token=tok1", "asgi.scope": {"headers": []}}, auth + ) + + emit.assert_not_awaited() + assert namespace_app.event_processor.enqueue.await_count == (1 if auth else 0) + + +@pytest.mark.asyncio +async def test_connect_whose_boot_event_fails_is_refused( + namespace: EventNamespace, + namespace_app: Mock, + otel_metrics, + caplog, + boot_event: dict[str, Any], +): + """A boot event that fails to process refuses the connect and undoes it. + + Raising instead would leave the session registered with Socket.IO yet + unanswered, so the client would neither hydrate nor retry. + """ + namespace_app.event_processor.enqueue.side_effect = ValueError("bad boot event") + environ = { + "QUERY_STRING": "token=tok1", + "asgi.scope": {"type": "websocket", "headers": []}, + } + + with caplog.at_level(logging.ERROR, logger="reflex.event_namespace"): + refused = await namespace.on_connect("sid1", environ, {"event": boot_event}) + # Let the token cleanup task run. + await asyncio.sleep(0) + + assert refused is False + assert "Error handling the boot event" in caplog.text + assert "sid1" not in namespace.sid_to_token + assert "tok1" not in namespace.token_to_sid + assert namespace._scopes == {} + (connections,) = metric_points(otel_metrics, otel.METRIC_WEBSOCKET_CONNECTIONS) + assert connections.value == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data", ['{"name": "state.on_click", "payload": {}, "router_data": {}}', 42] +) +async def test_undeserializable_event_disconnects_the_session( + namespace: EventNamespace, namespace_app: Mock, caplog, data: object +): + """A malformed event ends the session, as on the plain transport. + + Escaping into python-socketio, the error logged a traceback per frame, + outside any error budget, and the session stayed open. + """ + await namespace.on_connect( + "sid1", {"QUERY_STRING": "token=tok1", "asgi.scope": {"headers": []}} + ) + namespace.disconnect = AsyncMock() + + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + await namespace.on_event("sid1", data) + + namespace.disconnect.assert_awaited_once_with("sid1") + namespace_app.event_processor.enqueue.assert_not_awaited() + assert all(record.levelno <= logging.DEBUG for record in caplog.records) + + +@pytest.mark.asyncio +async def test_connect_with_a_malformed_boot_event_is_refused( + namespace: EventNamespace, namespace_app: Mock, caplog: pytest.LogCaptureFixture +): + """A boot event that does not deserialize refuses the connect, quietly. + + Disconnecting from inside the connect handler, as a later malformed event + does, would acknowledge a session that never hydrated. The input is + client-controlled, so it is logged at debug level, as on the plain + transport, rather than as a handler error with a traceback. + """ + namespace.disconnect = AsyncMock() + environ = {"QUERY_STRING": "token=tok1", "asgi.scope": {"headers": []}} + + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + refused = await namespace.on_connect("sid1", environ, {"event": "not an event"}) + # Let the token cleanup task run. + await asyncio.sleep(0) + + assert refused is False + namespace.disconnect.assert_not_awaited() + namespace_app.event_processor.enqueue.assert_not_awaited() + assert "tok1" not in namespace.token_to_sid + assert all(record.levelno <= logging.DEBUG for record in caplog.records) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("handler", ["event", "ping", "client_error"]) +async def test_a_failing_handler_is_logged_within_the_error_budget( + namespace: EventNamespace, + namespace_app: Mock, + caplog: pytest.LogCaptureFixture, + handler: str, +): + """A server-side handler failure is logged within budget and keeps the session. + + Escaping into python-socketio, it would be logged with a traceback per + frame, outside the budget that bounds client-triggered error logs. + """ + await namespace.on_connect( + "sid1", {"QUERY_STRING": "token=tok1", "asgi.scope": {"headers": []}} + ) + namespace.disconnect = AsyncMock() + failure = RuntimeError("server bug") + namespace_app.event_processor.enqueue.side_effect = failure + namespace_app.frontend_exception_handler = Mock(side_effect=failure) + namespace.emit = AsyncMock(side_effect=failure) + calls = { + "event": lambda: namespace.on_event( + "sid1", {"name": "state.on_click", "payload": {}, "router_data": {}} + ), + "ping": lambda: namespace.on_ping("sid1"), + "client_error": lambda: namespace.on_client_error( + "sid1", {"error_type": "x", "message": "y"} + ), + } + + with caplog.at_level(logging.ERROR, logger="reflex.event_namespace"): + for _ in range(namespace._MAX_CLIENT_ERRORS_PER_SID + 2): + await calls[handler]() + + assert len(caplog.records) == namespace._MAX_CLIENT_ERRORS_PER_SID + assert all("server bug" in repr(record.exc_info) for record in caplog.records) + namespace.disconnect.assert_not_awaited() + + +class AuditingNamespace(EventNamespace): + """A namespace an app customized through on_event, recording what it is handed.""" + + def __init__(self, namespace: str, app: Mock): + """Initialize the namespace.""" + super().__init__(namespace, app) + self.seen: list[Any] = [] + + async def on_event(self, sid: str, data: Any): + """Record the event, then handle it.""" + self.seen.append(data) + await super().on_event(sid, data) + + +@pytest.fixture +def auditing_namespace(namespace_app: Mock) -> AuditingNamespace: + """A namespace overriding on_event, serving the mock app. + + Returns: + The namespace. + """ + return AuditingNamespace("/_event", namespace_app) + + +@pytest.mark.asyncio +async def test_an_on_event_override_sees_the_boot_event_of_every_connect( + auditing_namespace: AuditingNamespace, + namespace_app: Mock, + boot_event: dict[str, Any], +): + """An on_event override sees the hydrate on connect and on reconnect. + + Apps override it for authorization checks, tenant routing or auditing, + which the boot event must not bypass. + """ + connected = await auditing_namespace.on_connect( + "sid1", _connect_environ("tok1"), {"event": boot_event} + ) + cleanup = auditing_namespace.on_disconnect("sid1") + assert cleanup is not None + await cleanup + reconnected = await auditing_namespace.on_connect( + "sid2", _connect_environ("tok1"), {"event": boot_event} + ) + + assert connected is None + assert reconnected is None + assert auditing_namespace.seen == [boot_event, boot_event] + assert [ + call.args[0] for call in namespace_app.event_processor.enqueue.await_args_list + ] == ["tok1", "tok1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("malformed", "errors"), + [ + pytest.param( + False, ["Error handling the boot event for session sid1."], id="raising" + ), + pytest.param(True, [], id="malformed"), + ], +) +async def test_a_failing_boot_event_refuses_the_connect_under_an_on_event_override( + auditing_namespace: AuditingNamespace, + namespace_app: Mock, + caplog: pytest.LogCaptureFixture, + boot_event: dict[str, Any], + malformed: bool, + errors: list[str], +): + """A boot event failing inside an on_event override still refuses the connect. + + on_event swallows a later event's failure, but acknowledged, a session + whose boot event failed would leave the client waiting for a hydrate that + never comes. A malformed one is refused quietly, without the disconnect a + later malformed event gets. + """ + boot: Any = "not an event" if malformed else boot_event + namespace_app.event_processor.enqueue.side_effect = ValueError("bad boot event") + auditing_namespace.disconnect = AsyncMock() + + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + refused = await auditing_namespace.on_connect( + "sid1", _connect_environ("tok1"), {"event": boot} + ) + # Let the token cleanup task run. + await asyncio.sleep(0) + + assert refused is False + assert auditing_namespace.seen == [boot] + auditing_namespace.disconnect.assert_not_awaited() + assert "tok1" not in auditing_namespace.token_to_sid + assert auditing_namespace._scopes == {} + assert [ + record.getMessage() + for record in caplog.records + if record.levelno > logging.DEBUG + ] == errors + + +@pytest.mark.asyncio +async def test_later_events_keep_the_failure_policy_under_an_on_event_override( + auditing_namespace: AuditingNamespace, + namespace_app: Mock, + caplog: pytest.LogCaptureFixture, + boot_event: dict[str, Any], +): + """Past the boot event, a failing event is logged and a malformed one disconnects. + + Only the boot event's failure is raised out of on_event; raised into + python-socketio, a later one would log a traceback per frame, outside the + session's error budget. + """ + await auditing_namespace.on_connect( + "sid1", _connect_environ("tok1"), {"event": boot_event} + ) + auditing_namespace.disconnect = AsyncMock() + namespace_app.event_processor.enqueue.side_effect = RuntimeError("server bug") + + with caplog.at_level(logging.ERROR, logger="reflex.event_namespace"): + await auditing_namespace.on_event( + "sid1", {"name": "state.on_click", "payload": {}, "router_data": {}} + ) + auditing_namespace.disconnect.assert_not_awaited() + await auditing_namespace.on_event("sid1", 42) + + assert len(auditing_namespace.seen) == 3 + assert [ + record.getMessage() + for record in caplog.records + if record.levelno >= logging.ERROR + ] == ["Error handling socket event 'event' for session sid1."] + auditing_namespace.disconnect.assert_awaited_once_with("sid1") + + +@pytest.mark.asyncio +async def test_a_handle_event_override_sees_every_event( + namespace_app: Mock, boot_event: dict[str, Any] +): + """A handle_event override sees the hydrate on connect and every later event.""" + seen: list[str] = [] + + class TenantNamespace(EventNamespace): + async def handle_event(self, sid, data, asgi_scope): + seen.append(data["name"]) + await super().handle_event(sid, data, asgi_scope) + + namespace = TenantNamespace("/_event", namespace_app) + await namespace.on_connect("sid1", _connect_environ("tok1"), {"event": boot_event}) + await namespace.on_event( + "sid1", {"name": "state.on_click", "payload": {}, "router_data": {}} + ) + + assert seen == ["state.hydrate_and_load", "state.on_click"] + + +def test_packets_with_an_unpaired_surrogate_stay_valid_utf8(): + """A packet carrying an unpaired surrogate is escaped, as plain frames are. + + A non-UTF-8 file name decodes to one; written raw, the text frame cannot + be encoded for sending and the update is lost. + """ + data = _sio_dumps({"file": "bad\udcffname"}, separators=(",", ":")) + + data.encode() + assert _sio_loads(data) == {"file": "bad\udcffname"} + + +@pytest.fixture +def frontend_errors(namespace: EventNamespace, namespace_app: Mock) -> list[str]: + """Capture what a linked session's client_error reports route to the app. + + Returns: + The reported messages. + """ + namespace.sid_to_token["known_sid"] = "some_token" + errors: list[str] = [] + namespace_app.frontend_exception_handler = lambda exc: errors.append(str(exc)) + return errors + + +def _connect_environ(token: str) -> dict[str, Any]: + """The environ of a connecting client carrying a token and two headers. + + Args: + token: The client token. + + Returns: + The environ. + """ + return { + "QUERY_STRING": f"token={token}", + "asgi.scope": { + "headers": [ + (b"origin", b"http://localhost:3000"), + (b"user-agent", b"test-agent"), + ], + "client": ("127.0.0.1", 1234), + }, + } + + +def _client_event() -> dict[str, Any]: + """An event as the frontend sends it. + + Returns: + The event payload. + """ + return { + "name": "state.hydrate", + "router_data": {"pathname": "/", "query": {}, "asPath": "/"}, + "payload": {}, + } + + +def _enqueued(namespace_app: Mock) -> list[Any]: + """The events the namespace handed to the event processor. + + Args: + namespace_app: The mock app. + + Returns: + The enqueued events, in order. + """ + return [ + call.args[1] for call in namespace_app.event_processor.enqueue.await_args_list + ] + + +def test_sio_json_records_message_sizes(otel_metrics): + """Socket.IO packet serialization records UTF-8 sizes in both directions.""" + data = _sio_dumps({"a": "é"}, separators=(",", ":")) + assert data == '{"a":"é"}' + assert _sio_loads(data) == {"a": "é"} + assert _sio_loads(data.encode()) == {"a": "é"} + # ASCII payloads take the fast path that sizes without encoding a copy. + ascii_data = _sio_dumps({"b": "x"}, separators=(",", ":")) + assert _sio_loads(ascii_data) == {"b": "x"} + points = { + p.attributes[otel.ATTR_NETWORK_IO_DIRECTION]: p.sum + for p in metric_points(otel_metrics, otel.METRIC_WEBSOCKET_MESSAGE_SIZE) + } + size = len(data.encode()) + assert size == len(data) + 1 + assert points == { + "transmit": size + len(ascii_data), + "receive": 2 * size + len(ascii_data), + } + + +@pytest.mark.asyncio +async def test_client_error_dispatch_missing_reports_actionable_error( + namespace: EventNamespace, frontend_errors: list[str] +): + """A dispatch_function_missing error reports the substate and the remedy.""" + await namespace.on_client_error( + "known_sid", + { + "error_type": constants.ClientErrorType.DISPATCH_MISSING, + "message": "Cannot process state update", + "substate": "reflex___state____state.my___state____my_state", + }, + ) + (message,) = frontend_errors + assert "reflex___state____state.my___state____my_state" in message + assert "skipped" in message.lower() + assert "rebuild" in message.lower() + + +@pytest.mark.asyncio +async def test_client_error_generic_reports_type_and_message( + namespace: EventNamespace, frontend_errors: list[str] +): + """A generic client error reports the error type and message.""" + await namespace.on_client_error( + "known_sid", + {"error_type": constants.ClientErrorType.STATE_UPDATE, "message": "boom"}, + ) + (message,) = frontend_errors + assert constants.ClientErrorType.STATE_UPDATE in message + assert "boom" in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("payload", ["not a dict", None, ["list"], 42]) +async def test_client_error_malformed_payload_is_ignored( + namespace: EventNamespace, frontend_errors: list[str], payload: Any +): + """Non-dict payloads are dropped without raising or reporting errors.""" + await namespace.on_client_error("known_sid", payload) + assert not frontend_errors + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sid", ["known_sid", "unknown_sid"]) +async def test_client_error_no_argument_emit_is_ignored( + namespace: EventNamespace, + frontend_errors: list[str], + caplog: pytest.LogCaptureFixture, + sid: str, +): + """A payload-less client_error emit is dropped like any malformed payload. + + python-socketio dispatches ``emit("client_error")`` without a payload as + ``on_client_error(sid)``, so a missing ``data`` argument must not raise a + TypeError before the anti-abuse guards run. + """ + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + await namespace.on_client_error(sid) + assert not frontend_errors + assert all(record.levelno <= logging.DEBUG for record in caplog.records) + + +@pytest.mark.asyncio +async def test_client_error_unknown_sid_does_not_report_error( + namespace: EventNamespace, + frontend_errors: list[str], + caplog: pytest.LogCaptureFixture, +): + """Errors from sockets without a linked token are not reported.""" + with caplog.at_level(logging.DEBUG, logger="reflex.event_namespace"): + await namespace.on_client_error( + "unknown_sid", + { + "error_type": constants.ClientErrorType.STATE_UPDATE, + "message": "spam from unauthenticated socket", + }, + ) + assert not frontend_errors + assert all(record.levelno < logging.ERROR for record in caplog.records) + + +@pytest.mark.asyncio +async def test_client_error_values_are_sanitized_and_truncated( + namespace: EventNamespace, frontend_errors: list[str] +): + """Control characters are stripped and long messages truncated before reporting.""" + evil = "\x1b[31mINJECT\x1b[0m\nFAKE LOG LINE\t" + "A" * 5000 + await namespace.on_client_error( + "known_sid", {"error_type": "custom_type", "message": evil} + ) + (message,) = frontend_errors + assert "\x1b" not in message + assert "\n" not in message + assert "\t" not in message + assert len(message) < 700 + + +@pytest.mark.asyncio +async def test_client_error_reporting_is_rate_limited_per_sid( + namespace: EventNamespace, frontend_errors: list[str] +): + """A single session cannot flood the backend logs with error reports.""" + for _ in range(20): + await namespace.on_client_error( + "known_sid", {"error_type": "custom_type", "message": "spam"} + ) + assert len(frontend_errors) == EventNamespace._MAX_CLIENT_ERRORS_PER_SID + # Disconnecting removes the counter so the mapping cannot grow unboundedly. + if (task := namespace.on_disconnect("known_sid")) is not None: + await task + assert "known_sid" not in namespace._client_error_budget.counts + + +@pytest.mark.asyncio +async def test_client_error_reporting_bounded_across_reconnects( + namespace: EventNamespace, + frontend_errors: list[str], + caplog: pytest.LogCaptureFixture, +): + """Reconnecting with fresh SIDs does not grant an unlimited report budget.""" + with caplog.at_level(logging.WARNING, logger="reflex.event_namespace"): + for reconnect in range(50): + sid = f"sid_{reconnect}" + namespace.sid_to_token[sid] = f"token_{reconnect}" + for _ in range(5): + await namespace.on_client_error( + sid, {"error_type": "custom_type", "message": "spam"} + ) + assert len(frontend_errors) == EventNamespace._MAX_CLIENT_ERRORS_PER_WINDOW + # Suppression is not silent: one warning is logged when the cap trips, so + # a flooding client cannot invisibly starve reports from other sessions. + assert ["suppressing" in r.getMessage() for r in caplog.records] == [True] + # Once the window elapses, errors are reported again (not silenced forever). + namespace._client_error_budget.window_start -= ( + EventNamespace._CLIENT_ERROR_WINDOW_SECONDS + 1 + ) + namespace.sid_to_token["sid_fresh"] = "token_fresh" + await namespace.on_client_error( + "sid_fresh", {"error_type": "custom_type", "message": "after window"} + ) + assert len(frontend_errors) == EventNamespace._MAX_CLIENT_ERRORS_PER_WINDOW + 1 + + +def test_client_error_event_name_matches_handler(): + """python-socketio dispatches events to on_ methods by name. + + This pins the handler to SocketEvent.CLIENT_ERROR. + """ + assert ( + f"on_{constants.SocketEvent.CLIENT_ERROR}" + == EventNamespace.on_client_error.__name__ + ) + + +@pytest.mark.asyncio +async def test_on_event_uses_connection_scoped_router_data( + token: str, namespace: EventNamespace, namespace_app: Mock +): + """on_event merges the connection-scoped router_data into every event. + + Headers, client IP, and session id cannot change for the lifetime of a + connection, so the decode is cached on the ASGI scope and never repeated + per event. + """ + environ = _connect_environ(token) + await namespace.on_connect("sid1", environ) + + await namespace.on_event("sid1", _client_event()) + + (event,) = _enqueued(namespace_app) + assert namespace_app.event_processor.enqueue.await_args.args[0] == token + assert event.router_data[constants.RouteVar.CLIENT_TOKEN] == token + assert event.router_data[constants.RouteVar.SESSION_ID] == "sid1" + assert event.router_data[constants.RouteVar.CLIENT_IP] == "127.0.0.1" + assert event.router_data[constants.RouteVar.HEADERS] == { + "origin": "http://localhost:3000", + "user-agent": "test-agent", + "asgi-scope-client": "127.0.0.1", + } + assert event.router_data[constants.RouteVar.PATH] == "/404" + assert event.router_data[constants.RouteVar.QUERY] == {} + + # A later event on the same connection reuses the decoded headers. + environ["asgi.scope"]["headers"] = [(b"origin", b"decoded-again")] + await namespace.on_event("sid1", _client_event()) + next_event = _enqueued(namespace_app)[-1] + assert ( + next_event.router_data[constants.RouteVar.HEADERS]["origin"] + == "http://localhost:3000" + ) + + +@pytest.mark.asyncio +async def test_on_event_does_not_share_the_cached_headers( + token: str, namespace: EventNamespace, namespace_app: Mock +): + """Each event gets its own headers mapping, not the cached one. + + The headers reach `state.router_data`, a plain mutable dict, so sharing + the cached mapping would let a handler mutating it corrupt the connection + cache for every later event on the socket. + """ + environ = _connect_environ(token) + await namespace.on_connect("sid1", environ) + + await namespace.on_event("sid1", _client_event()) + (event,) = _enqueued(namespace_app) + event_headers = event.router_data[constants.RouteVar.HEADERS] + cached_headers = environ["asgi.scope"]["_reflex_static_router_data"][ + constants.RouteVar.HEADERS + ] + + assert event_headers == cached_headers + assert event_headers is not cached_headers + # Mutating what the handler sees must not reach the connection cache. + event_headers["user-agent"] = "mutated" + assert cached_headers["user-agent"] == "test-agent" + + await namespace.on_event("sid1", _client_event()) + next_event = _enqueued(namespace_app)[-1] + assert ( + next_event.router_data[constants.RouteVar.HEADERS]["user-agent"] == "test-agent" + ) + + +@pytest.mark.asyncio +async def test_on_event_builds_router_data_without_connect( + token: str, namespace: EventNamespace, namespace_app: Mock +): + """on_event derives the connection data from the environ if connect was missed.""" + await namespace._token_manager.link_token_to_sid(token, "sid1") + environ = _connect_environ(token) + namespace_app.sio = Mock(get_environ=Mock(return_value=environ)) + + await namespace.on_event("sid1", _client_event()) + await namespace.on_event("sid1", _client_event()) + + # The connection-scoped entries are derived once and cached on the scope. + assert environ["asgi.scope"]["_reflex_static_router_data"] == { + constants.RouteVar.SESSION_ID: "sid1", + constants.RouteVar.CLIENT_IP: "127.0.0.1", + constants.RouteVar.HEADERS: { + "origin": "http://localhost:3000", + "user-agent": "test-agent", + "asgi-scope-client": "127.0.0.1", + }, + } + events = _enqueued(namespace_app) + assert len(events) == 2 + for event in events: + assert event.router_data[constants.RouteVar.SESSION_ID] == "sid1" + assert ( + event.router_data[constants.RouteVar.HEADERS]["user-agent"] == "test-agent" + ) diff --git a/tests/units/test_testing.py b/tests/units/test_testing.py index f1712907b11..d2e6c27b45c 100644 --- a/tests/units/test_testing.py +++ b/tests/units/test_testing.py @@ -1,15 +1,25 @@ """Unit tests for the included testing tools.""" +import asyncio +import contextlib +import socket import sys +import threading +import time +from collections.abc import Callable, Iterator +from pathlib import Path from types import ModuleType, SimpleNamespace from unittest import mock +import httpx import pytest import reflex_base.config from reflex_base.components.memo import MEMOS from reflex_base.constants import IS_WINDOWS from reflex_base.environment import environment from reflex_base.registry import RegistrationContext +from reflex_base.utils.types import ASGIApp +from starlette.applications import Starlette import reflex.constants import reflex.reflex as reflex_cli @@ -194,6 +204,247 @@ def test_app_harness_initialize_reloads_existing_imported_app( harness_mocks.get_and_validate_app.assert_called_once_with(reload=True) +def _asgi_app(on_startup: Callable[[], object] | None = None) -> ASGIApp: + """A minimal ASGI app with lifespan support that answers every request. + + Args: + on_startup: called (in a worker thread) during lifespan startup. + + Returns: + The ASGI app. + """ + + async def app(scope, receive, send): + if scope["type"] == "lifespan": + while (await receive())["type"] == "lifespan.startup": + if on_startup is not None: + await asyncio.to_thread(on_startup) + await send({"type": "lifespan.startup.complete"}) + await send({"type": "lifespan.shutdown.complete"}) + return + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"embedded"}) + + return app + + +@contextlib.contextmanager +def _running(server: reflex_testing._EmbeddedServer) -> Iterator[threading.Thread]: + """Run the server in a thread, stopping it on exit. + + Args: + server: the embedded server. + + Yields: + The thread serving the app. + """ + thread = threading.Thread(target=server.run) + thread.start() + try: + yield thread + finally: + server.should_exit = True + thread.join(timeout=15) + assert not thread.is_alive() + + +def _assert_port_taken(host: str, port: int): + """Assert that no socket can bind the port, even one asking to share it. + + Args: + host: the address of the port. + port: the port number. + """ + with socket.socket() as thief: + thief.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + if reuse_port := getattr(socket, "SO_REUSEPORT", None): + thief.setsockopt(socket.SOL_SOCKET, reuse_port, 1) + with pytest.raises(OSError): + thief.bind((host, port)) + thief.listen() + + +def _assert_port_free(host: str, port: int): + """Assert that the port can be bound again. + + Args: + host: the address of the port. + port: the port number. + """ + with socket.socket() as rebind: + rebind.bind((host, port)) + + +def _harness_with_backend( + tmp_path, server: reflex_testing._EmbeddedServer +) -> AppHarness: + """An AppHarness whose backend is the given server. + + Args: + tmp_path: the app root. + server: the backend server. + + Returns: + The harness. + """ + harness = AppHarness.create(root=tmp_path, app_name="embedded_app") + harness.backend = server + return harness + + +def test_embedded_server_owns_its_port(): + """The server holds its port from construction on, so no other socket can share it.""" + server = reflex_testing._EmbeddedServer(_asgi_app()) + _assert_port_taken(server.host, server.port) + with _running(server): + server.wait_started(timeout=15) + _assert_port_taken(server.host, server.port) + assert httpx.get(f"http://{server.host}:{server.port}/").content == b"embedded" + + +def _accepted_nodelay(client: socket.socket) -> int: + """Read TCP_NODELAY of the in-process server end of a connection. + + Args: + client: the client end of the connection. + + Returns: + The TCP_NODELAY value of the accepted socket. + """ + peer = client.getsockname() + for fd in Path("/dev/fd").iterdir(): + try: + sock = socket.fromfd(int(fd.name), socket.AF_INET, socket.SOCK_STREAM) + except OSError: # closed since the listing + continue + # Not every fd is a connected socket. + with sock, contextlib.suppress(OSError): + if sock.getpeername() == peer: + return sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY) + pytest.fail("The server end of the connection was not found.") + + +@pytest.mark.skipif(IS_WINDOWS, reason="Finds the accepted socket through /dev/fd.") +def test_embedded_server_connections_use_tcp_nodelay(): + """Accepted connections send small websocket frames without waiting for an ACK.""" + server = reflex_testing._EmbeddedServer(_asgi_app()) + with _running(server): + server.wait_started(timeout=15) + with socket.create_connection(server.getsockname()) as client: + client.sendall(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n") + assert client.recv(1024).startswith(b"HTTP/1.1 200") + assert _accepted_nodelay(client) + + +def test_poll_for_servers_waits_for_the_backend_itself(tmp_path): + """Readiness follows the app's startup, not a connectable port. + + Args: + tmp_path: pytest tmp_path fixture + """ + startup = threading.Event() + server = reflex_testing._EmbeddedServer(_asgi_app(on_startup=startup.wait)) + harness = _harness_with_backend(tmp_path, server) + # Another listener on the port must not pass for the backend. + foreign = socket.socket() + with contextlib.suppress(OSError): + foreign.bind((server.host, server.port)) + foreign.listen() + with foreign, _running(server): + try: + with pytest.raises(TimeoutError): + harness._poll_for_servers(timeout=0.5) + finally: + startup.set() + assert harness._poll_for_servers(timeout=15).getsockname() == ( + server.host, + server.port, + ) + assert httpx.get(f"http://{server.host}:{server.port}/").content == b"embedded" + + +def test_embedded_server_refuses_a_taken_requested_port(): + """A requested port that is in use fails loudly instead of moving.""" + with socket.socket() as blocker: + blocker.bind(("127.0.0.1", 0)) + blocker.listen() + port = blocker.getsockname()[1] + with pytest.raises(OSError): + reflex_testing._EmbeddedServer(_asgi_app(), port=port) + server = reflex_testing._EmbeddedServer(_asgi_app(), port=port) + assert server.getsockname() == ("127.0.0.1", port) + server.should_exit = True + server.run() + + +def test_failed_backend_startup_logs_and_fails_fast(tmp_path, capsys): + """A failing lifespan startup is logged and fails readiness without the timeout. + + Args: + tmp_path: pytest tmp_path fixture + capsys: pytest capsys fixture + """ + + @contextlib.asynccontextmanager + async def lifespan(app): + msg = "lifespan boom" + raise RuntimeError(msg) + yield + + server = reflex_testing._EmbeddedServer(Starlette(lifespan=lifespan)) + harness = _harness_with_backend(tmp_path, server) + with _running(server): + started = time.monotonic() + with pytest.raises(RuntimeError, match="stopped without serving"): + harness._poll_for_servers(timeout=10) + assert time.monotonic() - started < 5 + out = capsys.readouterr().out + assert "ASGI lifespan startup failed" in out + assert "RuntimeError: lifespan boom" in out + # The socket is released although no worker ever served it. + _assert_port_free(server.host, server.port) + + +def test_embedded_server_stop_releases_port(): + """Stopping a serving server runs its shutdown hook and frees the port.""" + shutdown = mock.AsyncMock() + server = reflex_testing._EmbeddedServer(_asgi_app()) + server.shutdown = shutdown + with _running(server): + server.wait_started(timeout=15) + shutdown.assert_awaited_once() + _assert_port_free(server.host, server.port) + + +def test_embedded_server_stopped_before_run_never_serves(): + """A server stopped before it runs returns at once and releases its port.""" + shutdown = mock.AsyncMock() + server = reflex_testing._EmbeddedServer(_asgi_app()) + server.shutdown = shutdown + server.should_exit = True + server.run() + shutdown.assert_awaited_once() + with pytest.raises(RuntimeError, match="stopped without serving"): + server.wait_started(timeout=0) + _assert_port_free(server.host, server.port) + + +def test_embedded_server_started_is_deprecated(monkeypatch): + """The uvicorn-style `started` flag still works, with a deprecation warning. + + Args: + monkeypatch: pytest monkeypatch fixture + """ + deprecate = mock.Mock() + monkeypatch.setattr(reflex_testing.console, "deprecate", deprecate) + server = reflex_testing._EmbeddedServer(_asgi_app()) + assert server.started is False + with _running(server): + server.wait_started(timeout=15) + assert server.started is True + assert deprecate.call_args.kwargs["feature_name"] == "AppHarness.backend.started" + + def test_app_harness_frontend_env_has_development_condition( tmp_path, monkeypatch: pytest.MonkeyPatch, harness_mocks ) -> None: diff --git a/tests/units/utils/test_build.py b/tests/units/utils/test_build.py index 9128d931ce5..f57a88ad6b5 100644 --- a/tests/units/utils/test_build.py +++ b/tests/units/utils/test_build.py @@ -292,13 +292,16 @@ def check_final_tree(directory: Path, formats: tuple[str, ...]): def _patch_env_json( - mocker: MockerFixture, tmp_path: Path, plugins: list[Plugin] | None = None + mocker: MockerFixture, + tmp_path: Path, + plugins: list[Plugin] | None = None, + transport: str = "websocket", ): web_dir = tmp_path / ".web" web_dir.mkdir() mocker.patch("reflex.utils.build.prerequisites.get_web_dir", return_value=web_dir) config = mocker.Mock() - config.transport = "websocket" + config.transport = transport config.plugins = plugins or [] mocker.patch("reflex.utils.build.get_config", return_value=config) mocker.patch("reflex.utils.build.is_in_app_harness", return_value=False) @@ -352,3 +355,33 @@ def update_env_json(self, **context): assert env["SHARED"] == "second" assert env["ONLY_FIRST"] == 1 assert env["ONLY_SECOND"] == 2 + + +@pytest.mark.parametrize( + ("transport", "engine_transport", "socketio"), + [ + ("websocket", "websocket", False), + ("socketio", "websocket", True), + ("polling", "polling", True), + ], +) +def test_set_env_json_keeps_transport_an_engine_io_name( + tmp_path: Path, + mocker: MockerFixture, + transport: str, + engine_transport: str, + socketio: bool, +): + """TRANSPORT stays a name socket.io-client accepts; SOCKETIO picks the protocol. + + Third-party frontend code hands TRANSPORT to socket.io-client to share the + app's Socket.IO connection, as xy's charts do, and "socketio" is no + engine.io transport. + """ + web_dir = _patch_env_json(mocker, tmp_path, transport=transport) + + build.set_env_json() + + env = json.loads((web_dir / "env.json").read_text()) + assert env["TRANSPORT"] == engine_transport + assert env["SOCKETIO"] is socketio diff --git a/tests/units/utils/test_exec.py b/tests/units/utils/test_exec.py index 23ab8c31f42..85ac772618c 100644 --- a/tests/units/utils/test_exec.py +++ b/tests/units/utils/test_exec.py @@ -2,6 +2,7 @@ import builtins import gc +import importlib import logging import multiprocessing import os @@ -19,6 +20,7 @@ from pytest_mock import MockerFixture from reflex_base.environment import environment from reflex_base.utils import serializers +from reflex_base.utils.decorator import once from reflex.utils import exec as exec_utils from reflex.utils import prerequisites, telemetry @@ -538,6 +540,199 @@ def test_arbitrate_ssr_env_var_wins(monkeypatch: pytest.MonkeyPatch): assert exec_utils.arbitrate_ssr(True) is False +@pytest.mark.parametrize("deflate", [True, False]) +def test_run_uvicorn_backend_passes_the_socket_policy( + tmp_path: Path, + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + deflate: bool, +): + """The dev server gets the app's websocket size and compression settings.""" + monkeypatch.setenv("REFLEX_SOCKET_PER_MESSAGE_DEFLATE", str(deflate).lower()) + mocker.patch.object( + exec_utils, + "get_dev_backend_reload_marker", + return_value=tmp_path / exec_utils.DEV_BACKEND_RELOAD_MARKER, + ) + mocker.patch.object(exec_utils, "get_app_instance", return_value="app:app") + mocker.patch.object(exec_utils, "get_reload_paths", return_value=[]) + uvicorn = pytest.importorskip("uvicorn") + run = mocker.patch.object(uvicorn, "run") + + exec_utils.run_uvicorn_backend( + host="0.0.0.0", port=8000, loglevel=exec_utils.LogLevel.INFO + ) + + kwargs = run.call_args.kwargs + assert kwargs["ws_per_message_deflate"] is deflate + assert kwargs["ws_max_size"] == exec_utils._uvicorn_ws_max_size() + + +@pytest.mark.parametrize("deflate", [True, False]) +def test_uvicorn_websocket_args_match_the_options( + monkeypatch: pytest.MonkeyPatch, deflate: bool +): + """Uvicorn's own CLI accepts the args, and reads the same policy from them. + + The Windows production backend passes these on a command line, where a + misspelled option is not a wrong setting but a server that refuses to + start, so they are checked against uvicorn's parser rather than a + hand-written expectation. + """ + pytest.importorskip("uvicorn") + from uvicorn.main import main as uvicorn_cli + + monkeypatch.setenv("REFLEX_SOCKET_PER_MESSAGE_DEFLATE", str(deflate).lower()) + options = exec_utils.uvicorn_websocket_options() + + context = uvicorn_cli.make_context( + "uvicorn", [*exec_utils._uvicorn_websocket_args(), "app:app"] + ) + + assert context.params["ws_per_message_deflate"] is deflate + assert "--no-ws-per-message-deflate" not in exec_utils._uvicorn_websocket_args() + assert context.params["ws_max_size"] == options["ws_max_size"] + + +def test_uvicorn_worker_carries_the_socket_policy(monkeypatch: pytest.MonkeyPatch): + """The gunicorn worker class applies the settings gunicorn cannot pass on.""" + pytest.importorskip("gunicorn") + pytest.importorskip("uvicorn") + monkeypatch.setenv("REFLEX_SOCKET_PER_MESSAGE_DEFLATE", "false") + name = "reflex.utils.uvicorn_worker" + # The class body reads the environment at import time, so import it fresh + # -- then put the cache back, or a later importer would be handed a module + # whose settings were baked from this test's environment. + cached = sys.modules.pop(name, None) + try: + worker = importlib.import_module(name) + finally: + if cached is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = cached + + assert ( + worker.ReflexUvicornWorker.CONFIG_KWARGS.items() + >= exec_utils.uvicorn_websocket_options().items() + ) + assert worker.ReflexUvicornWorker.CONFIG_KWARGS["ws_per_message_deflate"] is False + + +@pytest.fixture +def fresh_uvicorn_warnings(mocker: MockerFixture) -> None: + """Give each test its own `once` caches for the uvicorn warnings. + + They are cached for the life of the process, so a warning another test + already triggered would otherwise silently not be emitted again. + """ + for name in ("_warn_about_uvicorn_websockets", "_warn_user_about_uvicorn"): + warned = getattr(exec_utils, name) + # Re-wrapping would paper over a dropped `once`, so require it first: + # a single run asks `should_use_granian()` several times. + assert hasattr(warned, "__wrapped__"), f"{name} must stay `once`-cached" + mocker.patch.object(exec_utils, name, once(warned.__wrapped__)) + + +def test_auto_detected_uvicorn_is_announced_once( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + caplog, + fresh_uvicorn_warnings: None, +): + """A single run asks repeatedly; the notice belongs in the log once.""" + monkeypatch.delenv("REFLEX_USE_GRANIAN", raising=False) + mocker.patch.object( + exec_utils.importlib.util, "find_spec", side_effect=lambda name: object() + ) + + with caplog.at_level(logging.WARNING): + assert exec_utils.should_use_granian() is False + assert exec_utils.should_use_granian() is False + + assert caplog.text.count("This behavior will change in 0.8.0") == 1 + + +@pytest.mark.parametrize("use_granian", ["0", "1"]) +def test_forcing_uvicorn_warns_about_a_missing_websocket_library( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + caplog, + fresh_uvicorn_warnings: None, + use_granian: str, +): + """Choosing uvicorn explicitly still reports that it cannot serve websockets. + + That choice is the likeliest way to end up without a websocket library, so + the diagnostic cannot live only on the branch that auto-detects uvicorn. + """ + monkeypatch.setenv("REFLEX_USE_GRANIAN", use_granian) + mocker.patch.object( + exec_utils.importlib.util, + "find_spec", + side_effect=lambda name: ( + None if name in ("websockets", "wsproto") else object() + ), + ) + + with caplog.at_level(logging.WARNING): + assert exec_utils.should_use_granian() is (use_granian == "1") + + warned = "has no websocket protocol library" in caplog.text + assert warned is (use_granian == "0") + + +def test_auto_detection_skips_uvicorn_without_a_websocket_library( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + fresh_uvicorn_warnings: None, +): + """Uvicorn and gunicorn alone cannot serve the event socket; granian can. + + python-socketio no longer pulls in wsproto, so an install with uvicorn and + gunicorn but neither websocket library would start an app that never + connects, although the required granian would work. + """ + monkeypatch.delenv("REFLEX_USE_GRANIAN", raising=False) + mocker.patch.object( + exec_utils.importlib.util, + "find_spec", + side_effect=lambda name: ( + None if name in ("websockets", "wsproto") else object() + ), + ) + + assert exec_utils.should_use_granian() is True + + +@pytest.mark.parametrize( + ("configured", "effective"), + [(None, 1000 * 1000), ("50000000", 16 * 1024 * 1024), ("1000", 1000)], +) +def test_granian_caps_the_message_size_at_what_it_can_receive( + monkeypatch: pytest.MonkeyPatch, + caplog, + configured: str | None, + effective: int, +): + """Granian drops frames over 16 MiB, so a larger limit is never advertised. + + Clients size what they send by the limit the handshake advertises; above + what the server receives, an oversized message is lost with the socket. + """ + name = "REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE" + if configured is None: + monkeypatch.delenv(name, raising=False) + else: + monkeypatch.setenv(name, configured) + + with caplog.at_level(logging.WARNING): + exec_utils._cap_message_size_for_granian() + + assert environment.REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE.get() == effective + assert ("drops websocket frames over" in caplog.text) is (configured == "50000000") + + def _fake_granian_prod( mocker: MockerFixture, calls: list[str], telemetry_stopped: bool = True ): diff --git a/uv.lock b/uv.lock index 42feb8519e3..a280823df4c 100644 --- a/uv.lock +++ b/uv.lock @@ -9,10 +9,10 @@ resolution-markers = [ "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version < '3.12' and sys_platform == 'win32'", "python_full_version < '3.12' and sys_platform == 'emscripten'", @@ -39,9 +39,9 @@ reflex-components-core = false reflex-components-dataeditor = false reflex-base = false reflex-components-markdown = false +virtualenv = "2026-10-22T00:00:00Z" reflex-build-sdk = false certifi = "2026-10-21T00:00:00Z" -virtualenv = "2026-10-22T00:00:00Z" reflex-components-internal = false reflex-components-lucide = false reflex-components-radix = false @@ -835,10 +835,10 @@ resolution-markers = [ "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ @@ -1588,6 +1588,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/e5/0ae74c83c1cab2c14daadf28bd0143eb14bb81503834af987b66badc7b61/griffelib-2.3.2-py3-none-any.whl", hash = "sha256:8e710afededd5607f95bf3c8ccc175ee306e7084459faf27b9f9f44ef2a7a274", size = 169420, upload-time = "2026-10-06T09:54:32.038Z" }, ] +[[package]] +name = "gunicorn" +version = "26.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/8a/e4ef6ee11701b6cd64702848415ffb69eeff85cb388a3c6c7fe86f22f3f8/gunicorn-26.2.0.tar.gz", hash = "sha256:62b864895d9ebff0b2f9867ba04fe811c93121596540830c9c916d0769668447", size = 787921, upload-time = "2026-08-24T15:05:59.3Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/85/7522a52e5e2f42faf1a129113ab63e548c42e103e9af395b7bfe65e403e2/gunicorn-26.2.0-py3-none-any.whl", hash = "sha256:bd249d0b3f7972f7432f0a6b6ff3b3ee2d129f70cd1ff6c09a9dd9e29a2b88e3", size = 228389, upload-time = "2026-08-24T15:05:57.67Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -2613,10 +2622,10 @@ resolution-markers = [ "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/95/b0/c7453d0b6e2073c3264468b106ee1563750cecc910965e67357e3698c83e/numpy-2.5.4.tar.gz", hash = "sha256:9a94cf751c9ad8ebaa835bcd3d40dacf8534ad086b88c38029b65123c7999d2a", size = 20866315, upload-time = "2026-10-10T20:05:31.422Z" } @@ -3649,10 +3658,10 @@ resolution-markers = [ "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ @@ -4012,7 +4021,6 @@ dependencies = [ { name = "packaging" }, { name = "psutil", marker = "sys_platform == 'win32'" }, { name = "python-multipart" }, - { name = "python-socketio" }, { name = "redis" }, { name = "reflex-base" }, { name = "reflex-components-code" }, @@ -4044,10 +4052,17 @@ db = [ pydantic = [ { name = "reflex-base", extra = ["pydantic"] }, ] +socketio = [ + { name = "python-socketio" }, +] testing = [ { name = "psutil" }, { name = "selenium" }, +] +uvicorn = [ + { name = "gunicorn" }, { name = "uvicorn" }, + { name = "websockets" }, ] [package.dev-dependencies] @@ -4089,6 +4104,7 @@ dev = [ { name = "pytest-rerunfailures" }, { name = "pytest-split" }, { name = "python-dotenv" }, + { name = "python-socketio" }, { name = "pyyaml" }, { name = "reflex-bench" }, { name = "reflex-build-sdk", extra = ["aiohttp", "httpx", "httpx2"] }, @@ -4113,15 +4129,16 @@ dev = [ requires-dist = [ { name = "alembic", marker = "extra == 'db'", specifier = ">=1.15.2,<2.0" }, { name = "click", specifier = ">=8.2" }, - { name = "granian", extras = ["reload"], specifier = ">=2.7.4" }, + { name = "granian", extras = ["reload"], specifier = ">=2.8.1" }, { name = "greenlet", marker = "extra == 'db'", specifier = ">=3.3" }, + { name = "gunicorn", marker = "extra == 'uvicorn'", specifier = ">=23.0" }, { name = "httpx", specifier = ">=0.26,<1.0" }, { name = "packaging", specifier = ">=24.2,<27" }, { name = "psutil", marker = "sys_platform == 'win32'", specifier = ">=7.0.0,<8.0" }, { name = "psutil", marker = "extra == 'testing'", specifier = ">=7.0.0,<8.0" }, { name = "pydantic", marker = "extra == 'db'", specifier = ">=2.12.0,<3.0" }, { name = "python-multipart", specifier = ">=0.0.32,<1.0" }, - { name = "python-socketio", specifier = ">=5.12.0,<6.0" }, + { name = "python-socketio", marker = "extra == 'socketio'", specifier = ">=5.12.0,<6.0" }, { name = "redis", specifier = ">=6.4,<8.0" }, { name = "reflex-base", editable = "packages/reflex-base" }, { name = "reflex-base", extras = ["pydantic"], marker = "extra == 'pydantic'", editable = "packages/reflex-base" }, @@ -4143,10 +4160,11 @@ requires-dist = [ { name = "sqlmodel", marker = "extra == 'db'", specifier = ">=0.0.24" }, { name = "starlette", specifier = ">=1.3.1" }, { name = "typing-extensions", specifier = ">=4.13.0" }, - { name = "uvicorn", marker = "extra == 'testing'", specifier = ">=0.34.0,<1.0" }, + { name = "uvicorn", marker = "extra == 'uvicorn'", specifier = ">=0.20.0" }, + { name = "websockets", marker = "extra == 'uvicorn'", specifier = ">=13.0" }, { name = "wrapt", specifier = ">=1.17.0,<2.6" }, ] -provides-extras = ["db", "pydantic", "testing"] +provides-extras = ["db", "pydantic", "socketio", "testing", "uvicorn"] [package.metadata.requires-dev] bench = [ @@ -4186,6 +4204,7 @@ dev = [ { name = "pytest-rerunfailures" }, { name = "pytest-split" }, { name = "python-dotenv" }, + { name = "python-socketio" }, { name = "pyyaml" }, { name = "reflex-bench", editable = "packages/reflex-bench" }, { name = "reflex-build-sdk", extras = ["aiohttp", "httpx", "httpx2"], editable = "packages/reflex-build-sdk" }, @@ -4497,7 +4516,7 @@ dependencies = [ { name = "psycopg" }, { name = "psycopg", extra = ["binary"], marker = "python_full_version < '3.15'" }, { name = "python-frontmatter" }, - { name = "reflex" }, + { name = "reflex", extra = ["socketio"] }, { name = "reflex-components-internal" }, { name = "reflex-docgen" }, { name = "reflex-enterprise" }, @@ -4536,7 +4555,7 @@ requires-dist = [ { name = "psycopg", marker = "python_full_version >= '3.15'" }, { name = "psycopg", extras = ["binary"], marker = "python_full_version < '3.15'" }, { name = "python-frontmatter" }, - { name = "reflex", editable = "." }, + { name = "reflex", extras = ["socketio"], editable = "." }, { name = "reflex-components-internal", editable = "packages/reflex-components-internal" }, { name = "reflex-docgen", editable = "packages/reflex-docgen" }, { name = "reflex-enterprise", specifier = ">=0.9.4" }, @@ -4873,10 +4892,10 @@ resolution-markers = [ "python_full_version == '3.14.*' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [