Skip buffered upload handler when probe chunk detects client disconnect

Insert a b"\n" probe chunk before the real response body so that an early
client disconnect is detected before the upload handler is enqueued. Add
asyncio.sleep(0) yield point in _upload_buffered_file to let the
disconnect watcher fire first, and return early if disconnect was seen.
Includes a test covering the probe-based disconnect detection path.

Author: FarhanAliRaza

packages/reflex-components-core/src/reflex_components_core/core/_upload.py

@@ -441,18 +441,30 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         disconnect_task: asyncio.Task[None] | None = None
         use_watcher = spec_version >= (2, 4) and self._on_disconnect is not None
 
+        if use_watcher:
+            body_iterator = self.body_iterator
+
+            async def body_with_probe() -> AsyncGenerator[
+                str | bytes | memoryview[int], None
+            ]:
+                """Yield a tiny probe chunk before the real response body."""
+                yield b"\n"
+                async for chunk in body_iterator:
+                    yield chunk
+
+            self.body_iterator = body_with_probe()
+
         async def wrapped_receive() -> Message:
             message = await receive()
-            if message.get("type") == "http.disconnect":
+            if message["type"] == "http.disconnect":
                 self._handle_disconnect()
             return message
 
         try:
             if use_watcher:
-                # ASGI >= 2.4: use a dedicated task to watch for disconnect
-                # concurrently. Pass raw `receive` to Starlette — the watcher
-                # owns disconnect detection; using wrapped_receive here would
-                # race on the same receive callable.
+                # ASGI >= 2.4: Starlette does not call receive() while
+                # streaming. Use a dedicated task so disconnect fires the
+                # callback; pass raw receive to avoid racing wrapped_receive.
                 disconnect_task = asyncio.create_task(self._watch_disconnect(receive))
             try:
                 await super().__call__(
@@ -593,6 +605,10 @@ async def _ndjson_updates():
         Yields:
             Each state update as newline-delimited JSON.
         """
+        # Let the disconnect watcher run before we enqueue the upload handler.
+        await asyncio.sleep(0)
+        if disconnect_seen:
+            return
         # Enqueue the task on the main event loop, but emit deltas to the local queue.
         async for delta in app.event_processor.enqueue_stream_delta(
             token,

tests/units/test_app.py

@@ -30,6 +30,7 @@
 from reflex_components_radix.themes.typography.text import Text
 from starlette.applications import Starlette
 from starlette.datastructures import FormData, Headers, UploadFile
+from starlette.requests import ClientDisconnect
 from starlette.responses import StreamingResponse
 from starlette_admin.auth import AuthProvider
 
@@ -1379,6 +1380,74 @@ async def send(_message):  # noqa: RUF029
     assert bio.closed
 
 
+@pytest.mark.asyncio
+async def test_upload_file_skips_buffered_handler_when_disconnect_detected_on_probe(
+    token: str,
+):
+    """Buffered uploads skip handler dispatch when the probe send disconnects.
+
+    This models ASGI 2.4+ behavior where the upload request can finish parsing,
+    but the client disconnect is only surfaced on the first response-body send.
+
+    Args:
+        token: A token.
+    """
+    request_mock = unittest.mock.Mock()
+    request_mock.headers = {
+        "reflex-client-token": token,
+        "reflex-event-handler": f"{FileUploadState.get_full_name()}.multi_handle_upload",
+    }
+
+    bio = io.BytesIO(b"contents of image one")
+    file1 = UploadFile(filename="image1.jpg", file=bio)
+    form_data = FormData([("files", file1)])
+    original_close = form_data.close
+    form_close = AsyncMock(side_effect=original_close)
+    form_data.close = form_close
+
+    async def form():  # noqa: RUF029
+        return form_data
+
+    request_mock.form = form
+
+    msg = "upload handler should not be enqueued"
+    probe_chunk = b"\n"
+    asgi_24_scope = {"type": "http", "asgi": {"spec_version": "2.4"}}
+    enqueue_stream_delta = Mock(side_effect=AssertionError(msg))
+    app = Mock(
+        event_processor=Mock(enqueue_stream_delta=enqueue_stream_delta),
+    )
+
+    upload_fn = upload(app)
+    streaming_response = await upload_fn(request_mock)
+
+    assert isinstance(streaming_response, StreamingResponse)
+
+    async def receive():
+        await asyncio.sleep(0)
+        return {"type": "http.disconnect"}
+
+    async def send(message):
+        await asyncio.sleep(0)
+        if (
+            message.get("type") == "http.response.body"
+            and message.get("body") == probe_chunk
+        ):
+            err = "client disconnected"
+            raise OSError(err)
+
+    with pytest.raises(ClientDisconnect):
+        await streaming_response(
+            asgi_24_scope,
+            receive,
+            send,
+        )
+
+    assert enqueue_stream_delta.call_count == 0
+    assert form_close.await_count == 1
+    assert bio.closed
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "state",