Stop canceled uploads from finishing

Author: Gautam Manchandani

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

@@ -439,6 +439,9 @@ async def _emit_delta_impl(
                             yield await result
                             waiting_for.add(asyncio.create_task(deltas.get()))
                         break
+        except (asyncio.CancelledError, GeneratorExit):
+            task_future.cancel()
+            raise
         finally:
             for future in waiting_for:
                 future.cancel()

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

@@ -64,6 +64,14 @@ async def _multi_delta_handler():
         await asyncio.sleep(0.01)
 
 
+async def _delta_then_slow_logging_handler():
+    """Emit an initial delta, then block long enough for cancellation to matter."""
+    ctx = EventContext.get()
+    await ctx.emit_delta({"state": {"x": 1}})
+    await asyncio.sleep(0.5)
+    _CALL_LOG.append({"value": "completed_after_stream_close"})
+
+
 async def _slow_logging_handler(value: str = "default"):
     """A slow logging handler that pauses before recording.
 
@@ -117,6 +125,7 @@ async def _background_slow_logging_handler(value: str = "default"):
 chaining_event = EventHandler(fn=_chaining_handler)
 delta_event = EventHandler(fn=_delta_handler)
 multi_delta_event = EventHandler(fn=_multi_delta_handler)
+delta_then_slow_logging_event = EventHandler(fn=_delta_then_slow_logging_handler)
 slow_logging_event = EventHandler(fn=_slow_logging_handler)
 multi_chaining_event = EventHandler(fn=_multi_chaining_handler)
 background_slow_logging_event = EventHandler(fn=_background_slow_logging_handler)
@@ -140,6 +149,7 @@ def _register_handlers(forked_registration_context: RegistrationContext):
         chaining_event,
         delta_event,
         multi_delta_event,
+        delta_then_slow_logging_event,
         slow_logging_event,
         multi_chaining_event,
         background_slow_logging_event,
@@ -510,6 +520,20 @@ async def test_stream_delta_noop_handler_yields_nothing(token: str):
     assert collected == []
 
 
+async def test_stream_delta_cancels_event_when_generator_closes_early(token: str):
+    """Closing enqueue_stream_delta early cancels the in-flight event."""
+    ep = EventProcessor(graceful_shutdown_timeout=2)
+    ep.configure()
+    async with ep:
+        event = Event.from_event_type(delta_then_slow_logging_event())[0]
+        stream = ep.enqueue_stream_delta(token, event)
+        assert await anext(stream) == {"state": {"x": 1}}
+        await stream.aclose()
+        await asyncio.sleep(0.1)
+
+    assert _CALL_LOG == []
+
+
 async def test_stream_delta_not_configured_raises():
     """enqueue_stream_delta raises RuntimeError if processor is not configured."""
     ep = EventProcessor()