fixed buffer upload by shortcircuiting

Author: FarhanAliRaza

packages/reflex-base/src/reflex_base/event/processor/event_processor.py

@@ -380,6 +380,7 @@ async def enqueue_stream_delta(
         self,
         token: str,
         event: Event,
+        on_task_future: Callable[[EventFuture], None] | None = None,
     ) -> AsyncGenerator[Mapping[str, Any]]:
         """Enqueue an event to be processed and yield deltas emitted by the event handler.
 
@@ -393,6 +394,8 @@ async def enqueue_stream_delta(
         Args:
             token: The client token associated with the event.
             event: The event to be enqueued.
+            on_task_future: Optional callback invoked with the EventFuture for the
+                enqueued handler as soon as it is created.
 
         Yields:
             Deltas emitted by the event handler for the specified token.
@@ -425,6 +428,8 @@ async def _emit_delta_impl(
                 emit_delta_impl=_emit_delta_impl,
             ),
         )
+        if on_task_future is not None:
+            on_task_future(task_future)
         all_task_futures = asyncio.create_task(task_future.wait_all())
         waiting_for = {all_task_futures, asyncio.create_task(deltas.get())}
         try:

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

@@ -22,7 +22,8 @@
 from typing_extensions import Self
 
 if TYPE_CHECKING:
-    from reflex_base.utils.types import Receive, Scope, Send
+    from reflex_base.event.processor import EventFuture
+    from reflex_base.utils.types import Message, Receive, Scope, Send
 
     from reflex.app import App
 
@@ -403,20 +404,70 @@ class _UploadStreamingResponse(StreamingResponse):
     """Streaming response that always releases upload form resources."""
 
     _on_finish: Callable[[], Awaitable[None]]
+    _on_disconnect: Callable[[], None] | None
+    _disconnect_handled: bool
 
     def __init__(
         self,
         *args: Any,
         on_finish: Callable[[], Awaitable[None]],
+        on_disconnect: Callable[[], None] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, **kwargs)
         self._on_finish = on_finish
+        self._on_disconnect = on_disconnect
+        self._disconnect_handled = False
+
+    def _handle_disconnect(self) -> None:
+        """Run disconnect cleanup exactly once."""
+        if self._disconnect_handled or self._on_disconnect is None:
+            return
+        self._disconnect_handled = True
+        self._on_disconnect()
+
+    async def _watch_disconnect(self, receive: Receive) -> None:
+        """Wait for the client connection to close."""
+        while True:
+            message = await receive()
+            if message["type"] == "http.disconnect":
+                self._handle_disconnect()
+                return
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        spec_version = tuple(
+            map(int, scope.get("asgi", {}).get("spec_version", "2.0").split("."))
+        )
+        disconnect_task: asyncio.Task[None] | None = None
+        use_watcher = spec_version >= (2, 4) and self._on_disconnect is not None
+
+        async def wrapped_receive() -> Message:
+            message = await receive()
+            if message.get("type") == "http.disconnect":
+                self._handle_disconnect()
+            return message
+
         try:
-            await super().__call__(scope, receive, send)
+            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.
+                disconnect_task = asyncio.create_task(self._watch_disconnect(receive))
+            try:
+                await super().__call__(
+                    scope,
+                    wrapped_receive if not use_watcher else receive,
+                    send,
+                )
+            except ClientDisconnect:
+                self._handle_disconnect()
+                raise
         finally:
+            if disconnect_task is not None:
+                disconnect_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await disconnect_task
             await self._on_finish()
 
 
@@ -515,20 +566,46 @@ def _create_upload_event() -> Event:
         msg = "Upload event was not created."
         raise RuntimeError(msg)
 
+    task_future: EventFuture | None = None
+    disconnect_seen = False
+
+    def _try_cancel() -> None:
+        """Cancel the task future if it exists and is still running."""
+        if task_future is not None and not task_future.done():
+            task_future.cancel()
+
+    def _remember_task_future(future: EventFuture) -> None:
+        """Keep a handle to the upload task for disconnect cancellation."""
+        nonlocal task_future
+        task_future = future
+        if disconnect_seen:
+            _try_cancel()
+
+    def _cancel_upload_task() -> None:
+        """Cancel the queued upload handler when the client disconnects."""
+        nonlocal disconnect_seen
+        disconnect_seen = True
+        _try_cancel()
+
     async def _ndjson_updates():
         """Process the upload event, generating ndjson updates.
 
         Yields:
             Each state update as newline-delimited JSON.
         """
         # 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, event):
+        async for delta in app.event_processor.enqueue_stream_delta(
+            token,
+            event,
+            on_task_future=_remember_task_future,
+        ):
             yield json_dumps(StateUpdate(delta=delta)) + "\n"
 
     return _UploadStreamingResponse(
         _ndjson_updates(),
         media_type="application/x-ndjson",
         on_finish=_close_form_data,
+        on_disconnect=_cancel_upload_task,
     )
 
 

tests/units/reflex_base/event/processor/test_event_processor.py

@@ -518,6 +518,30 @@ async def test_stream_delta_not_configured_raises():
             pass
 
 
+async def test_stream_delta_calls_on_task_future(token: str):
+    """enqueue_stream_delta exposes the tracked EventFuture immediately.
+
+    Args:
+        token: The client token.
+    """
+    ep = EventProcessor(graceful_shutdown_timeout=2)
+    ep.configure()
+    captured = []
+    async with ep:
+        event = Event.from_event_type(noop_event())[0]
+        collected = [
+            d
+            async for d in ep.enqueue_stream_delta(
+                token,
+                event,
+                on_task_future=captured.append,
+            )
+        ]
+    assert collected == []
+    assert len(captured) == 1
+    assert captured[0].done()
+
+
 async def test_sequential_chained_events_run_in_order(token: str):
     """Chained events enqueued by a handler run in the order they were enqueued.
 

tests/units/test_app.py

@@ -1306,6 +1306,79 @@ async def send(_message):
     assert bio.closed
 
 
+@pytest.mark.asyncio
+async def test_upload_file_cancels_buffered_handler_on_disconnect_before_future_capture(
+    token: str,
+):
+    """Buffered uploads cancel the handler even if disconnect wins the race.
+
+    This exercises the ASGI 2.4 path where the response must watch
+    ``receive()`` directly because Starlette does not listen for disconnects
+    while streaming the response body.
+
+    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
+
+    cancelled = asyncio.Event()
+    task_future = Mock()
+    task_future.done = Mock(side_effect=cancelled.is_set)
+    task_future.cancel = Mock(side_effect=cancelled.set)
+
+    async def enqueue_stream_delta(_token, _event, on_task_future=None):
+        assert on_task_future is not None
+        on_task_future(task_future)
+        await cancelled.wait()
+        if False:  # pragma: no cover
+            yield {}
+
+    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):  # noqa: RUF029
+        return None
+
+    await asyncio.wait_for(
+        streaming_response(
+            {"type": "http", "asgi": {"spec_version": "2.4"}},
+            receive,
+            send,
+        ),
+        timeout=1,
+    )
+
+    assert task_future.cancel.call_count == 1
+    assert form_close.await_count == 1
+    assert bio.closed
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "state",