Fix buffered upload handler not cancelling on client disconnect (#6307)

* Fix buffered upload handler not cancelling on client disconnect

When a client disconnects during a buffered file upload, the enqueued
event handler was not being cancelled, leaving orphaned tasks. Add
disconnect detection to _UploadStreamingResponse (via ASGI 2.4 receive
watcher and ClientDisconnect handling) and propagate an on_task_future
callback through enqueue_stream_delta so the upload path can cancel the
EventFuture when the connection drops.

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

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

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

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Generalize disconnect-cancel for streaming event responses

Move disconnect detection out of upload-specific code into a reusable
DisconnectAwareStreamingResponse in reflex-base. Drop the on_task_future
callback from enqueue_stream_delta — the generator"s finally block already
cancels the task future when the body iterator is closed on disconnect.

* Add pre-2.4 streaming response disconnect test, remove bogus event processor tests

Add test coverage for the pre-2.4 ASGI disconnect path in
DisconnectAwareStreamingResponse.

Remove test_stream_delta_aclose_cancels_in_flight_event and
test_stream_delta_consumer_task_cancellation_cancels_in_flight_event — they
passed with or without the task_future.cancel() fix because _on_future_done
already cascades cancellation from the EventFuture to the handler task.

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Author: FarhanAliRaza

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

@@ -390,6 +390,10 @@ async def enqueue_stream_delta(
         Any frontend events or chained events are handled normally and deltas from chained events
         will not be yielded by this method.
 
+        If the consumer stops iterating early, the in-flight event future is
+        cancelled so the handler chain does not continue running in the
+        background.
+
         Args:
             token: The client token associated with the event.
             event: The event to be enqueued.
@@ -442,6 +446,9 @@ async def _emit_delta_impl(
         finally:
             for future in waiting_for:
                 future.cancel()
+            # Cancel the event chain if the streaming consumer exits early.
+            if not task_future.done():
+                task_future.cancel()
         # Raise any exceptions for the caller, waiting for all chained events.
         await task_future.wait_all()
 

packages/reflex-base/src/reflex_base/utils/streaming_response.py

@@ -0,0 +1,157 @@
+"""Internal helpers for streaming responses."""
+
+from __future__ import annotations
+
+import asyncio
+import builtins
+import contextlib
+import sys
+from collections.abc import Awaitable, Callable, Generator
+from functools import partial
+from typing import Any
+
+import anyio
+from starlette.requests import ClientDisconnect
+from starlette.responses import StreamingResponse
+
+from reflex_base.utils.types import Receive, Scope, Send
+
+_BASE_EXCEPTION_GROUP = getattr(builtins, "BaseExceptionGroup", None)
+
+
+def _parse_asgi_spec_version(scope: Scope) -> tuple[int, ...]:
+    """Parse the ASGI spec version from a scope.
+
+    Args:
+        scope: The ASGI scope.
+
+    Returns:
+        The parsed ASGI spec version, or ``(2, 0)`` if parsing fails.
+    """
+    raw_spec = scope.get("asgi", {}).get("spec_version", "2.0")
+    try:
+        return tuple(int(part) for part in str(raw_spec).split("."))
+    except (TypeError, ValueError):
+        return (2, 0)
+
+
+@contextlib.contextmanager
+def _collapse_excgroups() -> Generator[None, None, None]:
+    """Collapse single-item exception groups to their underlying exception."""
+    collapsed_exc: BaseException | None = None
+    try:
+        yield
+    except BaseException as exc:
+        collapsed_exc = exc
+        if sys.version_info >= (3, 11) and _BASE_EXCEPTION_GROUP is not None:
+            while isinstance(collapsed_exc, _BASE_EXCEPTION_GROUP):
+                nested_exceptions = getattr(collapsed_exc, "exceptions", None)
+                if (
+                    not isinstance(nested_exceptions, tuple)
+                    or len(nested_exceptions) != 1
+                    or not isinstance(nested_exceptions[0], BaseException)
+                ):
+                    break
+                collapsed_exc = nested_exceptions[0]
+        if collapsed_exc is exc:
+            raise
+    if collapsed_exc is not None:
+        raise collapsed_exc
+
+
+class DisconnectAwareStreamingResponse(StreamingResponse):
+    """Streaming response that cancels its body task on disconnect."""
+
+    _on_finish: Callable[[], Awaitable[None]]
+
+    def __init__(
+        self,
+        *args: Any,
+        on_finish: Callable[[], Awaitable[None]],
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the response.
+
+        Args:
+            args: Positional args forwarded to ``StreamingResponse``.
+            on_finish: Cleanup callback to run exactly once when the response ends.
+            kwargs: Keyword args forwarded to ``StreamingResponse``.
+        """
+        super().__init__(*args, **kwargs)
+        self._on_finish = on_finish
+
+    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":
+                return
+
+    async def _close_body_iterator(self) -> None:
+        """Close the body iterator if it supports ``aclose``."""
+        aclose = getattr(self.body_iterator, "aclose", None)
+        if aclose is not None:
+            await aclose()
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Serve the response and cancel the body task on disconnect."""
+        spec_version = _parse_asgi_spec_version(scope)
+
+        try:
+            if spec_version < (2, 4):
+                with _collapse_excgroups():
+                    async with anyio.create_task_group() as task_group:
+
+                        async def wrap(func: Callable[[], Awaitable[None]]) -> None:
+                            await func()
+                            task_group.cancel_scope.cancel()
+
+                        task_group.start_soon(wrap, partial(self.stream_response, send))
+                        await wrap(partial(self.listen_for_disconnect, receive))
+            else:
+                # Verified against Starlette 0.52.1: the ASGI >= 2.4 path in
+                # StreamingResponse.__call__ delegates straight to
+                # stream_response(send) and does not read from receive().
+                # Keep calling stream_response(send) directly here so the
+                # disconnect watcher remains the only receive() consumer; if
+                # Starlette changes that contract, re-check this logic.
+                stream_task = asyncio.create_task(self.stream_response(send))
+                disconnect_task = asyncio.create_task(self._watch_disconnect(receive))
+                should_close_body_iterator = False
+
+                try:
+                    done, _ = await asyncio.wait(
+                        {stream_task, disconnect_task},
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                    if disconnect_task in done and not stream_task.done():
+                        should_close_body_iterator = True
+                        stream_task.cancel()
+                        with contextlib.suppress(asyncio.CancelledError):
+                            await stream_task
+                    else:
+                        try:
+                            await stream_task
+                        except OSError as err:
+                            should_close_body_iterator = True
+                            raise ClientDisconnect from err
+                finally:
+                    if not disconnect_task.done():
+                        disconnect_task.cancel()
+                        with contextlib.suppress(asyncio.CancelledError):
+                            await disconnect_task
+                    if not stream_task.done():
+                        should_close_body_iterator = True
+                        stream_task.cancel()
+                        with contextlib.suppress(asyncio.CancelledError):
+                            await stream_task
+                    if should_close_body_iterator:
+                        await self._close_body_iterator()
+        finally:
+            await self._on_finish()
+
+        if self.background is not None:
+            await self.background()
+
+
+__all__ = ["DisconnectAwareStreamingResponse"]

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

@@ -6,13 +6,14 @@
 import contextlib
 import dataclasses
 from collections import deque
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncGenerator, AsyncIterator
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, BinaryIO, cast
 
 from python_multipart.multipart import MultipartParser, parse_options_header
 from reflex_base.utils import exceptions
 from reflex_base.utils.format import json_dumps
+from reflex_base.utils.streaming_response import DisconnectAwareStreamingResponse
 from starlette.datastructures import Headers
 from starlette.datastructures import UploadFile as StarletteUploadFile
 from starlette.exceptions import HTTPException
@@ -22,8 +23,6 @@
 from typing_extensions import Self
 
 if TYPE_CHECKING:
-    from reflex_base.utils.types import Receive, Scope, Send
-
     from reflex.app import App
 
 
@@ -399,27 +398,6 @@ async def parse(self) -> None:
         await self._flush_emitted_chunks()
 
 
-class _UploadStreamingResponse(StreamingResponse):
-    """Streaming response that always releases upload form resources."""
-
-    _on_finish: Callable[[], Awaitable[None]]
-
-    def __init__(
-        self,
-        *args: Any,
-        on_finish: Callable[[], Awaitable[None]],
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(*args, **kwargs)
-        self._on_finish = on_finish
-
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        try:
-            await super().__call__(scope, receive, send)
-        finally:
-            await self._on_finish()
-
-
 def _require_upload_headers(request: Request) -> tuple[str, str]:
     """Extract the required upload headers from a request.
 
@@ -525,7 +503,7 @@ async def _ndjson_updates():
         async for delta in app.event_processor.enqueue_stream_delta(token, event):
             yield json_dumps(StateUpdate(delta=delta)) + "\n"
 
-    return _UploadStreamingResponse(
+    return DisconnectAwareStreamingResponse(
         _ndjson_updates(),
         media_type="application/x-ndjson",
         on_finish=_close_form_data,