From 3c5811f6db3cf99d550e03fc5795cfdb18461ff2 Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Mon, 2 Nov 2026 23:13:00 -0800 Subject: [PATCH 1/8] fix(core): propagate early close through async streaming actions Closing Application.astream_iterate() (or abandoning an AsyncStreamingResultContainer) mid-action left the action's generator chain suspended until GC / loop.shutdown_asyncgens(), so resources held by the action (HTTP/SSE streams, cursors, spans) were released late, in an arbitrary context, with cleanup errors only logged. - Add AsyncStreamingResultContainer.aclose(): closes the action chain without draining it, runs the post-step callback once, idempotent. - astream_iterate closes the in-flight container when it is closed. - Every async forwarding layer (_arun_*_streaming_action, FunctionBasedStreamingAction, type_eraser, pydantic async_action_generator) now closes the generator it wraps. --- burr/core/action.py | 48 +++++++++++++----- burr/core/application.py | 13 ++++- burr/integrations/pydantic.py | 14 +++--- tests/core/test_application.py | 90 ++++++++++++++++++++++++++++++++++ 4 files changed, 147 insertions(+), 18 deletions(-) diff --git a/burr/core/action.py b/burr/core/action.py index b71940d47..240f0808b 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -18,6 +18,7 @@ import abc import ast import builtins +import contextlib import copy import inspect import sys @@ -152,8 +153,9 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @wraps(func) async def async_gen_wrapper(*args: Any, **kwargs: Any) -> Any: - async for item in func(*args, **kwargs): - yield item + async with contextlib.aclosing(func(*args, **kwargs)) as gen: + async for item in gen: + yield item return async_gen_wrapper @@ -1582,6 +1584,33 @@ def __init__( self._callback = callback self._result = None self._callback_realized = False + self._iterator: Optional[AsyncGenerator[StreamResultType, None]] = None + + async def _realize_callback(self, exc: Optional[BaseException]) -> None: + if self._result is None: + self._result = None, self._initial_state + if not self._callback_realized: + self._callback_realized = True + await self._callback(*self._result, exc) + + async def aclose(self) -> None: + """Stops the stream early without draining it. + + Closes the action's generator chain so its cleanup (``finally`` blocks, ``async with`` + exits) runs now rather than at garbage collection, then runs the post-step callback + exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a + no-op once the stream has completed. Exceptions raised by the action's cleanup propagate. + """ + exc: Optional[BaseException] = None + try: + await self.streaming_result_generator.aclose() + except BaseException as e: + exc = e + raise + finally: + await self._realize_callback(exc) + if self._iterator is not None: + await self._iterator.aclose() async def __anext__(self) -> StreamResultType: """Moves to the next state in the streaming result""" @@ -1606,15 +1635,11 @@ async def gen_fn(): except StopAsyncIteration: return finally: - if self._result is None: - self._result = None, self._initial_state - if not self._callback_realized: - exc = sys.exc_info()[1] - self._callback_realized = True - await self._callback(*self._result, exc) + await self._realize_callback(sys.exc_info()[1]) # return it as `__aiter__` cannot be async/have awaits :/ - return gen_fn() + self._iterator = gen_fn() + return self._iterator async def get(self) -> tuple[Optional[StreamResultType], State[StateType]]: # exhaust the generator @@ -1741,8 +1766,9 @@ def __init__( async def _a_stream_run_and_update( self, state: State, **run_kwargs ) -> AsyncGeneratorReturnType: - async for result in self._fn(state, **self._bound_params, **run_kwargs): - yield result + async with contextlib.aclosing(self._fn(state, **self._bound_params, **run_kwargs)) as gen: + async for result in gen: + yield result def _stream_run_and_update(self, state: State, **run_kwargs) -> GeneratorReturnType: yield from self._fn(state, **self._bound_params, **run_kwargs) diff --git a/burr/core/application.py b/burr/core/application.py index 91ef8b1b6..48ce3c9d2 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -456,6 +456,9 @@ async def _arun_single_step_streaming_action( e, exc_info=True, ) + finally: + # Unwind the action now on early close instead of leaving it to GC. + await generator.aclose() if state_update is None: raise ValueError( @@ -585,6 +588,9 @@ async def _arun_multi_step_streaming_action( e, exc_info=True, ) + finally: + # Unwind the action now on early close instead of leaving it to GC. + await generator.aclose() state_update = _run_reducer(action, state, result, action.name) _validate_result(result, action.name, action.schema) _validate_reducer_writes(action, state_update, action.name) @@ -2012,7 +2018,12 @@ async def astream_iterate( _, streaming_result = await self.astream_result( # Use astream_result halt_after=[next_action.name], halt_before=None, inputs=inputs ) - yield next_action, streaming_result + try: + yield next_action, streaming_result + except GeneratorExit: + # The caller closed us mid-action: unwind it rather than leaving it to GC. + await streaming_result.aclose() + raise # We need to ensure it's fully exhausted before going to the next action await streaming_result.get() # await the get call if self._should_halt_iterate(halt_before, halt_after, next_action): diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 54c175b81..08533c837 100644 --- a/burr/integrations/pydantic.py +++ b/burr/integrations/pydantic.py @@ -17,6 +17,7 @@ from __future__ import annotations +import contextlib import copy import inspect import types @@ -357,12 +358,13 @@ async def async_action_generator( state: State, **kwargs ) -> AsyncGenerator[tuple[dict, Optional[State]], None]: model_to_use = model_from_state(model=SubsetInputType, state=state) - async for partial, state_update in fn(state=model_to_use, **kwargs): - if state_update is None: - yield partial, None - else: - output = SubsetOutputType(**model_to_dict(state_update, include=writes)) - yield partial, merge_to_state(model=output, write_keys=writes, state=state) + async with contextlib.aclosing(fn(state=model_to_use, **kwargs)) as gen: + async for partial, state_update in gen: + if state_update is None: + yield partial, None + else: + output = SubsetOutputType(**model_to_dict(state_update, include=writes)) + yield partial, merge_to_state(model=output, write_keys=writes, state=state) is_async = inspect.isasyncgenfunction(fn) # This recreates the @streaming_action decorator diff --git a/tests/core/test_application.py b/tests/core/test_application.py index efb7d205b..1e1cb9149 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -17,6 +17,7 @@ import asyncio import collections +import contextlib import datetime import logging import typing @@ -31,6 +32,7 @@ Action, AsyncGenerator, AsyncStreamingAction, + AsyncStreamingResultContainer, Condition, Reducer, Result, @@ -40,6 +42,7 @@ action, default, expr, + streaming_action, ) from burr.core.application import ( PRIOR_STEP, @@ -2770,6 +2773,93 @@ async def test_astream_iterate(exhaust_intermediate_generators: bool): assert len(stream_event_tracker.post_stream_item_calls) == 20 +def _closing_probe_app(events: list, *hooks) -> Application: + """One streaming action whose nested producer has an awaiting ``finally`` (like an HTTP close).""" + + async def producer(): + try: + yield "a" + yield "b" + finally: + await asyncio.sleep(0) + events.append("producer_closed") + + @streaming_action(reads=[], writes=["x"]) + async def act(state: State): + async with contextlib.aclosing(producer()) as items: + async for item in items: + yield {"item": item}, None + events.append("advanced") + yield {"item": "done"}, state.update(x=1) + + return ( + ApplicationBuilder() + .with_actions(act) + .with_transitions() + .with_entrypoint("act") + .with_state(State()) + .with_hooks(*hooks) + .build() + ) + + +@pytest.mark.parametrize("consume_first_item", [True, False]) +async def test_astream_iterate_aclose_unwinds_in_flight_action(consume_first_item: bool): + """Closing astream_iterate mid-action runs the action's cleanup now (not at GC), + does not drain or advance it, and fires post_end_stream exactly once.""" + events = [] + stream_event_tracker = StreamEventCaptureTracker() + app = _closing_probe_app(events, stream_event_tracker) + steps = app.astream_iterate(halt_after=["act"]) + async for _, container in steps: + if consume_first_item: + async for _ in container: + break + break + await steps.aclose() + assert events == (["producer_closed"] if consume_first_item else []) + assert len(stream_event_tracker.post_end_stream_calls) == 1 + + +async def test_streaming_result_container_aclose_is_idempotent_and_get_returns_initial_state(): + events = [] + app = _closing_probe_app(events) + _, container = await app.astream_result(halt_after=["act"]) + async for _ in container: + break + await container.aclose() + await container.aclose() + result, state = await container.get() + assert events == ["producer_closed"] + assert result is None + assert "x" not in state + + +async def test_streaming_result_container_aclose_propagates_cleanup_errors(): + class CleanupError(Exception): + pass + + async def failing_cleanup(): + try: + yield {"item": "a"}, None + finally: + raise CleanupError("cleanup failed") + + seen = [] + + async def callback(result, state, exc): + seen.append(exc) + + container = AsyncStreamingResultContainer( + failing_cleanup(), State(), lambda r, s: (r, s), callback + ) + async for _ in container: + break + with pytest.raises(CleanupError): + await container.aclose() + assert len(seen) == 1 and isinstance(seen[0], CleanupError) + + async def test_astream_result_halt_after_run_through_streaming(): action_tracker = CallCaptureTracker() stream_event_tracker = StreamEventCaptureTrackerAsync() From ea69ad2fe0cd59798b9159af76378e094fd0c196 Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Tue, 3 Nov 2026 07:39:25 -0800 Subject: [PATCH 2/8] fix(core): report early-closed async streams as failed steps AsyncStreamingResultContainer.aclose() passed exception=None to the post-step callback when the stream was closed before completing, so post_run_step hooks saw an interrupted step as a success and persisters saved it as "completed" without its writes; resuming then skipped the action. Report a close-before-completion with GeneratorExit, matching the previous GC-finalizer behavior. - widen the callback's exception type to Optional[BaseException] - tests: hook exception on aclose and on task cancellation, persister status + resume, aclose after completion still reports success, multi-step AsyncStreamingAction and pydantic streaming paths - docs: "Stopping a stream early" section and astream_result docstring --- burr/core/action.py | 8 +- burr/core/application.py | 3 +- docs/concepts/streaming-actions.rst | 32 ++++++ tests/core/test_application.py | 127 ++++++++++++++++++++++- tests/integrations/test_burr_pydantic.py | 40 +++++++ 5 files changed, 204 insertions(+), 6 deletions(-) diff --git a/burr/core/action.py b/burr/core/action.py index 240f0808b..0383daa24 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -1565,7 +1565,7 @@ def __init__( [StreamResultType, State[StateType]], tuple[StreamResultType, State[StateType]] ], callback: Callable[ - [Optional[StreamResultType], State[StateType], Optional[Exception]], + [Optional[StreamResultType], State[StateType], Optional[BaseException]], typing.Coroutine[None, None, None], ], ): @@ -1600,8 +1600,12 @@ async def aclose(self) -> None: exits) runs now rather than at garbage collection, then runs the post-step callback exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a no-op once the stream has completed. Exceptions raised by the action's cleanup propagate. + + A stream closed before it completes is reported to the callback (and so to + ``post_run_step`` hooks and persisters) with a ``GeneratorExit``, i.e. as a failed step. """ - exc: Optional[BaseException] = None + # Closing before completion means the step did not finish: report it as failed. + exc: Optional[BaseException] = GeneratorExit() if self._result is None else None try: await self.streaming_result_generator.aclose() except BaseException as e: diff --git a/burr/core/application.py b/burr/core/application.py index 48ce3c9d2..f8b15441e 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -1700,7 +1700,8 @@ async def astream_result( The :py:class:`AsyncStreamingResultContainer ` is meant as a convenience -- specifically this allows for hooks, callbacks, etc... so you can take the control flow and still have state updated afterwards. Hooks/state update will be called after an exception - is thrown during streaming, or the stream is completed. Note that it is undefined behavior to attempt to execute another action while a stream is in progress. + is thrown during streaming, the stream is completed, or the stream is closed early (``await container.aclose()``, or the + task is cancelled), in which case the step is reported as failed. Note that it is undefined behavior to attempt to execute another action while a stream is in progress. To see how this works, let's take the following action (simplified as a single-node workflow) as an example: diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index 89b126935..c6fb4af59 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -224,6 +224,38 @@ This object is effectively a cached iterator. You can use it as follows: Thus you can run this in a web-service, a streamlit app, etc... +Stopping a stream early +~~~~~~~~~~~~~~~~~~~~~~~ + +If you stop consuming an async stream before it finishes (e.g. the client disconnected), close it so the +action's cleanup (``finally`` blocks, ``async with`` exits such as an open HTTP stream) runs right away, +rather than at garbage collection. Either close the container directly, or close ``astream_iterate`` +(``contextlib.aclosing`` does this for you): + +.. code-block:: python + + action, async_streaming_result = await application.astream_result( + halt_after='streaming_response', inputs={"prompt": prompt} + ) + async for result in async_streaming_result: + if should_stop(result): + break + await async_streaming_result.aclose() + +.. code-block:: python + + import contextlib + + async with contextlib.aclosing(application.astream_iterate(halt_after=["respond"])) as steps: + async for action, streaming_result in steps: + async for item in streaming_result: + ... + break # leaving the ``async with`` closes the in-flight action + +A step that is closed (or whose task is cancelled) before it completes does not update state, and is reported +to ``post_run_step`` hooks with a ``GeneratorExit`` exception -- persisters record it as ``failed``, so +resuming re-runs that action. Exceptions raised by the action's cleanup propagate out of ``aclose()``. + -------------- Considerations -------------- diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 1e1cb9149..36e8bfd37 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -2806,10 +2806,12 @@ async def act(state: State): @pytest.mark.parametrize("consume_first_item", [True, False]) async def test_astream_iterate_aclose_unwinds_in_flight_action(consume_first_item: bool): """Closing astream_iterate mid-action runs the action's cleanup now (not at GC), - does not drain or advance it, and fires post_end_stream exactly once.""" + does not drain or advance it, fires post_end_stream exactly once, and reports the + step as failed.""" events = [] stream_event_tracker = StreamEventCaptureTracker() - app = _closing_probe_app(events, stream_event_tracker) + call_tracker = CallCaptureTracker() + app = _closing_probe_app(events, stream_event_tracker, call_tracker) steps = app.astream_iterate(halt_after=["act"]) async for _, container in steps: if consume_first_item: @@ -2819,6 +2821,111 @@ async def test_astream_iterate_aclose_unwinds_in_flight_action(consume_first_ite await steps.aclose() assert events == (["producer_closed"] if consume_first_item else []) assert len(stream_event_tracker.post_end_stream_calls) == 1 + assert len(call_tracker.post_called) == 1 + assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + + +async def test_astream_iterate_cancelled_mid_stream_unwinds_and_reports_failure(): + """Cancelling a consumer that wraps astream_iterate in ``aclosing`` (e.g. on a client + disconnect) unwinds the action immediately and reports the step as failed.""" + events = [] + call_tracker = CallCaptureTracker() + app = _closing_probe_app(events, call_tracker) + first_item_seen = asyncio.Event() + + async def consume(): + async with contextlib.aclosing(app.astream_iterate(halt_after=["act"])) as steps: + async for _, container in steps: + async for _ in container: + first_item_seen.set() + await asyncio.sleep(10) + + task = asyncio.create_task(consume()) + await first_item_seen.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert events == ["producer_closed"] + assert len(call_tracker.post_called) == 1 + assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + + +async def test_astream_iterate_aclose_mid_stream_persists_failed_and_resumes_at_same_action(): + @streaming_action(reads=[], writes=["x"]) + async def act(state: State): + yield {"item": "a"}, None + yield {"item": "done"}, state.update(x=1) + + @action(reads=["x"], writes=["y"]) + def after(state: State) -> State: + return state.update(y=state["x"] + 1) + + persister = SQLLitePersister(db_path=":memory:", table_name="test_table") + persister.initialize() + + def build(resume: bool) -> Application: + builder = ( + ApplicationBuilder() + .with_actions(act=act, after=after) + .with_transitions(("act", "after")) + ) + if resume: + builder = builder.initialize_from( + persister, resume_at_next_action=True, default_state={}, default_entrypoint="act" + ) + else: + builder = builder.with_entrypoint("act").with_state(State()) + return builder.with_identifiers(app_id="app").with_state_persister(persister).build() + + steps = build(resume=False).astream_iterate(halt_after=["after"]) + async for _, container in steps: + async for _ in container: + break + break + await steps.aclose() + assert persister.load("", "app")["status"] == "failed" + assert build(resume=True).get_next_action().name == "act" + + +async def test_astream_iterate_aclose_unwinds_multi_step_streaming_action(): + events = [] + + class MultiStepAct(AsyncStreamingAction): + @property + def reads(self) -> list[str]: + return [] + + @property + def writes(self) -> list[str]: + return ["x"] + + async def stream_run(self, state: State, **run_kwargs): + try: + for i in range(3): + yield {"i": i} + finally: + await asyncio.sleep(0) + events.append("stream_run_closed") + + def update(self, result: dict, state: State) -> State: + return state.update(x=result["i"]) + + app = ( + ApplicationBuilder() + .with_actions(act=MultiStepAct()) + .with_transitions() + .with_entrypoint("act") + .with_state(State()) + .build() + ) + steps = app.astream_iterate(halt_after=["act"]) + async for _, container in steps: + async for _ in container: + break + break + await steps.aclose() + assert events == ["stream_run_closed"] + assert "x" not in app.state async def test_streaming_result_container_aclose_is_idempotent_and_get_returns_initial_state(): @@ -2835,6 +2942,19 @@ async def test_streaming_result_container_aclose_is_idempotent_and_get_returns_i assert "x" not in state +async def test_streaming_result_container_aclose_after_completion_reports_success(): + events = [] + call_tracker = CallCaptureTracker() + app = _closing_probe_app(events, call_tracker) + _, container = await app.astream_result(halt_after=["act"]) + result, state = await container.get() + await container.aclose() + assert result == {"item": "done"} + assert state["x"] == 1 + assert len(call_tracker.post_called) == 1 + assert call_tracker.post_called[0][1]["exception"] is None + + async def test_streaming_result_container_aclose_propagates_cleanup_errors(): class CleanupError(Exception): pass @@ -2857,7 +2977,8 @@ async def callback(result, state, exc): break with pytest.raises(CleanupError): await container.aclose() - assert len(seen) == 1 and isinstance(seen[0], CleanupError) + assert len(seen) == 1 + assert isinstance(seen[0], CleanupError) async def test_astream_result_halt_after_run_through_streaming(): diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 662ecb514..85c4b7a2a 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/tests/integrations/test_burr_pydantic.py @@ -824,3 +824,43 @@ async def final_result_streamed( assert state.data.count == 20 assert isinstance(result, IntermediateModel) assert result.result == 20 + + +async def test_pydantic_streaming_async_aclose_unwinds_action(): + events = [] + + @streaming_action.pydantic( + reads=["count"], + writes=["final_result"], + state_output_type=CounterExampleAppState, + state_input_type=CounterExampleAppState, + stream_type=IntermediateModel, + ) + async def final_result_streamed( + state: CounterExampleAppState, + ) -> AsyncGenerator[Tuple[IntermediateModel, Optional[CounterExampleAppState]], None]: + try: + for i in range(3): + yield IntermediateModel(result=i), None + state.final_result = state.count + yield IntermediateModel(result=state.count), state + finally: + await asyncio.sleep(0) + events.append("action_closed") + + app = ( + ApplicationBuilder() + .with_actions(final_result=final_result_streamed) + .with_entrypoint("final_result") + .with_typing(PydanticTypingSystem(CounterExampleAppState)) + .with_state(CounterExampleAppState()) + .with_transitions() + .build() + ) + steps = app.astream_iterate(halt_after=["final_result"]) + async for _, container in steps: + async for _ in container: + break + break + await steps.aclose() + assert events == ["action_closed"] From c285fb7cabb459448778059f8f2e0b786e06a426 Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Tue, 3 Nov 2026 08:59:23 -0800 Subject: [PATCH 3/8] fix(core): hold AsyncStreamingResultContainer's iterator weakly Storing the __aiter__ generator on the container created a container <-> iterator reference cycle (the iterator's closure references the container). When a caller broke out of the stream and dropped the container without calling aclose(), the iterator's finalizer -- and so the post-step callback/hooks/persisters -- was deferred until cyclic GC instead of running on refcount drop as it did before. Keep a weakref instead; aclose() still closes a live iterator. --- burr/core/action.py | 15 ++++++++++----- tests/core/test_application.py | 21 +++++++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/burr/core/action.py b/burr/core/action.py index 0383daa24..05e2cb521 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -25,6 +25,7 @@ import textwrap import types import typing +import weakref UnionType = types.UnionType @@ -1584,7 +1585,9 @@ def __init__( self._callback = callback self._result = None self._callback_realized = False - self._iterator: Optional[AsyncGenerator[StreamResultType, None]] = None + # Weak, so the container <-> iterator cycle doesn't defer the iterator's finalizer (and so + # the post-step callback) to cyclic GC when a caller abandons the stream without closing it. + self._iterator_ref: Optional[weakref.ref] = None async def _realize_callback(self, exc: Optional[BaseException]) -> None: if self._result is None: @@ -1613,8 +1616,9 @@ async def aclose(self) -> None: raise finally: await self._realize_callback(exc) - if self._iterator is not None: - await self._iterator.aclose() + iterator = self._iterator_ref() if self._iterator_ref is not None else None + if iterator is not None: + await iterator.aclose() async def __anext__(self) -> StreamResultType: """Moves to the next state in the streaming result""" @@ -1642,8 +1646,9 @@ async def gen_fn(): await self._realize_callback(sys.exc_info()[1]) # return it as `__aiter__` cannot be async/have awaits :/ - self._iterator = gen_fn() - return self._iterator + iterator = gen_fn() + self._iterator_ref = weakref.ref(iterator) + return iterator async def get(self) -> tuple[Optional[StreamResultType], State[StateType]]: # exhaust the generator diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 36e8bfd37..7a8c7c0cc 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -19,6 +19,7 @@ import collections import contextlib import datetime +import gc import logging import typing import uuid @@ -2955,6 +2956,26 @@ async def test_streaming_result_container_aclose_after_completion_reports_succes assert call_tracker.post_called[0][1]["exception"] is None +async def test_streaming_result_container_abandoned_without_aclose_reports_without_cyclic_gc(): + """Breaking out and dropping the container (no aclose) must still report the step via + refcounting alone -- the container must not form a reference cycle with its iterator, + which would defer the post-step hooks until cyclic GC runs.""" + call_tracker = CallCaptureTracker() + app = _closing_probe_app([], call_tracker) + gc.disable() + try: + _, container = await app.astream_result(halt_after=["act"]) + async for _ in container: + break + del container + for _ in range(5): # let the event loop run the scheduled asyncgen finalizer + await asyncio.sleep(0) + assert len(call_tracker.post_called) == 1 + assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + finally: + gc.enable() + + async def test_streaming_result_container_aclose_propagates_cleanup_errors(): class CleanupError(Exception): pass From 130676622b65a477c4682418f870d26a708f5598 Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Tue, 3 Nov 2026 09:20:24 -0800 Subject: [PATCH 4/8] fix(core): propagate early close through sync streaming actions Mirror the async fix on the sync path. Closing Application.stream_iterate() (or abandoning a StreamingResultContainer) mid-action left the action's generator suspended until GC, because each forwarding layer used a bare `for` loop and never closed the generator it read from. - StreamingResultContainer.close() (new, public): stops the stream without draining it, closes the action's generator chain, and runs the post-step callback exactly once. A close before completion is reported as a failed step (GeneratorExit); after completion it still reports success. Cleanup exceptions propagate. The __iter__ generator is held weakly to avoid a container <-> iterator reference cycle. - stream_iterate: on GeneratorExit at its yield, close the in-flight container before re-raising. - _run_single_step_streaming_action / _run_multi_step_streaming_action: `finally: generator.close()`. - pydantic action_generator: wrap the inner generator in contextlib.closing(...). (type_eraser and _stream_run_and_update already use `yield from`, which propagates close().) - callback exception typed Optional[BaseException]. - tests mirroring the async ones (close mid-action, persister status + resume, multi-step, idempotent close, close after completion, abandoned container without cyclic GC, cleanup errors, pydantic); docs cover sync. --- burr/core/action.py | 47 +++++- burr/core/application.py | 16 +- burr/integrations/pydantic.py | 13 +- docs/concepts/streaming-actions.rst | 34 +++- tests/core/test_application.py | 205 +++++++++++++++++++++++ tests/integrations/test_burr_pydantic.py | 39 +++++ 6 files changed, 331 insertions(+), 23 deletions(-) diff --git a/burr/core/action.py b/burr/core/action.py index 05e2cb521..e8eb99897 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -1475,7 +1475,7 @@ def __init__( streaming_result_generator: GeneratorReturnType, initial_state: State[StateType], process_result: Callable[[dict, State], tuple[dict, State]], - callback: Callable[[Optional[dict], State, Optional[Exception]], None], + callback: Callable[[Optional[dict], State, Optional[BaseException]], None], ): """Initializes a streaming result container. User will never call directly. @@ -1493,6 +1493,40 @@ def __init__( self._initial_state = initial_state self._result = None self._callback_realized = False + # Weak, so the container <-> iterator cycle doesn't defer the iterator's finalizer (and so + # the post-step callback) to cyclic GC when a caller abandons the stream without closing it. + self._iterator_ref: Optional[weakref.ref] = None + + def _realize_callback(self, exc: Optional[BaseException]) -> None: + if self._result is None: + self._result = None, self._initial_state + if not self._callback_realized: + self._callback_realized = True + self._callback(*self._result, exc) + + def close(self) -> None: + """Stops the stream early without draining it. + + Closes the action's generator chain so its cleanup (``finally`` blocks, ``with`` + exits) runs now rather than at garbage collection, then runs the post-step callback + exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a + no-op once the stream has completed. Exceptions raised by the action's cleanup propagate. + + A stream closed before it completes is reported to the callback (and so to + ``post_run_step`` hooks and persisters) with a ``GeneratorExit``, i.e. as a failed step. + """ + # Closing before completion means the step did not finish: report it as failed. + exc: Optional[BaseException] = GeneratorExit() if self._result is None else None + try: + self.streaming_result_generator.close() + except BaseException as e: + exc = e + raise + finally: + self._realize_callback(exc) + iterator = self._iterator_ref() if self._iterator_ref is not None else None + if iterator is not None: + iterator.close() def __next__(self) -> StreamResultType: if self._result is not None: @@ -1513,16 +1547,13 @@ def gen_fn(): except StopIteration: return finally: - if self._result is None: - self._result = None, self._initial_state - if not self._callback_realized: - exc = sys.exc_info()[1] - self._callback_realized = True - self._callback(*self._result, exc) + self._realize_callback(sys.exc_info()[1]) # We really don't need an internal generator function but this was done to keep it the same # as the async version - return gen_fn() + iterator = gen_fn() + self._iterator_ref = weakref.ref(iterator) + return iterator def get(self) -> Tuple[StreamResultType, State[StateType]]: # exhaust the generator diff --git a/burr/core/application.py b/burr/core/application.py index f8b15441e..7f6f4a7aa 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -387,6 +387,9 @@ def _run_single_step_streaming_action( e, exc_info=True, ) + finally: + # Unwind the action now on early close instead of leaving it to GC. + generator.close() if state_update is None: raise ValueError( @@ -530,6 +533,9 @@ def _run_multi_step_streaming_action( e, exc_info=True, ) + finally: + # Unwind the action now on early close instead of leaving it to GC. + generator.close() state_update = _run_reducer(action, state, result, action.name) _validate_result(result, action.name, action.schema) _validate_reducer_writes(action, state_update, action.name) @@ -1445,7 +1451,8 @@ def stream_result( The :py:class:`StreamingResultContainer ` is meant as a convenience -- specifically this allows for hooks, callbacks, etc... so you can take the control flow and still have state updated afterwards. Hooks/state update will be called after an exception - is thrown during streaming, or the stream is completed. Note that it is undefined behavior to attempt to execute another action while a stream is in progress. + is thrown during streaming, the stream is completed, or the stream is closed early (``container.close()``), in which case the + step is reported as failed. Note that it is undefined behavior to attempt to execute another action while a stream is in progress. To see how this works, let's take the following action (simplified as a single-node workflow) as an example: @@ -1982,7 +1989,12 @@ def stream_iterate( _, streaming_result = self.stream_result( halt_after=[next_action.name], halt_before=None, inputs=inputs ) - yield next_action, streaming_result + try: + yield next_action, streaming_result + except GeneratorExit: + # The caller closed us mid-action: unwind it rather than leaving it to GC. + streaming_result.close() + raise # We need to ensure it's fully exhausted before going to the next action streaming_result.get() if self._should_halt_iterate(halt_before, halt_after, next_action): diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 08533c837..e0457b233 100644 --- a/burr/integrations/pydantic.py +++ b/burr/integrations/pydantic.py @@ -347,12 +347,13 @@ def action_generator( state: State, **kwargs ) -> Generator[tuple[PartialType, Optional[State]], None, None]: model_to_use = model_from_state(model=SubsetInputType, state=state) - for partial, state_update in fn(state=model_to_use, **kwargs): - if state_update is None: - yield partial, None - else: - output = SubsetOutputType(**model_to_dict(state_update, include=writes)) - yield partial, merge_to_state(model=output, write_keys=writes, state=state) + with contextlib.closing(fn(state=model_to_use, **kwargs)) as gen: + for partial, state_update in gen: + if state_update is None: + yield partial, None + else: + output = SubsetOutputType(**model_to_dict(state_update, include=writes)) + yield partial, merge_to_state(model=output, write_keys=writes, state=state) async def async_action_generator( state: State, **kwargs diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index c6fb4af59..0d3009d65 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -227,10 +227,30 @@ Thus you can run this in a web-service, a streamlit app, etc... Stopping a stream early ~~~~~~~~~~~~~~~~~~~~~~~ -If you stop consuming an async stream before it finishes (e.g. the client disconnected), close it so the -action's cleanup (``finally`` blocks, ``async with`` exits such as an open HTTP stream) runs right away, -rather than at garbage collection. Either close the container directly, or close ``astream_iterate`` -(``contextlib.aclosing`` does this for you): +If you stop consuming a stream before it finishes (e.g. the client disconnected), close it so the +action's cleanup (``finally`` blocks, ``with``/``async with`` exits such as an open HTTP stream) runs right away, +rather than at garbage collection. Either close the container directly (``close()`` / ``await aclose()``), or +close ``stream_iterate`` / ``astream_iterate`` (``contextlib.closing`` / ``contextlib.aclosing`` does this for you): + +.. code-block:: python + + action, streaming_result = application.stream_result( + halt_after='streaming_response', inputs={"prompt": prompt} + ) + for result in streaming_result: + if should_stop(result): + break + streaming_result.close() + +.. code-block:: python + + import contextlib + + with contextlib.closing(application.stream_iterate(halt_after=["respond"])) as steps: + for action, streaming_result in steps: + for item in streaming_result: + ... + break # leaving the ``with`` closes the in-flight action .. code-block:: python @@ -252,9 +272,9 @@ rather than at garbage collection. Either close the container directly, or close ... break # leaving the ``async with`` closes the in-flight action -A step that is closed (or whose task is cancelled) before it completes does not update state, and is reported -to ``post_run_step`` hooks with a ``GeneratorExit`` exception -- persisters record it as ``failed``, so -resuming re-runs that action. Exceptions raised by the action's cleanup propagate out of ``aclose()``. +A step that is closed (or, for async, whose task is cancelled) before it completes does not update state, and is +reported to ``post_run_step`` hooks with a ``GeneratorExit`` exception -- persisters record it as ``failed``, so +resuming re-runs that action. Exceptions raised by the action's cleanup propagate out of ``close()`` / ``aclose()``. -------------- Considerations diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 7a8c7c0cc..842ceb681 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -40,6 +40,7 @@ SingleStepAction, SingleStepStreamingAction, StreamingAction, + StreamingResultContainer, action, default, expr, @@ -3002,6 +3003,210 @@ async def callback(result, state, exc): assert isinstance(seen[0], CleanupError) +class _ExceptionTypeCapture(PostRunStepHook): + """Records only the exception *type* -- holding the exception itself would keep its traceback, + and so the stream's frames, alive, defeating refcount-based cleanup checks.""" + + def __init__(self): + self.exception_types = [] + + def post_run_step(self, *, exception: Optional[BaseException], **future_kwargs: Any): + self.exception_types.append(type(exception) if exception is not None else None) + + +def _sync_closing_probe_app(events: list, *hooks) -> Application: + """Sync twin of ``_closing_probe_app``: one streaming action wrapping a nested producer.""" + + def producer(): + try: + yield "a" + yield "b" + finally: + events.append("producer_closed") + + @streaming_action(reads=[], writes=["x"]) + def act(state: State): + with contextlib.closing(producer()) as items: + for item in items: + yield {"item": item}, None + events.append("advanced") + yield {"item": "done"}, state.update(x=1) + + return ( + ApplicationBuilder() + .with_actions(act) + .with_transitions() + .with_entrypoint("act") + .with_state(State()) + .with_hooks(*hooks) + .build() + ) + + +@pytest.mark.parametrize("consume_first_item", [True, False]) +def test_stream_iterate_close_unwinds_in_flight_action(consume_first_item: bool): + """Closing stream_iterate mid-action runs the action's cleanup now (not at GC), does not + drain or advance it, fires post_end_stream exactly once, and reports the step as failed.""" + events = [] + stream_event_tracker = StreamEventCaptureTracker() + call_tracker = CallCaptureTracker() + app = _sync_closing_probe_app(events, stream_event_tracker, call_tracker) + steps = app.stream_iterate(halt_after=["act"]) + for _, container in steps: + if consume_first_item: + for _ in container: + break + break + steps.close() + assert events == (["producer_closed"] if consume_first_item else []) + assert len(stream_event_tracker.post_end_stream_calls) == 1 + assert len(call_tracker.post_called) == 1 + assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + + +def test_stream_iterate_close_mid_stream_persists_failed_and_resumes_at_same_action(): + @streaming_action(reads=[], writes=["x"]) + def act(state: State): + yield {"item": "a"}, None + yield {"item": "done"}, state.update(x=1) + + @action(reads=["x"], writes=["y"]) + def after(state: State) -> State: + return state.update(y=state["x"] + 1) + + persister = SQLLitePersister(db_path=":memory:", table_name="test_table") + persister.initialize() + + def build(resume: bool) -> Application: + builder = ( + ApplicationBuilder() + .with_actions(act=act, after=after) + .with_transitions(("act", "after")) + ) + if resume: + builder = builder.initialize_from( + persister, resume_at_next_action=True, default_state={}, default_entrypoint="act" + ) + else: + builder = builder.with_entrypoint("act").with_state(State()) + return builder.with_identifiers(app_id="app").with_state_persister(persister).build() + + with contextlib.closing(build(resume=False).stream_iterate(halt_after=["after"])) as steps: + for _, container in steps: + for _ in container: + break + break + assert persister.load("", "app")["status"] == "failed" + assert build(resume=True).get_next_action().name == "act" + + +def test_stream_iterate_close_unwinds_multi_step_streaming_action(): + events = [] + + class MultiStepAct(StreamingAction): + @property + def reads(self) -> list[str]: + return [] + + @property + def writes(self) -> list[str]: + return ["x"] + + def stream_run(self, state: State, **run_kwargs): + try: + for i in range(3): + yield {"i": i} + finally: + events.append("stream_run_closed") + + def update(self, result: dict, state: State) -> State: + return state.update(x=result["i"]) + + app = ( + ApplicationBuilder() + .with_actions(act=MultiStepAct()) + .with_transitions() + .with_entrypoint("act") + .with_state(State()) + .build() + ) + steps = app.stream_iterate(halt_after=["act"]) + for _, container in steps: + for _ in container: + break + break + steps.close() + assert events == ["stream_run_closed"] + assert "x" not in app.state + + +def test_streaming_result_container_close_is_idempotent_and_get_returns_initial_state(): + events = [] + app = _sync_closing_probe_app(events) + _, container = app.stream_result(halt_after=["act"]) + for _ in container: + break + container.close() + container.close() + result, state = container.get() + assert events == ["producer_closed"] + assert result is None + assert "x" not in state + + +def test_streaming_result_container_close_after_completion_reports_success(): + call_tracker = CallCaptureTracker() + app = _sync_closing_probe_app([], call_tracker) + _, container = app.stream_result(halt_after=["act"]) + result, state = container.get() + container.close() + assert result == {"item": "done"} + assert state["x"] == 1 + assert len(call_tracker.post_called) == 1 + assert call_tracker.post_called[0][1]["exception"] is None + + +def test_streaming_result_container_abandoned_without_close_reports_without_cyclic_gc(): + """Breaking out and dropping the container (no close) must unwind the action and report the + step via refcounting alone -- the container must not form a cycle with its iterator.""" + events = [] + exception_capture = _ExceptionTypeCapture() + app = _sync_closing_probe_app(events, exception_capture) + gc.disable() + try: + _, container = app.stream_result(halt_after=["act"]) + for _ in container: + break + del container + assert events == ["producer_closed"] + assert exception_capture.exception_types == [GeneratorExit] + finally: + gc.enable() + + +def test_streaming_result_container_close_propagates_cleanup_errors(): + class CleanupError(Exception): + pass + + def failing_cleanup(): + try: + yield {"item": "a"}, None + finally: + raise CleanupError("cleanup failed") + + seen = [] + container = StreamingResultContainer( + failing_cleanup(), State(), lambda r, s: (r, s), lambda r, s, exc: seen.append(exc) + ) + # Hold the iterator: breaking out of a ``for`` would drop it and report the close immediately. + items = iter(container) + next(items) + with pytest.raises(CleanupError): + container.close() + assert len(seen) == 1 + assert isinstance(seen[0], CleanupError) + + async def test_astream_result_halt_after_run_through_streaming(): action_tracker = CallCaptureTracker() stream_event_tracker = StreamEventCaptureTrackerAsync() diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 85c4b7a2a..d0b32bd53 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/tests/integrations/test_burr_pydantic.py @@ -864,3 +864,42 @@ async def final_result_streamed( break await steps.aclose() assert events == ["action_closed"] + + +def test_pydantic_streaming_sync_close_unwinds_action(): + events = [] + + @streaming_action.pydantic( + reads=["count"], + writes=["final_result"], + state_output_type=CounterExampleAppState, + state_input_type=CounterExampleAppState, + stream_type=IntermediateModel, + ) + def final_result_streamed( + state: CounterExampleAppState, + ) -> Generator[Tuple[IntermediateModel, Optional[CounterExampleAppState]], None, None]: + try: + for i in range(3): + yield IntermediateModel(result=i), None + state.final_result = state.count + yield IntermediateModel(result=state.count), state + finally: + events.append("action_closed") + + app = ( + ApplicationBuilder() + .with_actions(final_result=final_result_streamed) + .with_entrypoint("final_result") + .with_typing(PydanticTypingSystem(CounterExampleAppState)) + .with_state(CounterExampleAppState()) + .with_transitions() + .build() + ) + steps = app.stream_iterate(halt_after=["final_result"]) + for _, container in steps: + for _ in container: + break + break + steps.close() + assert events == ["action_closed"] From 0e23c3e3a599fef2b2739819ccdf19194003c394 Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Tue, 3 Nov 2026 10:32:16 -0800 Subject: [PATCH 5/8] docs(examples): add streaming-early-stop example Runnable documentation for stopping a streaming action early (user stop, moderation/budget check, client disconnect), covering the close API added in this PR: - application.py: one demo per pattern -- break + close()/aclose(), contextlib.closing/aclosing around stream_iterate/astream_iterate, cancelling the consuming task, what persisters record and that resuming re-runs the step, and cleanup errors raised by close(). Each demo prints a timeline (cleanup vs. post_run_step) so the behaviour is visible. - server.py: FastAPI SSE endpoint; on client disconnect the action's resources are released right away and hooks see CancelledError. - README.md: the patterns, a "what gets reported" table, and how to tell a deliberate stop from a real failure in a hook. Notes that async post_run_step hooks/persisters are cancelled on an ASGI disconnect, so disconnects are only recorded by sync hooks. - notebook.ipynb (executed) + statemachine.png, per examples conventions. - tests/core/test_streaming_early_stop_example.py runs the demos and asserts the timelines, so the example can't silently go stale. - docs: cancellation is reported as CancelledError (not GeneratorExit); link to the example. examples/README.md index + pyproject sdist exclude (matching streaming-overview). --- docs/concepts/streaming-actions.rst | 10 +- examples/README.md | 1 + examples/streaming-early-stop/README.md | 154 ++++++++ examples/streaming-early-stop/__init__.py | 16 + examples/streaming-early-stop/application.py | 264 ++++++++++++++ examples/streaming-early-stop/notebook.ipynb | 343 ++++++++++++++++++ examples/streaming-early-stop/server.py | 86 +++++ .../streaming-early-stop/statemachine.png | Bin 0 -> 3022 bytes pyproject.toml | 1 + .../core/test_streaming_early_stop_example.py | 90 +++++ 10 files changed, 962 insertions(+), 3 deletions(-) create mode 100644 examples/streaming-early-stop/README.md create mode 100644 examples/streaming-early-stop/__init__.py create mode 100644 examples/streaming-early-stop/application.py create mode 100644 examples/streaming-early-stop/notebook.ipynb create mode 100644 examples/streaming-early-stop/server.py create mode 100644 examples/streaming-early-stop/statemachine.png create mode 100644 tests/core/test_streaming_early_stop_example.py diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index 0d3009d65..4e30b8ee8 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -272,9 +272,13 @@ close ``stream_iterate`` / ``astream_iterate`` (``contextlib.closing`` / ``conte ... break # leaving the ``async with`` closes the in-flight action -A step that is closed (or, for async, whose task is cancelled) before it completes does not update state, and is -reported to ``post_run_step`` hooks with a ``GeneratorExit`` exception -- persisters record it as ``failed``, so -resuming re-runs that action. Exceptions raised by the action's cleanup propagate out of ``close()`` / ``aclose()``. +A step that is stopped before it completes does not update state, and is reported to ``post_run_step`` hooks +with the exception that stopped it -- ``GeneratorExit`` when it is closed, ``CancelledError`` when the consuming +task is cancelled (e.g. a web client disconnected). Persisters record it as ``failed``, so resuming re-runs that +action. Exceptions raised by the action's cleanup propagate out of ``close()`` / ``aclose()``. + +See the `streaming-early-stop example `_ +for runnable versions of these patterns, including a FastAPI endpoint that handles client disconnects. -------------- Considerations diff --git a/examples/README.md b/examples/README.md index 9d4447cc2..ea001c7e5 100644 --- a/examples/README.md +++ b/examples/README.md @@ -48,6 +48,7 @@ Note we have a few more in [other-examples](other-examples/), but those do not y - [multi-agent-collaboration](multi-agent-collaboration/) - This example shows how to use Burr to create a multi-agent collaboration. This is a clone of the following [LangGraph example](https://github.com/langchain-ai/langgraph/blob/main/examples/multi_agent/multi-agent-collaboration.ipynb). - [multi-modal-chatbot](multi-modal-chatbot/) - This example shows how to use Burr to create a multi-modal chatbot. This demonstrates how to use a model to delegate to other models conditionally. - [streaming-overview](streaming-overview/) - This example shows how we can use the streaming API to respond to return quicker results to the user and build a seamless experience +- [streaming-early-stop](streaming-early-stop/) - This example shows how to stop a streaming action early (user stop, moderation, client disconnect) so its resources are released right away and the step is reported as failed. - [integrations/bedrock](integrations/bedrock/) - Minimal graphs using Amazon Bedrock (`BedrockAction` and `BedrockStreamingAction`). - [tracing-and-spans](tracing-and-spans/) - This example shows how to use Burr to create a simple chatbot with additional visibility. This is a good starting point for understanding how to use Burr's tracing functionality. - [web-server](web-server/) - This example shows how to use Burr in a web server. This is a good starting point for understanding how to use Burr for interaction. diff --git a/examples/streaming-early-stop/README.md b/examples/streaming-early-stop/README.md new file mode 100644 index 000000000..7ba3e49d4 --- /dev/null +++ b/examples/streaming-early-stop/README.md @@ -0,0 +1,154 @@ + + +# Stopping a streaming action early + +Streaming actions often hold something open while they stream, such as an HTTP connection to an LLM, a +database cursor, or a tracing span. They release it in a `finally` block or a `with` exit. When you stop +reading a stream before it ends (the user hits "stop", a moderation check trips, a token budget runs out, +the client disconnects), you want three things: + +1. **The action's cleanup runs now**, not whenever the garbage collector gets to it. +2. **The step is reported honestly.** Hooks, trackers and persisters see it as *failed*, not completed, so + resuming the application re-runs it. +3. **The state is not updated.** A half-finished response is never committed. + +This example shows the patterns that give you all three. Every demo in [application.py](application.py) +prints a timeline of what happened, so you can see when the resource is released and what Burr +reports. Run it with: + +```bash +python application.py +``` + +![statemachine](statemachine.png) + +## The patterns + +### 1. Stop reading, then `close()` the container + +```python +action, container = app.stream_result(halt_after=["respond"]) +for item in container: + send_to_user(item["token"]) + if wants_to_stop(item): + break +container.close() # async: await container.aclose() + +result, state = container.get() # (None, state without the update) +``` + +`close()` / `aclose()` stops the action and runs its cleanup. Calling it more than once is safe. Once +the stream has finished normally, it does nothing (the step is still reported as completed). + +**Always call it after a `break`.** A `break` on its own stops *you* from reading, but the action's +cleanup only runs when the container is closed or garbage-collected. + +### 2. Use `contextlib.closing` / `contextlib.aclosing` around `stream_iterate` + +```python +import contextlib + +with contextlib.closing(app.stream_iterate(halt_after=["respond"])) as steps: + for action, container in steps: + for item in container: + ... + if wants_to_stop(item): + break + break +# leaving the block closed the in-flight action + +async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: + async for action, container in steps: + ... +``` + +This is the most robust option: however you leave the block (`break`, `return`, an exception, or +the task being cancelled), the in-flight action is closed. **Prefer this whenever you use +`stream_iterate` / `astream_iterate`.** + +### 3. Client disconnects (web servers): cancellation does the work + +When a client disconnects, ASGI servers such as Starlette/FastAPI *cancel* the task that is +producing the response. If your response generator uses pattern 2, that cancellation unwinds +the Burr action immediately. See [server.py](server.py): + +```python +@app.get("/chat") +async def chat(): + async def sse(): + async with contextlib.aclosing(burr_app.astream_iterate(halt_after=["respond"])) as steps: + async for _, container in steps: + async for item in container: + yield f"data: {item['token']}\n\n" + return StreamingResponse(sse(), media_type="text/event-stream") +``` + +```bash +uvicorn server:app +curl -N --max-time 1 localhost:8000/chat # hang up after a second +# server log: +# respond: open LLM connection +# respond: close LLM connection +# post_run_step: respond exception=CancelledError(...) +``` + +To stop a stream *from another task*, cancel the consuming task (as the server does). Don't call +`aclose()` on a container that another task is iterating: you can't close an async generator while +it's running. + +## What gets reported + +| How the stream ended | `post_run_step(exception=...)` | persister status | state updated? | on resume | +|---|---|---|---|---| +| read to the end | `None` | `completed` | yes | next action | +| `close()` / `aclose()` / `closing` before the end | `GeneratorExit()` | `failed` | no | **same action runs again** | +| task cancelled (e.g. client disconnect) | `CancelledError()` | `failed` | no | **same action runs again** | +| the action raised | that exception | `failed` | no | same action runs again | + +So in a `post_run_step` hook, `exception is not None` means "this step did not complete". If you +want to tell a deliberate stop apart from a real failure: + +```python +def post_run_step(self, *, exception, **kwargs): + if exception is None: + ... # completed + elif isinstance(exception, (GeneratorExit, asyncio.CancelledError)): + ... # stopped early by the caller / client + else: + ... # a real error +``` + +Note that `GeneratorExit` and `CancelledError` derive from `BaseException`, not `Exception`. + +## Things to know + +- **Cleanup errors.** If the action's `finally` raises while it is being closed, `close()` / `aclose()` + raises that error to you. See `cleanup_error_demo` in [application.py](application.py). +- **Use sync hooks/persisters to record client disconnects.** On a disconnect, Starlette/FastAPI + cancels the whole response task. Sync `post_run_step` hooks (and sync persisters) still run and + see the `CancelledError`. *Async* `post_run_step` hooks (and async persisters) are cancelled + before they start, so they never record the stopped step. Shielding inside the hook does not + help, because the hook body never runs. +- **Sync vs async.** Everything above works the same way for `stream_result`/`stream_iterate` + (with `close()` / `contextlib.closing`) and `astream_result`/`astream_iterate` (with `aclose()` / + `contextlib.aclosing`). + +See also the [streaming actions docs](https://burr.apache.org/concepts/streaming-actions/) and the +[streaming-fastapi](../streaming-fastapi) example for a full chatbot. diff --git a/examples/streaming-early-stop/__init__.py b/examples/streaming-early-stop/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/examples/streaming-early-stop/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/examples/streaming-early-stop/application.py b/examples/streaming-early-stop/application.py new file mode 100644 index 000000000..720878641 --- /dev/null +++ b/examples/streaming-early-stop/application.py @@ -0,0 +1,264 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Stopping a streaming action early -- the patterns to use. + +The application is a single streaming action, ``respond``, that pretends to stream tokens from an +LLM. It holds a resource while it streams (think: an open HTTP connection) and releases it in a +``finally`` block. Each ``*_demo`` function below shows one way to stop the stream early, and +prints what happened, in order, so you can see when the resource is released and what Burr +reports to hooks/persisters. + +Run it with ``python application.py``. See README.md for the explanation. +""" + +import asyncio +import contextlib +from typing import AsyncGenerator, Generator, List, Optional, Tuple + +from burr.core import Application, ApplicationBuilder, State +from burr.core.action import streaming_action +from burr.core.persistence import SQLLitePersister +from burr.lifecycle import PostRunStepHook + +TOKENS = ["Burr ", "streams ", "tokens ", "one ", "at ", "a ", "time."] + + +class Timeline: + """Collects what happened, in order, so the demos can print it.""" + + def __init__(self): + self.events: List[str] = [] + + def __call__(self, event: str): + self.events.append(event) + + def print(self, title: str): + print(f"\n=== {title}") + for event in self.events: + print(f" {event}") + + +class ReportStep(PostRunStepHook): + """Stands in for any hook/tracker/persister: shows how each step was reported.""" + + def __init__(self, log: Timeline): + self.log = log + + def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): + outcome = "completed" if exception is None else f"failed with {exception!r}" + self.log(f"post_run_step: {action.name} {outcome}") + + +def _respond_sync(log: Timeline, cleanup_error: bool = False): + @streaming_action(reads=[], writes=["response"]) + def respond(state: State) -> Generator[Tuple[dict, Optional[State]], None, None]: + log("respond: open LLM connection") + try: + response = "" + for token in TOKENS: + response += token + yield {"token": token}, None + yield {"token": ""}, state.update(response=response) + finally: + log("respond: close LLM connection") + if cleanup_error: + raise ConnectionError("LLM connection did not close cleanly") + + return respond + + +def _respond_async(log: Timeline, delay: float = 0.0): + @streaming_action(reads=[], writes=["response"]) + async def respond(state: State) -> AsyncGenerator[Tuple[dict, Optional[State]], None]: + log("respond: open LLM connection") + try: + response = "" + for token in TOKENS: + await asyncio.sleep(delay) # waiting for the next chunk from the network + response += token + yield {"token": token}, None + yield {"token": ""}, state.update(response=response) + finally: + log("respond: close LLM connection") + + return respond + + +def application( + respond, log: Timeline, persister: Optional[SQLLitePersister] = None +) -> Application: + builder = ( + ApplicationBuilder() + .with_actions(respond=respond) + .with_transitions() + .with_entrypoint("respond") + .with_state(State()) + .with_hooks(ReportStep(log)) + .with_identifiers(app_id="early-stop-demo") + ) + if persister is not None: + builder = builder.with_state_persister(persister) + return builder.build() + + +def wants_to_stop(token: str) -> bool: + """Stand-in for "the user hit stop", a moderation check, a token budget, ...""" + return token == "tokens " + + +# -------------------------------------------------------------------------------------------- +# Pattern 1: stop reading, then close() the container. +# -------------------------------------------------------------------------------------------- + + +def stop_with_close_demo(): + log = Timeline() + app = application(_respond_sync(log), log) + _, container = app.stream_result(halt_after=["respond"]) + for item in container: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + container.close() # releases the action's resources *now* + log("caller: closed the container") + result, state = container.get() # the step did not finish: no result, state not updated + log(f"caller: get() -> result={result!r}, 'response' in state={'response' in state}") + log.print("Pattern 1: break, then container.close()") + + +async def stop_with_aclose_demo(): + log = Timeline() + app = application(_respond_async(log), log) + _, container = await app.astream_result(halt_after=["respond"]) + async for item in container: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + await container.aclose() + log("caller: closed the container") + log.print("Pattern 1 (async): break, then await container.aclose()") + + +# -------------------------------------------------------------------------------------------- +# Pattern 2: wrap stream_iterate/astream_iterate in contextlib.closing/aclosing. +# -------------------------------------------------------------------------------------------- + + +def stop_iterate_with_closing_demo(): + log = Timeline() + app = application(_respond_sync(log), log) + with contextlib.closing(app.stream_iterate(halt_after=["respond"])) as steps: + for _, container in steps: + for item in container: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + break + log("caller: left the `with` block") + log.print("Pattern 2: contextlib.closing(app.stream_iterate(...))") + + +# -------------------------------------------------------------------------------------------- +# Pattern 3: the client disconnects -> cancel the task that is consuming the stream. +# -------------------------------------------------------------------------------------------- + + +async def client_disconnect_demo(): + log = Timeline() + app = application(_respond_async(log, delay=0.05), log) + + async def stream_to_client(): + async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: + async for _, container in steps: + async for item in container: + log(f"server: sent {item['token']!r}") + + task = asyncio.create_task(stream_to_client()) + await asyncio.sleep(0.12) + log("server: client disconnected -> task.cancel()") + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + log.print("Pattern 3: client disconnect, cancel the consuming task") + + +# -------------------------------------------------------------------------------------------- +# What a stopped step means for persistence: it is recorded as failed, and resuming re-runs it. +# -------------------------------------------------------------------------------------------- + + +def resume_after_stop_demo(): + log = Timeline() + persister = SQLLitePersister(db_path=":memory:", table_name="burr_state") + persister.initialize() + app = application(_respond_sync(log), log, persister=persister) + _, container = app.stream_result(halt_after=["respond"]) + for item in container: + if wants_to_stop(item["token"]): + break + container.close() + log(f"persister: last step saved as {persister.load('', 'early-stop-demo')['status']!r}") + + resumed = ( + ApplicationBuilder() + .with_actions(respond=_respond_sync(log)) + .with_transitions() + .initialize_from( + persister, resume_at_next_action=True, default_state={}, default_entrypoint="respond" + ) + .with_identifiers(app_id="early-stop-demo") + .build() + ) + log(f"resume: next action is {resumed.get_next_action().name!r} (it runs again)") + log.print("Persistence: a stopped step is 'failed' and is re-run on resume") + + +# -------------------------------------------------------------------------------------------- +# Cleanup errors: close() raises them to the caller. +# -------------------------------------------------------------------------------------------- + + +def cleanup_error_demo(): + log = Timeline() + app = application(_respond_sync(log, cleanup_error=True), log) + _, container = app.stream_result(halt_after=["respond"]) + for item in container: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + try: + container.close() + except ConnectionError as e: + log(f"caller: close() raised {e!r}") + log.print("Cleanup errors: close() raises them to the caller") + + +def visualize(): + """Writes statemachine.png (the app is a single streaming action).""" + application(_respond_sync(Timeline()), Timeline()).visualize( + output_file_path="statemachine", include_conditions=True, format="png" + ) + + +if __name__ == "__main__": + stop_with_close_demo() + asyncio.run(stop_with_aclose_demo()) + stop_iterate_with_closing_demo() + asyncio.run(client_disconnect_demo()) + resume_after_stop_demo() + cleanup_error_demo() diff --git a/examples/streaming-early-stop/notebook.ipynb b/examples/streaming-early-stop/notebook.ipynb new file mode 100644 index 000000000..0223729da --- /dev/null +++ b/examples/streaming-early-stop/notebook.ipynb @@ -0,0 +1,343 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "1c749fcb", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:53.338155Z", + "iopub.status.busy": "2026-11-03T18:29:53.337954Z", + "iopub.status.idle": "2026-11-03T18:29:53.342511Z", + "shell.execute_reply": "2026-11-03T18:29:53.342120Z" + } + }, + "outputs": [], + "source": [ + "# Licensed to the Apache Software Foundation (ASF) under one\n", + "# or more contributor license agreements. See the NOTICE file\n", + "# distributed with this work for additional information\n", + "# regarding copyright ownership. The ASF licenses this file\n", + "# to you under the Apache License, Version 2.0 (the\n", + "# \"License\"); you may not use this file except in compliance\n", + "# with the License. You may obtain a copy of the License at\n", + "#\n", + "# http://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing,\n", + "# software distributed under the License is distributed on an\n", + "# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n", + "# KIND, either express or implied. See the License for the\n", + "# specific language governing permissions and limitations\n", + "# under the License." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ab14c810", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:53.344766Z", + "iopub.status.busy": "2026-11-03T18:29:53.344598Z", + "iopub.status.idle": "2026-11-03T18:29:56.705276Z", + "shell.execute_reply": "2026-11-03T18:29:56.704928Z" + } + }, + "outputs": [], + "source": [ + "!pip install apache-burr" + ] + }, + { + "cell_type": "markdown", + "id": "5f51adcc", + "metadata": {}, + "source": [ + "# Stopping a streaming action early\n", + "\n", + "When you stop reading a stream before it ends, you want the action's cleanup to run *now*, the step to be reported as failed (so resuming re-runs it), and the state not to be updated. Each demo below prints a timeline of what happened, in order. The code is in [application.py](application.py), and the explanation is in [README.md](README.md)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b3869f03", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:56.706970Z", + "iopub.status.busy": "2026-11-03T18:29:56.706861Z", + "iopub.status.idle": "2026-11-03T18:29:57.083361Z", + "shell.execute_reply": "2026-11-03T18:29:57.083012Z" + } + }, + "outputs": [], + "source": [ + "import asyncio\n", + "\n", + "import application as demo" + ] + }, + { + "cell_type": "markdown", + "id": "01f71955", + "metadata": {}, + "source": [ + "## Pattern 1: `break`, then `container.close()` / `await container.aclose()`" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "aef6f0e1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:57.085325Z", + "iopub.status.busy": "2026-11-03T18:29:57.085085Z", + "iopub.status.idle": "2026-11-03T18:29:57.087941Z", + "shell.execute_reply": "2026-11-03T18:29:57.087560Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Pattern 1: break, then container.close()\n", + " respond: open LLM connection\n", + " caller: got 'Burr '\n", + " caller: got 'streams '\n", + " caller: got 'tokens '\n", + " post_run_step: respond failed with GeneratorExit()\n", + " respond: close LLM connection\n", + " caller: closed the container\n", + " caller: get() -> result=None, 'response' in state=False\n" + ] + } + ], + "source": [ + "demo.stop_with_close_demo()" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "30199433", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:57.089611Z", + "iopub.status.busy": "2026-11-03T18:29:57.089482Z", + "iopub.status.idle": "2026-11-03T18:29:57.091864Z", + "shell.execute_reply": "2026-11-03T18:29:57.091625Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Pattern 1 (async): break, then await container.aclose()\n", + " respond: open LLM connection\n", + " caller: got 'Burr '\n", + " caller: got 'streams '\n", + " caller: got 'tokens '\n", + " respond: close LLM connection\n", + " post_run_step: respond failed with GeneratorExit()\n", + " caller: closed the container\n" + ] + } + ], + "source": [ + "await demo.stop_with_aclose_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "9540ca87", + "metadata": {}, + "source": [ + "## Pattern 2: `contextlib.closing(app.stream_iterate(...))`\n", + "\n", + "However you leave the block, the in-flight action is closed." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "170e909a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:57.093140Z", + "iopub.status.busy": "2026-11-03T18:29:57.093026Z", + "iopub.status.idle": "2026-11-03T18:29:57.094760Z", + "shell.execute_reply": "2026-11-03T18:29:57.094556Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Pattern 2: contextlib.closing(app.stream_iterate(...))\n", + " respond: open LLM connection\n", + " caller: got 'Burr '\n", + " caller: got 'streams '\n", + " caller: got 'tokens '\n", + " post_run_step: respond failed with GeneratorExit()\n", + " respond: close LLM connection\n", + " caller: left the `with` block\n" + ] + } + ], + "source": [ + "demo.stop_iterate_with_closing_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "ef962268", + "metadata": {}, + "source": [ + "## Pattern 3: a client disconnect cancels the consuming task\n", + "\n", + "The hook receives `CancelledError`. See [server.py](server.py) for the FastAPI version." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "4ce17cbd", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:57.095983Z", + "iopub.status.busy": "2026-11-03T18:29:57.095913Z", + "iopub.status.idle": "2026-11-03T18:29:57.219020Z", + "shell.execute_reply": "2026-11-03T18:29:57.218734Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Pattern 3: client disconnect, cancel the consuming task\n", + " respond: open LLM connection\n", + " server: sent 'Burr '\n", + " server: sent 'streams '\n", + " server: client disconnected -> task.cancel()\n", + " respond: close LLM connection\n", + " post_run_step: respond failed with CancelledError()\n" + ] + } + ], + "source": [ + "await demo.client_disconnect_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "1869862c", + "metadata": {}, + "source": [ + "## Persistence: a stopped step is saved as `failed`, and resuming re-runs it" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "4fa878d4", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:57.220438Z", + "iopub.status.busy": "2026-11-03T18:29:57.220348Z", + "iopub.status.idle": "2026-11-03T18:29:57.222888Z", + "shell.execute_reply": "2026-11-03T18:29:57.222643Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Persistence: a stopped step is 'failed' and is re-run on resume\n", + " respond: open LLM connection\n", + " post_run_step: respond failed with GeneratorExit()\n", + " respond: close LLM connection\n", + " persister: last step saved as 'failed'\n", + " resume: next action is 'respond' (it runs again)\n" + ] + } + ], + "source": [ + "demo.resume_after_stop_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "7469b4f3", + "metadata": {}, + "source": [ + "## Cleanup errors are raised by `close()`" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "e307a750", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T18:29:57.224179Z", + "iopub.status.busy": "2026-11-03T18:29:57.224097Z", + "iopub.status.idle": "2026-11-03T18:29:57.225865Z", + "shell.execute_reply": "2026-11-03T18:29:57.225642Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Cleanup errors: close() raises them to the caller\n", + " respond: open LLM connection\n", + " caller: got 'Burr '\n", + " caller: got 'streams '\n", + " caller: got 'tokens '\n", + " post_run_step: respond failed with GeneratorExit()\n", + " respond: close LLM connection\n", + " caller: close() raised ConnectionError('LLM connection did not close cleanly')\n" + ] + } + ], + "source": [ + "demo.cleanup_error_demo()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/streaming-early-stop/server.py b/examples/streaming-early-stop/server.py new file mode 100644 index 000000000..60bfd7f30 --- /dev/null +++ b/examples/streaming-early-stop/server.py @@ -0,0 +1,86 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Client disconnects in a streaming FastAPI endpoint. + +Run with ``uvicorn server:app``, then ``curl -N --max-time 1 localhost:8000/chat`` (curl hangs up +after a second, like a user closing the tab). The server log shows the LLM connection being closed +right away and the step reported to hooks as failed with ``CancelledError``. +""" + +import asyncio +import contextlib +import logging +from typing import AsyncGenerator, Optional, Tuple + +from fastapi import FastAPI +from starlette.responses import StreamingResponse + +from burr.core import ApplicationBuilder, State +from burr.core.action import streaming_action +from burr.lifecycle import PostRunStepHook + +logger = logging.getLogger("uvicorn.error") +app = FastAPI() + + +@streaming_action(reads=[], writes=["response"]) +async def respond(state: State) -> AsyncGenerator[Tuple[dict, Optional[State]], None]: + logger.info("respond: open LLM connection") + try: + response = "" + for i in range(50): + await asyncio.sleep(0.1) # waiting for the next chunk from the LLM + token = f"token-{i} " + response += token + yield {"token": token}, None + yield {"token": ""}, state.update(response=response) + finally: + logger.info("respond: close LLM connection") + + +class LogStep(PostRunStepHook): + # A *sync* hook: see the README for why async post-step hooks need care on disconnect. + def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): + logger.info("post_run_step: %s exception=%r", action.name, exception) + + +def chat_application(): + return ( + ApplicationBuilder() + .with_actions(respond=respond) + .with_transitions() + .with_entrypoint("respond") + .with_state(State()) + .with_hooks(LogStep()) + .build() + ) + + +@app.get("/chat") +async def chat() -> StreamingResponse: + burr_app = chat_application() + + async def sse(): + # `aclosing` makes leaving this block -- normally, on error, or because the server + # cancelled us when the client went away -- close the in-flight action immediately. + async with contextlib.aclosing(burr_app.astream_iterate(halt_after=["respond"])) as steps: + async for _, container in steps: + async for item in container: + yield f"data: {item['token']}\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") diff --git a/examples/streaming-early-stop/statemachine.png b/examples/streaming-early-stop/statemachine.png new file mode 100644 index 0000000000000000000000000000000000000000..1b88069e340e6e293a29a9704f5dcfdbd7f1944c GIT binary patch literal 3022 zcmV;<3o-PGP)f z^*!I`j&sgE7bO%5;R}5S{iy{ERlXS4Prew}PcZ!}*4EZWr_*a|YuRk}zzUKXSS*%6 zAdpBTZEfxTlo$4e`uh4qhYro0>43rX({VXOPha2P-=9LE^vMcTFQ2L$4u?!8lVf91 zX!O)MPS!JKk2Nz_)6!NNt~}Ts8zlIAesd$Myts(|^7r2#KVUNG_V&{f5)MzAG^sbm zdu7~>8#jW3L+V*9w@pC{JiLdlSv|nWeR(!+e|}DupPwI@OePQrJ?H-ye2gRrg_UfM4yq@*6i&92^`zmhGd)2?T-_Zf=hrr|md$V#4GpgQim+Y3Fj| zBZ6}>o@P9MZe?ZFZJv)Bw|)Ee*n5+M_9a*dG*$=TB-i-U0n?T8alc%6*Ux4)7Gi1suGWMGOnei1px3$ z!(=LID4;lOcyn{Jc%+kY$Y2 z$`N;Eow=Pm%0YFPr)7c*m zyv};j#A1!Np5)}ZaW)iXUDo=e{#nm?$1FY>v%83v$K|vbnUGvJY@Rl60g{xsGaQSmt5e;=2w2j_%r?cJrF!Vwcfo<~482 z?*DwTtf=7N<(mNT2PLC~{_6UPqZU)9eeLFzo0)zlF1o6`G-z)u0OV&=cdqqPRa2ku zz6z(PNPY5ff6#_iL0dfnLIB_mgWkkqQ6Amb(l?yxyreKM=TXWP4!dc4;&A|=Q7PYV z@WJ8n_VebWQ0ViAW0gmY{KUoqKt(Cjdb&g856ARJj|G4;anbK9-^E?IMY6O40I$HU zQR`QoKDf_usjH&Wu&1flZ5-x=?2kvvY8#G9-uGQ)*&Dr4qXD4$ZTY&$sHLmd1AupM z_>MIm53gTd;ORZa%pz%D6b_G%z4QxFQwsoG)&)eZ^SXTO@QekEk+SbA-?^?2@D1ID zLZO5LL5Qn!-ir*RmE`CsEDm=l<(`(F0RSvt>mRseo|L_Y0lCL{2Y-XO#%^p#zx|7e zg(cz|0AR3Krxh#tJl=1=-V+Eq0DwlNF!Ns{sf)Z;U&}4hF)%_7hpVU{Cnf-(&}g@f zfkL4$H}e@YpT^9mO`JLnaSZ@47|fhyZhRguH}e@1!6_=P+#H6CKv8J4nf2rrHoLB- zx*#v7v>@Mkr7z+d0O%SSJ1%jNzSq*8aU+RpCKi^+|H@LPP$=YbI5Ar`i864S*#J=U zt_-iJ}_a50)XnbI(Vb`CNX=3r?HPx~;Y9zN8dyj~ z0fQ0!MyL{IEm~$cmIMIq{z1oo&GZNe!7D1=y?lO`zi)uEBdhKMQkNdH~S;=afO^(7&CJW7W>C<&&1oJQ7LEz1?%avdX{8jJ`sb#yn2x-iga}F({KMGO*at+ z(l`tj>l+%$X=cZT2Qu?%)o;sgoKCuO?68&n%<*zXrtu8AnIJ{#kGoQxHe;pmTp@rQv0G*Y-3UTJSdbV-7 zr(^d%zjq4&6cjM?R(Sesi9n%HO2d`6#wNx`hHUfs`&}&#kN4RUVKeJ*<1Fo_t7&RS z2X8>;1{St<0a4LN((2G{f{u>Vb0_a!I)_B&dw6f$vAZ{tR)=lpwsIa^z4-9@Wu%pj z-|mU~CP>Pjc6zUR_3Bl3clX-`6;f8(+S1I-r*k2${7a5w?R@dHTQ%Y!=A^ROl zKi+yQ;}&}lrZd=M~^or zYH1Hs9;&;A`8*!Gp`n6V#AM_H2p29|6c!rVb-ea5`d}P_Kp@D<%1TR1%gf8FdH=qt zX#kd?KM7bYmPjC)nv!g7ZJnH)BxY3qN}rAU0>aSStt2OWF|MC{F|MEd1?4KzZvP;T Q761SM07*qoM6N<$f<9H@ZU6uP literal 0 HcmV?d00001 diff --git a/pyproject.toml b/pyproject.toml index 07be284f1..f85c2b3f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -330,6 +330,7 @@ exclude = [ "examples/recursive/**", "examples/simple-chatbot-intro/**", "examples/simulation/**", + "examples/streaming-early-stop/**", "examples/streaming-overview/**", "examples/talks/**", "examples/templates/**", diff --git a/tests/core/test_streaming_early_stop_example.py b/tests/core/test_streaming_early_stop_example.py new file mode 100644 index 000000000..d781af506 --- /dev/null +++ b/tests/core/test_streaming_early_stop_example.py @@ -0,0 +1,90 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Keeps examples/streaming-early-stop honest: each demo must show what its README claims.""" + +import importlib.util +import pathlib + +import pytest + +EXAMPLE = pathlib.Path(__file__).parents[2] / "examples" / "streaming-early-stop" / "application.py" + + +@pytest.fixture(scope="module") +def demo(): + spec = importlib.util.spec_from_file_location("streaming_early_stop_example", EXAMPLE) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _timeline(capsys) -> list[str]: + return [line.strip() for line in capsys.readouterr().out.splitlines() if line.startswith(" ")] + + +def test_stop_with_close(demo, capsys): + demo.stop_with_close_demo() + events = _timeline(capsys) + assert "respond: close LLM connection" in events + assert "post_run_step: respond failed with GeneratorExit()" in events + assert events.index("respond: close LLM connection") < events.index( + "caller: closed the container" + ) + assert events[-1] == "caller: get() -> result=None, 'response' in state=False" + + +async def test_stop_with_aclose(demo, capsys): + await demo.stop_with_aclose_demo() + events = _timeline(capsys) + assert events.index("respond: close LLM connection") < events.index( + "caller: closed the container" + ) + assert "post_run_step: respond failed with GeneratorExit()" in events + + +def test_stop_iterate_with_closing(demo, capsys): + demo.stop_iterate_with_closing_demo() + events = _timeline(capsys) + assert events.index("respond: close LLM connection") < events.index( + "caller: left the `with` block" + ) + assert "post_run_step: respond failed with GeneratorExit()" in events + + +async def test_client_disconnect(demo, capsys): + await demo.client_disconnect_demo() + events = _timeline(capsys) + assert events[-2:] == [ + "respond: close LLM connection", + "post_run_step: respond failed with CancelledError()", + ] + + +def test_resume_after_stop(demo, capsys): + demo.resume_after_stop_demo() + events = _timeline(capsys) + assert "persister: last step saved as 'failed'" in events + assert events[-1] == "resume: next action is 'respond' (it runs again)" + + +def test_cleanup_error(demo, capsys): + demo.cleanup_error_demo() + events = _timeline(capsys) + assert events[-1] == ( + "caller: close() raised ConnectionError('LLM connection did not close cleanly')" + ) From 2d5734afcdb5086eba57da7c15a4cd062dfd9a12 Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Tue, 3 Nov 2026 12:11:24 -0800 Subject: [PATCH 6/8] fix(core): run cleanup before reporting on every early-stop path Make stopping a stream early behave the same however it happens: the action's cleanup runs first, then the step is reported to hooks and persisters exactly once, as failed, with no state update -- independent of timing. What users write is plain Python: sync code just `break`s; async code `break`s inside `async with contextlib.aclosing(...)` (or calls `aclose()`). - `for ... break` (or dropping the iterator) now closes the action before the post-step callback runs, instead of reporting while the action's resource is still open and leaving cleanup to GC. This also removes the container <-> iterator cycle, so the weakref workaround is gone. - async: `aclose()` racing the event loop's async-generator finalizer (after `async for ... break`) waits for that close and gets its outcome, so cleanup errors are no longer lost or reported after the hook, whatever the number of awaits in between. - Cleanup errors: hooks receive them; `close()`/`aclose()` raise them, also on repeated calls. After a bare `break` there is no caller to raise to, so they are logged at warning. - First error wins: a cleanup error raised while another error is already stopping the stream no longer replaces it; it is logged at warning (shared `_close_action_generator` helpers in the runners). - `close()`/`aclose()` while the stream is producing an item (from a hook, or from another task) raises a RuntimeError that says to cancel the consuming task, instead of recording a bogus failure while the stream carries on (and possibly commits). - Hook/callback types: `exception` is `Optional[BaseException]` (GeneratorExit/CancelledError are BaseExceptions). Tests: the early-stop tests in test_application.py are consolidated into one section on shared scaffolding (sync + async, single- and multi-step) and cover the timing race, cleanup errors, first-error-wins, cancellation and close-while-running. Docs and the streaming-early-stop example now teach "sync: break; async: aclosing; servers: let cancellation do it". --- burr/core/action.py | 127 ++- burr/core/application.py | 55 +- burr/lifecycle/base.py | 16 +- burr/tracking/client.py | 2 +- docs/concepts/streaming-actions.rst | 49 +- examples/streaming-early-stop/README.md | 113 +-- examples/streaming-early-stop/application.py | 100 ++- examples/streaming-early-stop/notebook.ipynb | 230 ++++-- examples/streaming-early-stop/server.py | 2 +- tests/core/test_application.py | 749 +++++++++--------- .../core/test_streaming_early_stop_example.py | 63 +- 11 files changed, 887 insertions(+), 619 deletions(-) diff --git a/burr/core/action.py b/burr/core/action.py index e8eb99897..0d49e878c 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -17,15 +17,16 @@ import abc import ast +import asyncio import builtins import contextlib import copy import inspect +import logging import sys import textwrap import types import typing -import weakref UnionType = types.UnionType @@ -56,6 +57,8 @@ from burr.core.state import State +logger = logging.getLogger(__name__) + def _validate_declared_reads(fn: Callable, declared_reads: list[str]) -> None: if not declared_reads: @@ -1426,6 +1429,30 @@ def is_async(self) -> bool: StreamResultType = TypeVar("StreamResultType") +def _resolve_close_future(future: "asyncio.Future", exc: Optional[BaseException]) -> None: + # A close interrupted by cancellation is not an error to hand to later callers. + if exc is not None and not isinstance(exc, asyncio.CancelledError): + future.set_exception(exc) + future.exception() # mark retrieved: the first closer already raised it + else: + future.set_result(None) + + +def _log_abandoned_cleanup_error() -> None: + # Nobody is left to raise this to: the iterator is being closed by ``break`` or garbage + # collection. Hooks have already seen it, and ``close()``/``aclose()`` will raise it. + logger.warning( + "Cleanup of a streaming action raised after the stream was stopped early", exc_info=True + ) + + +_CLOSE_WHILE_RUNNING = ( + "Cannot close a streaming result while it is producing an item (closing from a hook, or from " + "another task). Stop it from the code that consumes it (break out of the loop), or cancel the " + "task that consumes it." +) + + class StreamingResultContainer(Generic[StateType, StreamResultType], Iterator[StreamResultType]): """Container for a streaming result. This allows you to: @@ -1493,9 +1520,9 @@ def __init__( self._initial_state = initial_state self._result = None self._callback_realized = False - # Weak, so the container <-> iterator cycle doesn't defer the iterator's finalizer (and so - # the post-step callback) to cyclic GC when a caller abandons the stream without closing it. - self._iterator_ref: Optional[weakref.ref] = None + self._closed = False + # Cleanup error from the first close; later ``close()`` calls raise it again. + self._close_error: Optional[BaseException] = None def _realize_callback(self, exc: Optional[BaseException]) -> None: if self._result is None: @@ -1510,23 +1537,33 @@ def close(self) -> None: Closes the action's generator chain so its cleanup (``finally`` blocks, ``with`` exits) runs now rather than at garbage collection, then runs the post-step callback exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a - no-op once the stream has completed. Exceptions raised by the action's cleanup propagate. + no-op once the stream has completed. An exception raised by the action's cleanup + propagates, and is raised again by later calls -- including after ``for ... break`` + already closed the stream. A stream closed before it completes is reported to the callback (and so to ``post_run_step`` hooks and persisters) with a ``GeneratorExit``, i.e. as a failed step. + + Raises ``RuntimeError`` if the stream is producing an item at that moment (e.g. ``close`` + is called from a hook, or from another task). To stop a stream from another task, cancel + the task that is consuming it. """ + if self._closed: + if self._close_error is not None: + raise self._close_error + return + if self.streaming_result_generator.gi_running: + raise RuntimeError(_CLOSE_WHILE_RUNNING) + self._closed = True # Closing before completion means the step did not finish: report it as failed. exc: Optional[BaseException] = GeneratorExit() if self._result is None else None try: self.streaming_result_generator.close() except BaseException as e: - exc = e + exc = self._close_error = e raise finally: self._realize_callback(exc) - iterator = self._iterator_ref() if self._iterator_ref is not None else None - if iterator is not None: - iterator.close() def __next__(self) -> StreamResultType: if self._result is not None: @@ -1546,14 +1583,21 @@ def gen_fn(): yield out except StopIteration: return + except GeneratorExit: + # The caller stopped iterating (``break``, or dropped the iterator): unwind the + # action *before* reporting the step. + if not self._closed: + try: + self.close() + except Exception: + _log_abandoned_cleanup_error() + raise finally: self._realize_callback(sys.exc_info()[1]) # We really don't need an internal generator function but this was done to keep it the same # as the async version - iterator = gen_fn() - self._iterator_ref = weakref.ref(iterator) - return iterator + return gen_fn() def get(self) -> Tuple[StreamResultType, State[StateType]]: # exhaust the generator @@ -1616,9 +1660,11 @@ def __init__( self._callback = callback self._result = None self._callback_realized = False - # Weak, so the container <-> iterator cycle doesn't defer the iterator's finalizer (and so - # the post-step callback) to cyclic GC when a caller abandons the stream without closing it. - self._iterator_ref: Optional[weakref.ref] = None + # True while ``__anext__`` is waiting on the action for the next item. + self._producing = False + # Set by the first ``aclose()`` (explicit, or the iterator's finalizer); later calls wait on + # it, so the action is closed once and every caller sees the same outcome. + self._close_future: Optional[asyncio.Future] = None async def _realize_callback(self, exc: Optional[BaseException]) -> None: if self._result is None: @@ -1633,30 +1679,49 @@ async def aclose(self) -> None: Closes the action's generator chain so its cleanup (``finally`` blocks, ``async with`` exits) runs now rather than at garbage collection, then runs the post-step callback exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a - no-op once the stream has completed. Exceptions raised by the action's cleanup propagate. + no-op once the stream has completed. An exception raised by the action's cleanup + propagates, and is raised again by later calls. A stream closed before it completes is reported to the callback (and so to ``post_run_step`` hooks and persisters) with a ``GeneratorExit``, i.e. as a failed step. + + Safe to call while the iterator's own finalizer is closing the stream (e.g. after + ``async for ... break``): it waits for that close and returns, or raises, the same outcome. + + Raises ``RuntimeError`` if the stream is producing an item at that moment (e.g. ``aclose`` + is called from a hook, or from another task). To stop a stream from another task, cancel + the task that is consuming it. """ + if self._close_future is not None: + await asyncio.shield(self._close_future) + return + if self._producing: + raise RuntimeError(_CLOSE_WHILE_RUNNING) + self._close_future = asyncio.get_running_loop().create_future() # Closing before completion means the step did not finish: report it as failed. exc: Optional[BaseException] = GeneratorExit() if self._result is None else None + raised = False try: await self.streaming_result_generator.aclose() except BaseException as e: - exc = e + exc, raised = e, True raise finally: - await self._realize_callback(exc) - iterator = self._iterator_ref() if self._iterator_ref is not None else None - if iterator is not None: - await iterator.aclose() + try: + await self._realize_callback(exc) + finally: + _resolve_close_future(self._close_future, exc if raised else None) async def __anext__(self) -> StreamResultType: """Moves to the next state in the streaming result""" if self._result is not None: # we're done, and we've run through it raise StopAsyncIteration - result, state = await self.streaming_result_generator.__anext__() + self._producing = True + try: + result, state = await self.streaming_result_generator.__anext__() + finally: + self._producing = False if state is not None: # we're done -- we've hit the last one self._result = self._process_result(result, state) raise StopAsyncIteration @@ -1673,13 +1738,23 @@ async def gen_fn(): yield await self.__anext__() except StopAsyncIteration: return + except GeneratorExit: + # The caller stopped iterating (``break``, or dropped the iterator): unwind the + # action *before* reporting the step. If an explicit ``aclose()`` got there + # first, wait for it so the step isn't reported while cleanup is still running. + if self._close_future is None: + try: + await self.aclose() + except Exception: + _log_abandoned_cleanup_error() + else: + await asyncio.wait([self._close_future]) + raise finally: await self._realize_callback(sys.exc_info()[1]) # return it as `__aiter__` cannot be async/have awaits :/ - iterator = gen_fn() - self._iterator_ref = weakref.ref(iterator) - return iterator + return gen_fn() async def get(self) -> tuple[Optional[StreamResultType], State[StateType]]: # exhaust the generator @@ -1699,7 +1774,7 @@ async def just_results() -> AsyncGeneratorReturnType: yield results, final_state async def empty_callback( - result: Optional[StreamResultType], state: State, exc: Optional[Exception] + result: Optional[StreamResultType], state: State, exc: Optional[BaseException] ): pass diff --git a/burr/core/application.py b/burr/core/application.py index 7f6f4a7aa..e7139e919 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -23,6 +23,7 @@ import inspect import logging import pprint +import sys import uuid from concurrent.futures import Executor, ThreadPoolExecutor from contextlib import AbstractContextManager @@ -328,6 +329,44 @@ def _run_single_step_action( return out +def _close_action_generator( + generator: Generator, action_name: str, in_flight: Optional[BaseException] +) -> None: + """Closes an action's generator so its cleanup runs now (early close, or an error), rather than + at garbage collection. A cleanup error must not replace an error that is already propagating: + that cleanup error is logged instead.""" + if in_flight is None or isinstance(in_flight, GeneratorExit): + generator.close() + return + try: + generator.close() + except Exception: + logger.warning( + "Cleanup of streaming action '%s' raised while handling %r", + action_name, + in_flight, + exc_info=True, + ) + + +async def _aclose_action_generator( + generator: AsyncGenerator, action_name: str, in_flight: Optional[BaseException] +) -> None: + """Async version of :py:func:`_close_action_generator`.""" + if in_flight is None or isinstance(in_flight, GeneratorExit): + await generator.aclose() + return + try: + await generator.aclose() + except Exception: + logger.warning( + "Cleanup of streaming action '%s' raised while handling %r", + action_name, + in_flight, + exc_info=True, + ) + + def _run_single_step_streaming_action( action: SingleStepStreamingAction, state: State, @@ -388,8 +427,7 @@ def _run_single_step_streaming_action( exc_info=True, ) finally: - # Unwind the action now on early close instead of leaving it to GC. - generator.close() + _close_action_generator(generator, action.name, sys.exc_info()[1]) if state_update is None: raise ValueError( @@ -460,8 +498,7 @@ async def _arun_single_step_streaming_action( exc_info=True, ) finally: - # Unwind the action now on early close instead of leaving it to GC. - await generator.aclose() + await _aclose_action_generator(generator, action.name, sys.exc_info()[1]) if state_update is None: raise ValueError( @@ -534,8 +571,7 @@ def _run_multi_step_streaming_action( exc_info=True, ) finally: - # Unwind the action now on early close instead of leaving it to GC. - generator.close() + _close_action_generator(generator, action.name, sys.exc_info()[1]) state_update = _run_reducer(action, state, result, action.name) _validate_result(result, action.name, action.schema) _validate_reducer_writes(action, state_update, action.name) @@ -595,8 +631,7 @@ async def _arun_multi_step_streaming_action( exc_info=True, ) finally: - # Unwind the action now on early close instead of leaving it to GC. - await generator.aclose() + await _aclose_action_generator(generator, action.name, sys.exc_info()[1]) state_update = _run_reducer(action, state, result, action.name) _validate_result(result, action.name, action.schema) _validate_reducer_writes(action, state_update, action.name) @@ -1579,7 +1614,7 @@ def process_result( def callback( result: Optional[dict], state: State[ApplicationStateType], - exc: Optional[Exception] = None, + exc: Optional[BaseException] = None, ): self._adapter_set.call_all_lifecycle_hooks_sync( "post_run_step", @@ -1835,7 +1870,7 @@ def process_result( async def callback( result: Optional[dict], state: State[ApplicationStateType], - exc: Optional[Exception] = None, + exc: Optional[BaseException] = None, ): await self._adapter_set.call_all_lifecycle_hooks_sync_and_async( "post_run_step", diff --git a/burr/lifecycle/base.py b/burr/lifecycle/base.py index 66d8bd7e6..61f0fcb62 100644 --- a/burr/lifecycle/base.py +++ b/burr/lifecycle/base.py @@ -98,7 +98,7 @@ def post_run_step( state: "State", action: "Action", result: Optional[Dict[str, Any]], - exception: Exception, + exception: Optional[BaseException], **future_kwargs: Any, ): """Run after a step is executed. @@ -107,7 +107,9 @@ def post_run_step( :param action: Action that was executed :param result: Result of the action :param sequence_id: Sequence ID of the action - :param exception: Exception that was raised + :param exception: Exception that was raised, or ``None`` if the step completed. A streaming + step stopped early is reported with ``GeneratorExit`` (closed) or ``CancelledError`` + (its task was cancelled). :param future_kwargs: Future keyword arguments """ pass @@ -127,7 +129,7 @@ async def post_run_step( state: "State", action: "Action", result: Optional[dict], - exception: Exception, + exception: Optional[BaseException], **future_kwargs: Any, ): """Async run after a step is executed @@ -136,7 +138,9 @@ async def post_run_step( :param action: Action that was executed :param result: Result of the action :param sequence_id: Sequence ID of the action - :param exception: Exception that was raised + :param exception: Exception that was raised, or ``None`` if the step completed. A streaming + step stopped early is reported with ``GeneratorExit`` (closed) or ``CancelledError`` + (its task was cancelled). :param future_kwargs: Future keyword arguments """ pass @@ -354,7 +358,7 @@ def post_run_execute_call( partition_key: str, state: "State", method: ExecuteMethod, - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs, ): pass @@ -372,7 +376,7 @@ async def post_run_execute_call( partition_key: str, state: "State", method: ExecuteMethod, - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs, ): pass diff --git a/burr/tracking/client.py b/burr/tracking/client.py index 5375cdb4b..f4335661f 100644 --- a/burr/tracking/client.py +++ b/burr/tracking/client.py @@ -95,7 +95,7 @@ def flock(*args, **kwargs): ) -def _format_exception(exception: Exception) -> Optional[str]: +def _format_exception(exception: Optional[BaseException]) -> Optional[str]: if exception is None: return None return "".join(traceback.format_exception(type(exception), exception, exception.__traceback__)) diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index 4e30b8ee8..a11fbb7f0 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -227,10 +227,11 @@ Thus you can run this in a web-service, a streamlit app, etc... Stopping a stream early ~~~~~~~~~~~~~~~~~~~~~~~ -If you stop consuming a stream before it finishes (e.g. the client disconnected), close it so the -action's cleanup (``finally`` blocks, ``with``/``async with`` exits such as an open HTTP stream) runs right away, -rather than at garbage collection. Either close the container directly (``close()`` / ``await aclose()``), or -close ``stream_iterate`` / ``astream_iterate`` (``contextlib.closing`` / ``contextlib.aclosing`` does this for you): +When you stop reading a stream before it finishes (the user hit stop, a check tripped, the client +disconnected), Burr stops the action right away: its cleanup (``finally`` blocks, ``with``/``async with`` +exits such as an open HTTP stream) runs, then the step is reported. You don't need to drain the stream. + +**Sync: just** ``break``. This works the same for ``stream_result`` and ``stream_iterate``: .. code-block:: python @@ -239,43 +240,41 @@ close ``stream_iterate`` / ``astream_iterate`` (``contextlib.closing`` / ``conte ) for result in streaming_result: if should_stop(result): - break - streaming_result.close() + break # the action's cleanup runs here + +**Async:** ``break`` **inside** ``contextlib.aclosing``. Python never closes an async generator synchronously +when you ``break``; it leaves that to the event loop, at some later point. ``aclosing`` makes it happen as you +leave the block: .. code-block:: python import contextlib - with contextlib.closing(application.stream_iterate(halt_after=["respond"])) as steps: - for action, streaming_result in steps: - for item in streaming_result: - ... - break # leaving the ``with`` closes the in-flight action - -.. code-block:: python - action, async_streaming_result = await application.astream_result( halt_after='streaming_response', inputs={"prompt": prompt} ) - async for result in async_streaming_result: - if should_stop(result): - break - await async_streaming_result.aclose() - -.. code-block:: python - - import contextlib + async with contextlib.aclosing(async_streaming_result): + async for result in async_streaming_result: + if should_stop(result): + break async with contextlib.aclosing(application.astream_iterate(halt_after=["respond"])) as steps: async for action, streaming_result in steps: async for item in streaming_result: ... - break # leaving the ``async with`` closes the in-flight action + +``await async_streaming_result.aclose()`` (or ``streaming_result.close()``) does the same explicitly. + +**Web servers:** use the async pattern in your response generator. When a client disconnects, the server +cancels the task that is producing the response, and the cancellation unwinds the action. To stop a stream from +*another* task, cancel the task that consumes it: closing a stream while it is producing an item (from a hook, +or from another task) raises a ``RuntimeError``. A step that is stopped before it completes does not update state, and is reported to ``post_run_step`` hooks with the exception that stopped it -- ``GeneratorExit`` when it is closed, ``CancelledError`` when the consuming -task is cancelled (e.g. a web client disconnected). Persisters record it as ``failed``, so resuming re-runs that -action. Exceptions raised by the action's cleanup propagate out of ``close()`` / ``aclose()``. +task is cancelled. Persisters record it as ``failed``, so resuming re-runs that action. If the action's cleanup +itself raises, hooks receive that error, and ``close()`` / ``aclose()`` raise it. A cleanup error never replaces +an error that was already stopping the stream; it is logged instead. See the `streaming-early-stop example `_ for runnable versions of these patterns, including a FastAPI endpoint that handles client disconnects. diff --git a/examples/streaming-early-stop/README.md b/examples/streaming-early-stop/README.md index 7ba3e49d4..631e1a43d 100644 --- a/examples/streaming-early-stop/README.md +++ b/examples/streaming-early-stop/README.md @@ -22,16 +22,23 @@ Streaming actions often hold something open while they stream, such as an HTTP connection to an LLM, a database cursor, or a tracing span. They release it in a `finally` block or a `with` exit. When you stop reading a stream before it ends (the user hits "stop", a moderation check trips, a token budget runs out, -the client disconnects), you want three things: +the client disconnects), Burr: -1. **The action's cleanup runs now**, not whenever the garbage collector gets to it. -2. **The step is reported honestly.** Hooks, trackers and persisters see it as *failed*, not completed, so - resuming the application re-runs it. -3. **The state is not updated.** A half-finished response is never committed. +1. **runs the action's cleanup right away**, not whenever the garbage collector gets to it; +2. **then reports the step as failed** to hooks, trackers and persisters, so resuming the application + re-runs it; +3. **doesn't update the state**, so a half-finished response is never committed. -This example shows the patterns that give you all three. Every demo in [application.py](application.py) -prints a timeline of what happened, so you can see when the resource is released and what Burr -reports. Run it with: +All you have to do is stop iterating the way you would with any Python generator: + +| | do this | +|---|---| +| sync | `break` | +| async | `break` inside `async with contextlib.aclosing(...)` | +| web server | the async pattern; the server cancels your task when the client disconnects | + +Every demo in [application.py](application.py) prints a timeline of what happened, so you can see +when the resource is released and what Burr reports. Run it with: ```bash python application.py @@ -39,54 +46,65 @@ python application.py ![statemachine](statemachine.png) -## The patterns - -### 1. Stop reading, then `close()` the container +## Sync: just `break` ```python action, container = app.stream_result(halt_after=["respond"]) for item in container: send_to_user(item["token"]) if wants_to_stop(item): - break -container.close() # async: await container.aclose() + break # the action's cleanup runs here, then the step is reported result, state = container.get() # (None, state without the update) ``` -`close()` / `aclose()` stops the action and runs its cleanup. Calling it more than once is safe. Once -the stream has finished normally, it does nothing (the step is still reported as completed). +``` +respond: open LLM connection +caller: got 'Burr ' +caller: got 'streams ' +caller: got 'tokens ' +respond: close LLM connection <- at the break +post_run_step: respond failed with GeneratorExit() +caller: after the loop +``` -**Always call it after a `break`.** A `break` on its own stops *you* from reading, but the action's -cleanup only runs when the container is closed or garbage-collected. +Breaking out of `app.stream_iterate(...)` works the same way. -### 2. Use `contextlib.closing` / `contextlib.aclosing` around `stream_iterate` +## Async: `break` inside `contextlib.aclosing` ```python import contextlib -with contextlib.closing(app.stream_iterate(halt_after=["respond"])) as steps: - for action, container in steps: - for item in container: - ... - if wants_to_stop(item): - break - break -# leaving the block closed the in-flight action +action, container = await app.astream_result(halt_after=["respond"]) +async with contextlib.aclosing(container): + async for item in container: + if wants_to_stop(item): + break +# cleanup has run and the step has been reported async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: async for action, container in steps: - ... + async for item in container: + ... +``` + +Why the extra `async with`? Python can't run an async generator's cleanup at the `break`, because the +cleanup may need to `await`. Without `aclosing`, Python hands it to the event loop, which runs it at some +later point, after your code has moved on (`async_break_without_aclose_demo`): + +``` +respond: open LLM connection +caller: after the loop -- nothing has been cleaned up yet +respond: close LLM connection <- whenever the event loop gets to it +post_run_step: respond failed with GeneratorExit() ``` -This is the most robust option: however you leave the block (`break`, `return`, an exception, or -the task being cancelled), the in-flight action is closed. **Prefer this whenever you use -`stream_iterate` / `astream_iterate`.** +`await container.aclose()` after the loop does the same as `aclosing`, if you prefer it. -### 3. Client disconnects (web servers): cancellation does the work +## Web servers: client disconnects When a client disconnects, ASGI servers such as Starlette/FastAPI *cancel* the task that is -producing the response. If your response generator uses pattern 2, that cancellation unwinds +producing the response. If your response generator uses the async pattern, that cancellation unwinds the Burr action immediately. See [server.py](server.py): ```python @@ -109,21 +127,22 @@ curl -N --max-time 1 localhost:8000/chat # hang up after a second # post_run_step: respond exception=CancelledError(...) ``` -To stop a stream *from another task*, cancel the consuming task (as the server does). Don't call -`aclose()` on a container that another task is iterating: you can't close an async generator while -it's running. +**To stop a stream from somewhere else** (another task, a "stop" endpoint), cancel the task that is +consuming it. Calling `close()` / `aclose()` while the stream is producing an item (from another task, +or from inside a hook) raises a `RuntimeError` that tells you to do this. ## What gets reported | How the stream ended | `post_run_step(exception=...)` | persister status | state updated? | on resume | |---|---|---|---|---| | read to the end | `None` | `completed` | yes | next action | -| `close()` / `aclose()` / `closing` before the end | `GeneratorExit()` | `failed` | no | **same action runs again** | +| `break` / `close()` / `aclose()` before the end | `GeneratorExit()` | `failed` | no | **same action runs again** | | task cancelled (e.g. client disconnect) | `CancelledError()` | `failed` | no | **same action runs again** | | the action raised | that exception | `failed` | no | same action runs again | +| stopped early, and the action's cleanup raised | the cleanup error | `failed` | no | same action runs again | -So in a `post_run_step` hook, `exception is not None` means "this step did not complete". If you -want to tell a deliberate stop apart from a real failure: +So in a `post_run_step` hook, `exception is not None` means "this step did not complete". To tell a +deliberate stop apart from a real failure: ```python def post_run_step(self, *, exception, **kwargs): @@ -139,16 +158,14 @@ Note that `GeneratorExit` and `CancelledError` derive from `BaseException`, not ## Things to know -- **Cleanup errors.** If the action's `finally` raises while it is being closed, `close()` / `aclose()` - raises that error to you. See `cleanup_error_demo` in [application.py](application.py). -- **Use sync hooks/persisters to record client disconnects.** On a disconnect, Starlette/FastAPI - cancels the whole response task. Sync `post_run_step` hooks (and sync persisters) still run and - see the `CancelledError`. *Async* `post_run_step` hooks (and async persisters) are cancelled - before they start, so they never record the stopped step. Shielding inside the hook does not - help, because the hook body never runs. -- **Sync vs async.** Everything above works the same way for `stream_result`/`stream_iterate` - (with `close()` / `contextlib.closing`) and `astream_result`/`astream_iterate` (with `aclose()` / - `contextlib.aclosing`). +- **Cleanup errors.** If the action's `finally` raises while it is being stopped, hooks receive that error, + and Burr logs a warning. A `break` itself doesn't raise it. If you want the error in your own code, call + `close()` / `aclose()` afterwards and they raise it (`cleanup_error_demo`). A cleanup error never + replaces an error that was already stopping the stream (such as a bug in a hook); it is logged instead. +- **Use sync hooks/persisters to record client disconnects.** On a disconnect, Starlette/FastAPI cancels + the whole response task. Sync `post_run_step` hooks (and sync persisters) still run and see the + `CancelledError`. *Async* `post_run_step` hooks (and async persisters) are cancelled before they start, + so they never record the stopped step. See also the [streaming actions docs](https://burr.apache.org/concepts/streaming-actions/) and the [streaming-fastapi](../streaming-fastapi) example for a full chatbot. diff --git a/examples/streaming-early-stop/application.py b/examples/streaming-early-stop/application.py index 720878641..3bd26cbac 100644 --- a/examples/streaming-early-stop/application.py +++ b/examples/streaming-early-stop/application.py @@ -19,15 +19,19 @@ The application is a single streaming action, ``respond``, that pretends to stream tokens from an LLM. It holds a resource while it streams (think: an open HTTP connection) and releases it in a -``finally`` block. Each ``*_demo`` function below shows one way to stop the stream early, and +``finally`` block. Each ``*_demo`` function below stops the stream early in a different way, and prints what happened, in order, so you can see when the resource is released and what Burr reports to hooks/persisters. +In short: sync code just ``break``s; async code ``break``s inside +``async with contextlib.aclosing(...)``. + Run it with ``python application.py``. See README.md for the explanation. """ import asyncio import contextlib +import logging from typing import AsyncGenerator, Generator, List, Optional, Tuple from burr.core import Application, ApplicationBuilder, State @@ -122,59 +126,84 @@ def wants_to_stop(token: str) -> bool: # -------------------------------------------------------------------------------------------- -# Pattern 1: stop reading, then close() the container. +# Sync: just break. # -------------------------------------------------------------------------------------------- -def stop_with_close_demo(): +def sync_break_demo(): log = Timeline() app = application(_respond_sync(log), log) _, container = app.stream_result(halt_after=["respond"]) for item in container: log(f"caller: got {item['token']!r}") if wants_to_stop(item["token"]): - break - container.close() # releases the action's resources *now* - log("caller: closed the container") + break # the action's cleanup runs here, then the step is reported + log("caller: after the loop") result, state = container.get() # the step did not finish: no result, state not updated log(f"caller: get() -> result={result!r}, 'response' in state={'response' in state}") - log.print("Pattern 1: break, then container.close()") + log.print("Sync: break") -async def stop_with_aclose_demo(): +def sync_iterate_break_demo(): log = Timeline() - app = application(_respond_async(log), log) - _, container = await app.astream_result(halt_after=["respond"]) - async for item in container: - log(f"caller: got {item['token']!r}") - if wants_to_stop(item["token"]): - break - await container.aclose() - log("caller: closed the container") - log.print("Pattern 1 (async): break, then await container.aclose()") + app = application(_respond_sync(log), log) + for _, container in app.stream_iterate(halt_after=["respond"]): + for item in container: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + break + log("caller: after the loop") + log.print("Sync: break out of stream_iterate") # -------------------------------------------------------------------------------------------- -# Pattern 2: wrap stream_iterate/astream_iterate in contextlib.closing/aclosing. +# Async: break inside `async with contextlib.aclosing(...)`. # -------------------------------------------------------------------------------------------- -def stop_iterate_with_closing_demo(): +async def async_aclosing_demo(): log = Timeline() - app = application(_respond_sync(log), log) - with contextlib.closing(app.stream_iterate(halt_after=["respond"])) as steps: - for _, container in steps: - for item in container: + app = application(_respond_async(log), log) + _, container = await app.astream_result(halt_after=["respond"]) + async with contextlib.aclosing(container): + async for item in container: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + log("caller: left the `async with` block") + log.print("Async: break inside contextlib.aclosing(container)") + + +async def async_iterate_aclosing_demo(): + log = Timeline() + app = application(_respond_async(log), log) + async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: + async for _, container in steps: + async for item in container: log(f"caller: got {item['token']!r}") if wants_to_stop(item["token"]): break break - log("caller: left the `with` block") - log.print("Pattern 2: contextlib.closing(app.stream_iterate(...))") + log("caller: left the `async with` block") + log.print("Async: break inside contextlib.aclosing(app.astream_iterate(...))") + + +async def async_break_without_aclose_demo(): + log = Timeline() + app = application(_respond_async(log), log) + _, container = await app.astream_result(halt_after=["respond"]) + async for item in container: + if wants_to_stop(item["token"]): + break + log("caller: after the loop -- nothing has been cleaned up yet") + await asyncio.sleep(0.01) # ...the event loop gets to it later, in the background + log("caller: some time later") + log.print("Async without aclosing: cleanup is left to the event loop (avoid)") # -------------------------------------------------------------------------------------------- -# Pattern 3: the client disconnects -> cancel the task that is consuming the stream. +# Web servers: the client disconnects -> the server cancels the task consuming the stream. # -------------------------------------------------------------------------------------------- @@ -194,7 +223,7 @@ async def stream_to_client(): task.cancel() with contextlib.suppress(asyncio.CancelledError): await task - log.print("Pattern 3: client disconnect, cancel the consuming task") + log.print("Web server: client disconnect cancels the consuming task") # -------------------------------------------------------------------------------------------- @@ -211,7 +240,6 @@ def resume_after_stop_demo(): for item in container: if wants_to_stop(item["token"]): break - container.close() log(f"persister: last step saved as {persister.load('', 'early-stop-demo')['status']!r}") resumed = ( @@ -229,7 +257,7 @@ def resume_after_stop_demo(): # -------------------------------------------------------------------------------------------- -# Cleanup errors: close() raises them to the caller. +# Cleanup errors: hooks always see them; close() raises them if you want them. # -------------------------------------------------------------------------------------------- @@ -238,14 +266,13 @@ def cleanup_error_demo(): app = application(_respond_sync(log, cleanup_error=True), log) _, container = app.stream_result(halt_after=["respond"]) for item in container: - log(f"caller: got {item['token']!r}") if wants_to_stop(item["token"]): - break + break # the error is reported to hooks and logged; the break itself doesn't raise try: container.close() except ConnectionError as e: log(f"caller: close() raised {e!r}") - log.print("Cleanup errors: close() raises them to the caller") + log.print("Cleanup errors: reported to hooks, and raised by close()") def visualize(): @@ -256,9 +283,12 @@ def visualize(): if __name__ == "__main__": - stop_with_close_demo() - asyncio.run(stop_with_aclose_demo()) - stop_iterate_with_closing_demo() + logging.basicConfig(level=logging.ERROR) # hide the cleanup-error warning in this demo + sync_break_demo() + sync_iterate_break_demo() + asyncio.run(async_aclosing_demo()) + asyncio.run(async_iterate_aclosing_demo()) + asyncio.run(async_break_without_aclose_demo()) asyncio.run(client_disconnect_demo()) resume_after_stop_demo() cleanup_error_demo() diff --git a/examples/streaming-early-stop/notebook.ipynb b/examples/streaming-early-stop/notebook.ipynb index 0223729da..f93d89c23 100644 --- a/examples/streaming-early-stop/notebook.ipynb +++ b/examples/streaming-early-stop/notebook.ipynb @@ -3,13 +3,13 @@ { "cell_type": "code", "execution_count": 1, - "id": "1c749fcb", + "id": "00ae3074", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:53.338155Z", - "iopub.status.busy": "2026-11-03T18:29:53.337954Z", - "iopub.status.idle": "2026-11-03T18:29:53.342511Z", - "shell.execute_reply": "2026-11-03T18:29:53.342120Z" + "iopub.execute_input": "2026-11-03T19:18:54.703130Z", + "iopub.status.busy": "2026-11-03T19:18:54.703026Z", + "iopub.status.idle": "2026-11-03T19:18:54.706150Z", + "shell.execute_reply": "2026-11-03T19:18:54.705826Z" } }, "outputs": [], @@ -35,13 +35,13 @@ { "cell_type": "code", "execution_count": null, - "id": "ab14c810", + "id": "bb1abd63", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:53.344766Z", - "iopub.status.busy": "2026-11-03T18:29:53.344598Z", - "iopub.status.idle": "2026-11-03T18:29:56.705276Z", - "shell.execute_reply": "2026-11-03T18:29:56.704928Z" + "iopub.execute_input": "2026-11-03T19:18:54.707835Z", + "iopub.status.busy": "2026-11-03T19:18:54.707725Z", + "iopub.status.idle": "2026-11-03T19:18:59.113436Z", + "shell.execute_reply": "2026-11-03T19:18:59.112668Z" } }, "outputs": [], @@ -51,51 +51,56 @@ }, { "cell_type": "markdown", - "id": "5f51adcc", + "id": "6e825efb", "metadata": {}, "source": [ "# Stopping a streaming action early\n", "\n", - "When you stop reading a stream before it ends, you want the action's cleanup to run *now*, the step to be reported as failed (so resuming re-runs it), and the state not to be updated. Each demo below prints a timeline of what happened, in order. The code is in [application.py](application.py), and the explanation is in [README.md](README.md)." + "When you stop reading a stream before it ends, Burr runs the action's cleanup right away, then reports the step as failed (so resuming re-runs it), and doesn't update the state. **Sync code just `break`s. Async code `break`s inside `async with contextlib.aclosing(...)`.**\n", + "\n", + "Each demo below prints a timeline of what happened, in order. The code is in [application.py](application.py), and the explanation is in [README.md](README.md)." ] }, { "cell_type": "code", "execution_count": 3, - "id": "b3869f03", + "id": "5203dcf7", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:56.706970Z", - "iopub.status.busy": "2026-11-03T18:29:56.706861Z", - "iopub.status.idle": "2026-11-03T18:29:57.083361Z", - "shell.execute_reply": "2026-11-03T18:29:57.083012Z" + "iopub.execute_input": "2026-11-03T19:18:59.116209Z", + "iopub.status.busy": "2026-11-03T19:18:59.115942Z", + "iopub.status.idle": "2026-11-03T19:18:59.713846Z", + "shell.execute_reply": "2026-11-03T19:18:59.713427Z" } }, "outputs": [], "source": [ "import asyncio\n", + "import logging\n", + "\n", + "import application as demo\n", "\n", - "import application as demo" + "logging.basicConfig(level=logging.ERROR)" ] }, { "cell_type": "markdown", - "id": "01f71955", + "id": "2bc86470", "metadata": {}, "source": [ - "## Pattern 1: `break`, then `container.close()` / `await container.aclose()`" + "## Sync: just `break`" ] }, { "cell_type": "code", "execution_count": 4, - "id": "aef6f0e1", + "id": "28becda8", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:57.085325Z", - "iopub.status.busy": "2026-11-03T18:29:57.085085Z", - "iopub.status.idle": "2026-11-03T18:29:57.087941Z", - "shell.execute_reply": "2026-11-03T18:29:57.087560Z" + "iopub.execute_input": "2026-11-03T19:18:59.716713Z", + "iopub.status.busy": "2026-11-03T19:18:59.716419Z", + "iopub.status.idle": "2026-11-03T19:18:59.719557Z", + "shell.execute_reply": "2026-11-03T19:18:59.719255Z" } }, "outputs": [ @@ -104,32 +109,32 @@ "output_type": "stream", "text": [ "\n", - "=== Pattern 1: break, then container.close()\n", + "=== Sync: break\n", " respond: open LLM connection\n", " caller: got 'Burr '\n", " caller: got 'streams '\n", " caller: got 'tokens '\n", - " post_run_step: respond failed with GeneratorExit()\n", " respond: close LLM connection\n", - " caller: closed the container\n", + " post_run_step: respond failed with GeneratorExit()\n", + " caller: after the loop\n", " caller: get() -> result=None, 'response' in state=False\n" ] } ], "source": [ - "demo.stop_with_close_demo()" + "demo.sync_break_demo()" ] }, { "cell_type": "code", "execution_count": 5, - "id": "30199433", + "id": "ef586cdc", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:57.089611Z", - "iopub.status.busy": "2026-11-03T18:29:57.089482Z", - "iopub.status.idle": "2026-11-03T18:29:57.091864Z", - "shell.execute_reply": "2026-11-03T18:29:57.091625Z" + "iopub.execute_input": "2026-11-03T19:18:59.722774Z", + "iopub.status.busy": "2026-11-03T19:18:59.722528Z", + "iopub.status.idle": "2026-11-03T19:18:59.725193Z", + "shell.execute_reply": "2026-11-03T19:18:59.724691Z" } }, "outputs": [ @@ -138,41 +143,72 @@ "output_type": "stream", "text": [ "\n", - "=== Pattern 1 (async): break, then await container.aclose()\n", + "=== Sync: break out of stream_iterate\n", " respond: open LLM connection\n", " caller: got 'Burr '\n", " caller: got 'streams '\n", " caller: got 'tokens '\n", " respond: close LLM connection\n", " post_run_step: respond failed with GeneratorExit()\n", - " caller: closed the container\n" + " caller: after the loop\n" ] } ], "source": [ - "await demo.stop_with_aclose_demo()" + "demo.sync_iterate_break_demo()" ] }, { "cell_type": "markdown", - "id": "9540ca87", + "id": "3ce77bb1", "metadata": {}, "source": [ - "## Pattern 2: `contextlib.closing(app.stream_iterate(...))`\n", - "\n", - "However you leave the block, the in-flight action is closed." + "## Async: `break` inside `contextlib.aclosing`" ] }, { "cell_type": "code", "execution_count": 6, - "id": "170e909a", + "id": "90a16b98", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T19:18:59.727245Z", + "iopub.status.busy": "2026-11-03T19:18:59.727111Z", + "iopub.status.idle": "2026-11-03T19:18:59.730054Z", + "shell.execute_reply": "2026-11-03T19:18:59.729717Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Async: break inside contextlib.aclosing(container)\n", + " respond: open LLM connection\n", + " caller: got 'Burr '\n", + " caller: got 'streams '\n", + " caller: got 'tokens '\n", + " respond: close LLM connection\n", + " post_run_step: respond failed with GeneratorExit()\n", + " caller: left the `async with` block\n" + ] + } + ], + "source": [ + "await demo.async_aclosing_demo()" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "9b4dd2b0", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:57.093140Z", - "iopub.status.busy": "2026-11-03T18:29:57.093026Z", - "iopub.status.idle": "2026-11-03T18:29:57.094760Z", - "shell.execute_reply": "2026-11-03T18:29:57.094556Z" + "iopub.execute_input": "2026-11-03T19:18:59.732184Z", + "iopub.status.busy": "2026-11-03T19:18:59.732053Z", + "iopub.status.idle": "2026-11-03T19:18:59.734460Z", + "shell.execute_reply": "2026-11-03T19:18:59.734169Z" } }, "outputs": [ @@ -181,41 +217,80 @@ "output_type": "stream", "text": [ "\n", - "=== Pattern 2: contextlib.closing(app.stream_iterate(...))\n", + "=== Async: break inside contextlib.aclosing(app.astream_iterate(...))\n", " respond: open LLM connection\n", " caller: got 'Burr '\n", " caller: got 'streams '\n", " caller: got 'tokens '\n", + " respond: close LLM connection\n", " post_run_step: respond failed with GeneratorExit()\n", + " caller: left the `async with` block\n" + ] + } + ], + "source": [ + "await demo.async_iterate_aclosing_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "38cbda4d", + "metadata": {}, + "source": [ + "Without `aclosing`, Python leaves the async cleanup to the event loop, to run at some later point:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "d14ced01", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T19:18:59.735901Z", + "iopub.status.busy": "2026-11-03T19:18:59.735787Z", + "iopub.status.idle": "2026-11-03T19:18:59.749733Z", + "shell.execute_reply": "2026-11-03T19:18:59.749049Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Async without aclosing: cleanup is left to the event loop (avoid)\n", + " respond: open LLM connection\n", + " caller: after the loop -- nothing has been cleaned up yet\n", " respond: close LLM connection\n", - " caller: left the `with` block\n" + " post_run_step: respond failed with GeneratorExit()\n", + " caller: some time later\n" ] } ], "source": [ - "demo.stop_iterate_with_closing_demo()" + "await demo.async_break_without_aclose_demo()" ] }, { "cell_type": "markdown", - "id": "ef962268", + "id": "e6b971fa", "metadata": {}, "source": [ - "## Pattern 3: a client disconnect cancels the consuming task\n", + "## Web servers: a client disconnect cancels the consuming task\n", "\n", "The hook receives `CancelledError`. See [server.py](server.py) for the FastAPI version." ] }, { "cell_type": "code", - "execution_count": 7, - "id": "4ce17cbd", + "execution_count": 9, + "id": "144fdfac", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:57.095983Z", - "iopub.status.busy": "2026-11-03T18:29:57.095913Z", - "iopub.status.idle": "2026-11-03T18:29:57.219020Z", - "shell.execute_reply": "2026-11-03T18:29:57.218734Z" + "iopub.execute_input": "2026-11-03T19:18:59.751621Z", + "iopub.status.busy": "2026-11-03T19:18:59.751469Z", + "iopub.status.idle": "2026-11-03T19:18:59.875523Z", + "shell.execute_reply": "2026-11-03T19:18:59.875184Z" } }, "outputs": [ @@ -224,7 +299,7 @@ "output_type": "stream", "text": [ "\n", - "=== Pattern 3: client disconnect, cancel the consuming task\n", + "=== Web server: client disconnect cancels the consuming task\n", " respond: open LLM connection\n", " server: sent 'Burr '\n", " server: sent 'streams '\n", @@ -240,7 +315,7 @@ }, { "cell_type": "markdown", - "id": "1869862c", + "id": "df4f3273", "metadata": {}, "source": [ "## Persistence: a stopped step is saved as `failed`, and resuming re-runs it" @@ -248,14 +323,14 @@ }, { "cell_type": "code", - "execution_count": 8, - "id": "4fa878d4", + "execution_count": 10, + "id": "5a6ec507", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:57.220438Z", - "iopub.status.busy": "2026-11-03T18:29:57.220348Z", - "iopub.status.idle": "2026-11-03T18:29:57.222888Z", - "shell.execute_reply": "2026-11-03T18:29:57.222643Z" + "iopub.execute_input": "2026-11-03T19:18:59.877311Z", + "iopub.status.busy": "2026-11-03T19:18:59.877179Z", + "iopub.status.idle": "2026-11-03T19:18:59.880004Z", + "shell.execute_reply": "2026-11-03T19:18:59.879722Z" } }, "outputs": [ @@ -266,8 +341,8 @@ "\n", "=== Persistence: a stopped step is 'failed' and is re-run on resume\n", " respond: open LLM connection\n", - " post_run_step: respond failed with GeneratorExit()\n", " respond: close LLM connection\n", + " post_run_step: respond failed with GeneratorExit()\n", " persister: last step saved as 'failed'\n", " resume: next action is 'respond' (it runs again)\n" ] @@ -279,22 +354,22 @@ }, { "cell_type": "markdown", - "id": "7469b4f3", + "id": "94a900f7", "metadata": {}, "source": [ - "## Cleanup errors are raised by `close()`" + "## Cleanup errors: reported to hooks, and raised by `close()`" ] }, { "cell_type": "code", - "execution_count": 9, - "id": "e307a750", + "execution_count": 11, + "id": "ef2ae65c", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T18:29:57.224179Z", - "iopub.status.busy": "2026-11-03T18:29:57.224097Z", - "iopub.status.idle": "2026-11-03T18:29:57.225865Z", - "shell.execute_reply": "2026-11-03T18:29:57.225642Z" + "iopub.execute_input": "2026-11-03T19:18:59.881678Z", + "iopub.status.busy": "2026-11-03T19:18:59.881554Z", + "iopub.status.idle": "2026-11-03T19:18:59.883520Z", + "shell.execute_reply": "2026-11-03T19:18:59.883266Z" } }, "outputs": [ @@ -303,13 +378,10 @@ "output_type": "stream", "text": [ "\n", - "=== Cleanup errors: close() raises them to the caller\n", + "=== Cleanup errors: reported to hooks, and raised by close()\n", " respond: open LLM connection\n", - " caller: got 'Burr '\n", - " caller: got 'streams '\n", - " caller: got 'tokens '\n", - " post_run_step: respond failed with GeneratorExit()\n", " respond: close LLM connection\n", + " post_run_step: respond failed with ConnectionError('LLM connection did not close cleanly')\n", " caller: close() raised ConnectionError('LLM connection did not close cleanly')\n" ] } diff --git a/examples/streaming-early-stop/server.py b/examples/streaming-early-stop/server.py index 60bfd7f30..f8eac66e4 100644 --- a/examples/streaming-early-stop/server.py +++ b/examples/streaming-early-stop/server.py @@ -54,7 +54,7 @@ async def respond(state: State) -> AsyncGenerator[Tuple[dict, Optional[State]], class LogStep(PostRunStepHook): - # A *sync* hook: see the README for why async post-step hooks need care on disconnect. + # A *sync* hook: on a client disconnect, async post_run_step hooks are cancelled before they run. def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): logger.info("post_run_step: %s exception=%r", action.name, exception) diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 842ceb681..6c337cae2 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -33,14 +33,12 @@ Action, AsyncGenerator, AsyncStreamingAction, - AsyncStreamingResultContainer, Condition, Reducer, Result, SingleStepAction, SingleStepStreamingAction, StreamingAction, - StreamingResultContainer, action, default, expr, @@ -2775,436 +2773,467 @@ async def test_astream_iterate(exhaust_intermediate_generators: bool): assert len(stream_event_tracker.post_stream_item_calls) == 20 -def _closing_probe_app(events: list, *hooks) -> Application: - """One streaming action whose nested producer has an awaiting ``finally`` (like an HTTP close).""" +# --- Stopping a streaming action early --------------------------------------------------------- +# Every way of stopping a stream early (break, close()/aclose(), closing stream_iterate, task +# cancellation, dropping the container) runs the action's cleanup *first*, then reports the step to +# hooks/persisters exactly once as failed, without updating state -- independent of timing. - async def producer(): - try: - yield "a" - yield "b" - finally: - await asyncio.sleep(0) - events.append("producer_closed") - @streaming_action(reads=[], writes=["x"]) - async def act(state: State): - async with contextlib.aclosing(producer()) as items: - async for item in items: - yield {"item": item}, None - events.append("advanced") - yield {"item": "done"}, state.update(x=1) +class _EarlyStopCleanupError(Exception): + pass - return ( - ApplicationBuilder() - .with_actions(act) - .with_transitions() - .with_entrypoint("act") - .with_state(State()) - .with_hooks(*hooks) - .build() - ) +class _EarlyStopLog: + """Ordered record of what the action and the hooks did. Hooks store exception *types* only: + holding the exception would keep its traceback, and so the stream's frames, alive.""" -@pytest.mark.parametrize("consume_first_item", [True, False]) -async def test_astream_iterate_aclose_unwinds_in_flight_action(consume_first_item: bool): - """Closing astream_iterate mid-action runs the action's cleanup now (not at GC), - does not drain or advance it, fires post_end_stream exactly once, and reports the - step as failed.""" - events = [] - stream_event_tracker = StreamEventCaptureTracker() - call_tracker = CallCaptureTracker() - app = _closing_probe_app(events, stream_event_tracker, call_tracker) - steps = app.astream_iterate(halt_after=["act"]) - async for _, container in steps: - if consume_first_item: - async for _ in container: - break - break - await steps.aclose() - assert events == (["producer_closed"] if consume_first_item else []) - assert len(stream_event_tracker.post_end_stream_calls) == 1 - assert len(call_tracker.post_called) == 1 - assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + def __init__(self): + self.events: list = [] + self.post_end_stream_calls = 0 + def hooks(self, is_async: bool) -> list: + log = self -async def test_astream_iterate_cancelled_mid_stream_unwinds_and_reports_failure(): - """Cancelling a consumer that wraps astream_iterate in ``aclosing`` (e.g. on a client - disconnect) unwinds the action immediately and reports the step as failed.""" - events = [] - call_tracker = CallCaptureTracker() - app = _closing_probe_app(events, call_tracker) - first_item_seen = asyncio.Event() + def record(exception): + log.events.append(("reported", type(exception) if exception is not None else None)) - async def consume(): - async with contextlib.aclosing(app.astream_iterate(halt_after=["act"])) as steps: - async for _, container in steps: - async for _ in container: - first_item_seen.set() - await asyncio.sleep(10) + class SyncHook(PostRunStepHook, PostEndStreamHook): + def post_run_step(self, *, exception, **future_kwargs): + record(exception) - task = asyncio.create_task(consume()) - await first_item_seen.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert events == ["producer_closed"] - assert len(call_tracker.post_called) == 1 - assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + def post_end_stream(self, **future_kwargs): + log.post_end_stream_calls += 1 + class AsyncHook(PostRunStepHookAsync, PostEndStreamHook): + async def post_run_step(self, *, exception, **future_kwargs): + await asyncio.sleep(0) + record(exception) -async def test_astream_iterate_aclose_mid_stream_persists_failed_and_resumes_at_same_action(): - @streaming_action(reads=[], writes=["x"]) - async def act(state: State): - yield {"item": "a"}, None - yield {"item": "done"}, state.update(x=1) + def post_end_stream(self, **future_kwargs): + log.post_end_stream_calls += 1 - @action(reads=["x"], writes=["y"]) - def after(state: State) -> State: - return state.update(y=state["x"] + 1) + return [AsyncHook() if is_async else SyncHook()] - persister = SQLLitePersister(db_path=":memory:", table_name="test_table") - persister.initialize() - def build(resume: bool) -> Application: - builder = ( - ApplicationBuilder() - .with_actions(act=act, after=after) - .with_transitions(("act", "after")) - ) - if resume: - builder = builder.initialize_from( - persister, resume_at_next_action=True, default_state={}, default_entrypoint="act" - ) - else: - builder = builder.with_entrypoint("act").with_state(State()) - return builder.with_identifiers(app_id="app").with_state_persister(persister).build() - - steps = build(resume=False).astream_iterate(halt_after=["after"]) - async for _, container in steps: - async for _ in container: - break - break - await steps.aclose() - assert persister.load("", "app")["status"] == "failed" - assert build(resume=True).get_next_action().name == "act" +def _early_stop_action( + log: _EarlyStopLog, is_async: bool, cleanup_raises: bool = False, delay: float = 0.0 +): + """A streaming action holding a resource that it releases in ``finally``.""" + def cleanup(): + log.events.append("cleanup") + if cleanup_raises: + raise _EarlyStopCleanupError("cleanup failed") -async def test_astream_iterate_aclose_unwinds_multi_step_streaming_action(): - events = [] + if is_async: - class MultiStepAct(AsyncStreamingAction): - @property - def reads(self) -> list[str]: - return [] + @streaming_action(reads=[], writes=["x"]) + async def act(state: State): + try: + for item in ["a", "b", "c"]: + await asyncio.sleep(delay) + yield {"item": item}, None + await asyncio.sleep(delay) + yield {"item": "done"}, state.update(x=1) + finally: + await asyncio.sleep(0) # e.g. closing an HTTP stream + cleanup() - @property - def writes(self) -> list[str]: - return ["x"] + else: - async def stream_run(self, state: State, **run_kwargs): + @streaming_action(reads=[], writes=["x"]) + def act(state: State): try: - for i in range(3): - yield {"i": i} + for item in ["a", "b", "c"]: + yield {"item": item}, None + yield {"item": "done"}, state.update(x=1) finally: - await asyncio.sleep(0) - events.append("stream_run_closed") + cleanup() - def update(self, result: dict, state: State) -> State: - return state.update(x=result["i"]) + return act - app = ( - ApplicationBuilder() - .with_actions(act=MultiStepAct()) - .with_transitions() - .with_entrypoint("act") - .with_state(State()) - .build() - ) - steps = app.astream_iterate(halt_after=["act"]) - async for _, container in steps: - async for _ in container: - break - break - await steps.aclose() - assert events == ["stream_run_closed"] - assert "x" not in app.state +def _early_stop_multi_step_action(log: _EarlyStopLog, is_async: bool) -> Action: + if is_async: -async def test_streaming_result_container_aclose_is_idempotent_and_get_returns_initial_state(): - events = [] - app = _closing_probe_app(events) - _, container = await app.astream_result(halt_after=["act"]) - async for _ in container: - break - await container.aclose() - await container.aclose() - result, state = await container.get() - assert events == ["producer_closed"] - assert result is None - assert "x" not in state + class MultiStepAct(AsyncStreamingAction): + reads = [] + writes = ["x"] + async def stream_run(self, state: State, **run_kwargs): + try: + for i in range(3): + yield {"i": i} + finally: + await asyncio.sleep(0) + log.events.append("cleanup") -async def test_streaming_result_container_aclose_after_completion_reports_success(): - events = [] - call_tracker = CallCaptureTracker() - app = _closing_probe_app(events, call_tracker) - _, container = await app.astream_result(halt_after=["act"]) - result, state = await container.get() - await container.aclose() - assert result == {"item": "done"} - assert state["x"] == 1 - assert len(call_tracker.post_called) == 1 - assert call_tracker.post_called[0][1]["exception"] is None - - -async def test_streaming_result_container_abandoned_without_aclose_reports_without_cyclic_gc(): - """Breaking out and dropping the container (no aclose) must still report the step via - refcounting alone -- the container must not form a reference cycle with its iterator, - which would defer the post-step hooks until cyclic GC runs.""" - call_tracker = CallCaptureTracker() - app = _closing_probe_app([], call_tracker) - gc.disable() - try: - _, container = await app.astream_result(halt_after=["act"]) - async for _ in container: - break - del container - for _ in range(5): # let the event loop run the scheduled asyncgen finalizer - await asyncio.sleep(0) - assert len(call_tracker.post_called) == 1 - assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) - finally: - gc.enable() + def update(self, result: dict, state: State) -> State: + return state.update(x=result["i"]) + else: -async def test_streaming_result_container_aclose_propagates_cleanup_errors(): - class CleanupError(Exception): - pass + class MultiStepAct(StreamingAction): + reads = [] + writes = ["x"] - async def failing_cleanup(): - try: - yield {"item": "a"}, None - finally: - raise CleanupError("cleanup failed") + def stream_run(self, state: State, **run_kwargs): + try: + for i in range(3): + yield {"i": i} + finally: + log.events.append("cleanup") - seen = [] + def update(self, result: dict, state: State) -> State: + return state.update(x=result["i"]) - async def callback(result, state, exc): - seen.append(exc) + return MultiStepAct() - container = AsyncStreamingResultContainer( - failing_cleanup(), State(), lambda r, s: (r, s), callback - ) - async for _ in container: - break - with pytest.raises(CleanupError): - await container.aclose() - assert len(seen) == 1 - assert isinstance(seen[0], CleanupError) +@action(reads=["x"], writes=["y"]) +def _early_stop_after(state: State) -> State: + return state.update(y=state["x"] + 1) -class _ExceptionTypeCapture(PostRunStepHook): - """Records only the exception *type* -- holding the exception itself would keep its traceback, - and so the stream's frames, alive, defeating refcount-based cleanup checks.""" - def __init__(self): - self.exception_types = [] +def _early_stop_app( + act: Action, + log: _EarlyStopLog, + is_async: bool, + persister: Optional[SQLLitePersister] = None, + resume: bool = False, + hooks: tuple = (), +) -> Application: + builder = ( + ApplicationBuilder() + .with_actions(act=act, after=_early_stop_after) + .with_transitions(("act", "after")) + .with_hooks(*log.hooks(is_async), *hooks) + .with_identifiers(app_id="early-stop") + ) + if resume: + builder = builder.initialize_from( + persister, resume_at_next_action=True, default_state={}, default_entrypoint="act" + ) + else: + builder = builder.with_entrypoint("act").with_state(State()) + if persister is not None: + builder = builder.with_state_persister(persister) + return builder.build() - def post_run_step(self, *, exception: Optional[BaseException], **future_kwargs: Any): - self.exception_types.append(type(exception) if exception is not None else None) +def _early_stop_persister() -> SQLLitePersister: + persister = SQLLitePersister(db_path=":memory:", table_name="early_stop") + persister.initialize() + return persister -def _sync_closing_probe_app(events: list, *hooks) -> Application: - """Sync twin of ``_closing_probe_app``: one streaming action wrapping a nested producer.""" - def producer(): - try: - yield "a" - yield "b" - finally: - events.append("producer_closed") +STOPPED = [("reported", GeneratorExit)] - @streaming_action(reads=[], writes=["x"]) - def act(state: State): - with contextlib.closing(producer()) as items: - for item in items: - yield {"item": item}, None - events.append("advanced") - yield {"item": "done"}, state.update(x=1) - return ( - ApplicationBuilder() - .with_actions(act) - .with_transitions() - .with_entrypoint("act") - .with_state(State()) - .with_hooks(*hooks) - .build() - ) +@pytest.fixture +def no_cyclic_gc(): + """Early-close behaviour must not depend on when the cyclic GC happens to run.""" + gc.disable() + yield + gc.enable() -@pytest.mark.parametrize("consume_first_item", [True, False]) -def test_stream_iterate_close_unwinds_in_flight_action(consume_first_item: bool): - """Closing stream_iterate mid-action runs the action's cleanup now (not at GC), does not - drain or advance it, fires post_end_stream exactly once, and reports the step as failed.""" - events = [] - stream_event_tracker = StreamEventCaptureTracker() - call_tracker = CallCaptureTracker() - app = _sync_closing_probe_app(events, stream_event_tracker, call_tracker) - steps = app.stream_iterate(halt_after=["act"]) - for _, container in steps: - if consume_first_item: - for _ in container: - break +def test_stream_result_break_runs_cleanup_then_reports_failed_step(no_cyclic_gc): + """Plain ``for ... break`` is all a sync caller needs: no explicit close().""" + log = _EarlyStopLog() + app = _early_stop_app(_early_stop_action(log, is_async=False), log, is_async=False) + _, container = app.stream_result(halt_after=["act"]) + for _ in container: break - steps.close() - assert events == (["producer_closed"] if consume_first_item else []) - assert len(stream_event_tracker.post_end_stream_calls) == 1 - assert len(call_tracker.post_called) == 1 - assert isinstance(call_tracker.post_called[0][1]["exception"], GeneratorExit) + assert log.events == ["cleanup", *STOPPED] + assert log.post_end_stream_calls == 1 + result, state = container.get() + assert result is None + assert "x" not in state and "x" not in app.state -def test_stream_iterate_close_mid_stream_persists_failed_and_resumes_at_same_action(): - @streaming_action(reads=[], writes=["x"]) - def act(state: State): - yield {"item": "a"}, None - yield {"item": "done"}, state.update(x=1) +def test_stream_result_dropped_mid_stream_runs_cleanup_then_reports(no_cyclic_gc): + log = _EarlyStopLog() + app = _early_stop_app(_early_stop_action(log, is_async=False), log, is_async=False) + _, container = app.stream_result(halt_after=["act"]) + items = iter(container) + next(items) + del items, container + assert log.events == ["cleanup", *STOPPED] - @action(reads=["x"], writes=["y"]) - def after(state: State) -> State: - return state.update(y=state["x"] + 1) - persister = SQLLitePersister(db_path=":memory:", table_name="test_table") - persister.initialize() +@pytest.mark.parametrize("awaits_before_aclose", [0, 1, 2, 3, 5]) +async def test_astream_result_break_then_aclose_runs_cleanup_then_reports( + awaits_before_aclose: int, +): + """``async for ... break`` hands the iterator to the event loop's async-generator finalizer; + an aclose() racing it -- however many awaits later -- must give the same outcome.""" + log = _EarlyStopLog() + app = _early_stop_app(_early_stop_action(log, is_async=True), log, is_async=True) + _, container = await app.astream_result(halt_after=["act"]) + async for _ in container: + break + for _ in range(awaits_before_aclose): + await asyncio.sleep(0) + await container.aclose() + assert log.events == ["cleanup", *STOPPED] + assert log.post_end_stream_calls == 1 + result, state = await container.get() + assert result is None + assert "x" not in state and "x" not in app.state + + +async def test_astream_result_dropped_mid_stream_runs_cleanup_then_reports(no_cyclic_gc): + log = _EarlyStopLog() + app = _early_stop_app(_early_stop_action(log, is_async=True), log, is_async=True) + _, container = await app.astream_result(halt_after=["act"]) + async for _ in container: + break + del container + for _ in range(10): # let the event loop run the async-generator finalizer + await asyncio.sleep(0) + assert log.events == ["cleanup", *STOPPED] - def build(resume: bool) -> Application: - builder = ( - ApplicationBuilder() - .with_actions(act=act, after=after) - .with_transitions(("act", "after")) - ) - if resume: - builder = builder.initialize_from( - persister, resume_at_next_action=True, default_state={}, default_entrypoint="act" - ) - else: - builder = builder.with_entrypoint("act").with_state(State()) - return builder.with_identifiers(app_id="app").with_state_persister(persister).build() - with contextlib.closing(build(resume=False).stream_iterate(halt_after=["after"])) as steps: +@pytest.mark.parametrize("multi_step", [False, True], ids=["single_step", "multi_step"]) +@pytest.mark.parametrize("consume_first_item", [True, False]) +def test_stream_iterate_closing_unwinds_in_flight_action( + consume_first_item: bool, multi_step: bool +): + log = _EarlyStopLog() + act = ( + _early_stop_multi_step_action(log, is_async=False) + if multi_step + else _early_stop_action(log, is_async=False) + ) + app = _early_stop_app(act, log, is_async=False) + with contextlib.closing(app.stream_iterate(halt_after=["after"])) as steps: for _, container in steps: - for _ in container: - break + if consume_first_item: + for _ in container: + break break - assert persister.load("", "app")["status"] == "failed" - assert build(resume=True).get_next_action().name == "act" + # An action that never started has nothing to clean up. + assert log.events == (["cleanup"] if consume_first_item else []) + STOPPED + assert log.post_end_stream_calls == 1 + assert "x" not in app.state -def test_stream_iterate_close_unwinds_multi_step_streaming_action(): - events = [] +@pytest.mark.parametrize("multi_step", [False, True], ids=["single_step", "multi_step"]) +@pytest.mark.parametrize("consume_first_item", [True, False]) +async def test_astream_iterate_aclosing_unwinds_in_flight_action( + consume_first_item: bool, multi_step: bool +): + log = _EarlyStopLog() + act = ( + _early_stop_multi_step_action(log, is_async=True) + if multi_step + else _early_stop_action(log, is_async=True) + ) + app = _early_stop_app(act, log, is_async=True) + async with contextlib.aclosing(app.astream_iterate(halt_after=["after"])) as steps: + async for _, container in steps: + if consume_first_item: + async for _ in container: + break + break + assert log.events == (["cleanup"] if consume_first_item else []) + STOPPED + assert log.post_end_stream_calls == 1 + assert "x" not in app.state - class MultiStepAct(StreamingAction): - @property - def reads(self) -> list[str]: - return [] - @property - def writes(self) -> list[str]: - return ["x"] +async def test_astream_iterate_cancelled_consumer_unwinds_and_reports_cancelled_error(): + """A web server cancels the response task when the client disconnects.""" + log = _EarlyStopLog() + persister = _early_stop_persister() + app = _early_stop_app( + _early_stop_action(log, is_async=True, delay=0.05), log, is_async=True, persister=persister + ) + first_item_seen = asyncio.Event() - def stream_run(self, state: State, **run_kwargs): - try: - for i in range(3): - yield {"i": i} - finally: - events.append("stream_run_closed") + async def respond_to_client(): + async with contextlib.aclosing(app.astream_iterate(halt_after=["after"])) as steps: + async for _, container in steps: + async for _ in container: + first_item_seen.set() - def update(self, result: dict, state: State) -> State: - return state.update(x=result["i"]) + task = asyncio.create_task(respond_to_client()) + await first_item_seen.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert log.events == ["cleanup", ("reported", asyncio.CancelledError)] + assert persister.load("", "early-stop")["status"] == "failed" - app = ( - ApplicationBuilder() - .with_actions(act=MultiStepAct()) - .with_transitions() - .with_entrypoint("act") - .with_state(State()) - .build() + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_stopped_step_persists_failed_and_resume_reruns_it(is_async: bool): + log = _EarlyStopLog() + persister = _early_stop_persister() + app = _early_stop_app( + _early_stop_action(log, is_async), log, is_async=is_async, persister=persister ) - steps = app.stream_iterate(halt_after=["act"]) - for _, container in steps: - for _ in container: - break + if is_async: + async with contextlib.aclosing(app.astream_iterate(halt_after=["after"])) as steps: + async for _, container in steps: + async for _ in container: + break + break + else: + with contextlib.closing(app.stream_iterate(halt_after=["after"])) as steps: + for _, container in steps: + for _ in container: + break + break + assert persister.load("", "early-stop")["status"] == "failed" + resumed = _early_stop_app( + _early_stop_action(log, is_async), log, is_async, persister=persister, resume=True + ) + assert resumed.get_next_action().name == "act" + + +def test_stream_result_close_is_idempotent_and_noop_after_completion(): + log = _EarlyStopLog() + app = _early_stop_app(_early_stop_action(log, is_async=False), log, is_async=False) + _, stopped = app.stream_result(halt_after=["act"]) + next(iter(stopped)) + stopped.close() + stopped.close() + assert log.events == ["cleanup", *STOPPED] + + log.events.clear() + _, completed = app.stream_result(halt_after=["act"]) + result, state = completed.get() + completed.close() + assert result == {"item": "done"} and state["x"] == 1 + assert log.events == ["cleanup", ("reported", None)] + + +async def test_astream_result_aclose_is_idempotent_and_noop_after_completion(): + log = _EarlyStopLog() + app = _early_stop_app(_early_stop_action(log, is_async=True), log, is_async=True) + _, stopped = await app.astream_result(halt_after=["act"]) + await stopped.__aiter__().__anext__() + await stopped.aclose() + await stopped.aclose() + assert log.events == ["cleanup", *STOPPED] + + log.events.clear() + _, completed = await app.astream_result(halt_after=["act"]) + result, state = await completed.get() + await completed.aclose() + assert result == {"item": "done"} and state["x"] == 1 + assert log.events == ["cleanup", ("reported", None)] + + +def test_stream_result_cleanup_error_is_reported_and_raised_by_close(no_cyclic_gc, caplog): + """After ``for ... break`` there is nobody to raise a cleanup error to, so it's reported to + hooks and logged at the break -- and close() raises it, for callers that want it.""" + log = _EarlyStopLog() + app = _early_stop_app( + _early_stop_action(log, is_async=False, cleanup_raises=True), log, is_async=False + ) + _, container = app.stream_result(halt_after=["act"]) + for _ in container: break - steps.close() - assert events == ["stream_run_closed"] - assert "x" not in app.state + assert log.events == ["cleanup", ("reported", _EarlyStopCleanupError)] + assert "Cleanup of a streaming action raised" in caplog.text + with pytest.raises(_EarlyStopCleanupError): + container.close() + with pytest.raises(_EarlyStopCleanupError): + container.close() -def test_streaming_result_container_close_is_idempotent_and_get_returns_initial_state(): - events = [] - app = _sync_closing_probe_app(events) - _, container = app.stream_result(halt_after=["act"]) - for _ in container: +@pytest.mark.parametrize("awaits_before_aclose", [0, 1, 2, 3, 5]) +async def test_astream_result_cleanup_error_is_reported_and_raised_by_aclose( + awaits_before_aclose: int, +): + log = _EarlyStopLog() + app = _early_stop_app( + _early_stop_action(log, is_async=True, cleanup_raises=True), log, is_async=True + ) + _, container = await app.astream_result(halt_after=["act"]) + async for _ in container: break - container.close() - container.close() - result, state = container.get() - assert events == ["producer_closed"] - assert result is None - assert "x" not in state + for _ in range(awaits_before_aclose): + await asyncio.sleep(0) + with pytest.raises(_EarlyStopCleanupError): + await container.aclose() + assert log.events == ["cleanup", ("reported", _EarlyStopCleanupError)] -def test_streaming_result_container_close_after_completion_reports_success(): - call_tracker = CallCaptureTracker() - app = _sync_closing_probe_app([], call_tracker) - _, container = app.stream_result(halt_after=["act"]) - result, state = container.get() - container.close() - assert result == {"item": "done"} - assert state["x"] == 1 - assert len(call_tracker.post_called) == 1 - assert call_tracker.post_called[0][1]["exception"] is None - - -def test_streaming_result_container_abandoned_without_close_reports_without_cyclic_gc(): - """Breaking out and dropping the container (no close) must unwind the action and report the - step via refcounting alone -- the container must not form a cycle with its iterator.""" - events = [] - exception_capture = _ExceptionTypeCapture() - app = _sync_closing_probe_app(events, exception_capture) - gc.disable() - try: - _, container = app.stream_result(halt_after=["act"]) - for _ in container: - break - del container - assert events == ["producer_closed"] - assert exception_capture.exception_types == [GeneratorExit] - finally: - gc.enable() +class _BrokenStreamItemHook(PostStreamItemHook): + def post_stream_item(self, **future_kwargs): + raise KeyError("bug in a hook") -def test_streaming_result_container_close_propagates_cleanup_errors(): - class CleanupError(Exception): - pass +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_cleanup_error_does_not_mask_the_error_that_stopped_the_stream( + is_async: bool, caplog +): + log = _EarlyStopLog() + app = _early_stop_app( + _early_stop_action(log, is_async, cleanup_raises=True), + log, + is_async, + hooks=(_BrokenStreamItemHook(),), + ) + with pytest.raises(KeyError, match="bug in a hook"): + if is_async: + _, container = await app.astream_result(halt_after=["act"]) + async for _ in container: + pass + else: + _, container = app.stream_result(halt_after=["act"]) + for _ in container: + pass + assert log.events == ["cleanup", ("reported", KeyError)] + assert "Cleanup of streaming action 'act' raised while handling" in caplog.text - def failing_cleanup(): - try: - yield {"item": "a"}, None - finally: - raise CleanupError("cleanup failed") - seen = [] - container = StreamingResultContainer( - failing_cleanup(), State(), lambda r, s: (r, s), lambda r, s, exc: seen.append(exc) +def test_stream_result_close_from_a_hook_raises_a_clear_error(): + """Closing while the stream is producing an item is not supported: stop from the consumer.""" + log = _EarlyStopLog() + holder = {} + + class CloseFromHook(PostStreamItemHook): + def post_stream_item(self, **future_kwargs): + holder["container"].close() + + app = _early_stop_app( + _early_stop_action(log, is_async=False), log, is_async=False, hooks=(CloseFromHook(),) ) - # Hold the iterator: breaking out of a ``for`` would drop it and report the close immediately. - items = iter(container) - next(items) - with pytest.raises(CleanupError): - container.close() - assert len(seen) == 1 - assert isinstance(seen[0], CleanupError) + _, holder["container"] = app.stream_result(halt_after=["act"]) + with pytest.raises(RuntimeError, match="Cannot close a streaming result while it is producing"): + for _ in holder["container"]: + pass + assert log.events == ["cleanup", ("reported", RuntimeError)] + + +@pytest.mark.parametrize("during", ["middle_item", "final_item"]) +async def test_astream_result_aclose_from_another_task_raises_and_stream_continues(during: str): + """Closing from another task while the consumer is awaiting the next item raises a clear + error and leaves the stream alone: the step's outcome is decided by its consumer.""" + log = _EarlyStopLog() + persister = _early_stop_persister() + app = _early_stop_app( + _early_stop_action(log, is_async=True, delay=0.05), log, is_async=True, persister=persister + ) + _, container = await app.astream_result(halt_after=["act"]) + + async def consume(): + async for _ in container: + pass + + async def close_from_elsewhere(): + await asyncio.sleep(0.07 if during == "middle_item" else 0.17) + with pytest.raises(RuntimeError, match="cancel the task that consumes it"): + await container.aclose() + + await asyncio.gather(consume(), close_from_elsewhere()) + assert log.events == ["cleanup", ("reported", None)] + assert persister.load("", "early-stop")["status"] == "completed" + assert app.state["x"] == 1 async def test_astream_result_halt_after_run_through_streaming(): diff --git a/tests/core/test_streaming_early_stop_example.py b/tests/core/test_streaming_early_stop_example.py index d781af506..ebe4db94c 100644 --- a/tests/core/test_streaming_early_stop_example.py +++ b/tests/core/test_streaming_early_stop_example.py @@ -24,6 +24,9 @@ EXAMPLE = pathlib.Path(__file__).parents[2] / "examples" / "streaming-early-stop" / "application.py" +CLOSED = "respond: close LLM connection" +STOPPED = "post_run_step: respond failed with GeneratorExit()" + @pytest.fixture(scope="module") def demo(): @@ -37,42 +40,45 @@ def _timeline(capsys) -> list[str]: return [line.strip() for line in capsys.readouterr().out.splitlines() if line.startswith(" ")] -def test_stop_with_close(demo, capsys): - demo.stop_with_close_demo() +def _cleanup_then_report_then(events: list[str], caller_line: str): + """The action is cleaned up, then the step is reported, all before the caller moves on.""" + i = events.index(CLOSED) + assert events[i : i + 3] == [CLOSED, STOPPED, caller_line] + + +def test_sync_break(demo, capsys): + demo.sync_break_demo() events = _timeline(capsys) - assert "respond: close LLM connection" in events - assert "post_run_step: respond failed with GeneratorExit()" in events - assert events.index("respond: close LLM connection") < events.index( - "caller: closed the container" - ) + _cleanup_then_report_then(events, "caller: after the loop") assert events[-1] == "caller: get() -> result=None, 'response' in state=False" -async def test_stop_with_aclose(demo, capsys): - await demo.stop_with_aclose_demo() - events = _timeline(capsys) - assert events.index("respond: close LLM connection") < events.index( - "caller: closed the container" - ) - assert "post_run_step: respond failed with GeneratorExit()" in events +def test_sync_iterate_break(demo, capsys): + demo.sync_iterate_break_demo() + _cleanup_then_report_then(_timeline(capsys), "caller: after the loop") + + +async def test_async_aclosing(demo, capsys): + await demo.async_aclosing_demo() + _cleanup_then_report_then(_timeline(capsys), "caller: left the `async with` block") -def test_stop_iterate_with_closing(demo, capsys): - demo.stop_iterate_with_closing_demo() +async def test_async_iterate_aclosing(demo, capsys): + await demo.async_iterate_aclosing_demo() + _cleanup_then_report_then(_timeline(capsys), "caller: left the `async with` block") + + +async def test_async_break_without_aclose_defers_cleanup(demo, capsys): + await demo.async_break_without_aclose_demo() events = _timeline(capsys) - assert events.index("respond: close LLM connection") < events.index( - "caller: left the `with` block" + assert events.index("caller: after the loop -- nothing has been cleaned up yet") < ( + events.index(CLOSED) ) - assert "post_run_step: respond failed with GeneratorExit()" in events async def test_client_disconnect(demo, capsys): await demo.client_disconnect_demo() - events = _timeline(capsys) - assert events[-2:] == [ - "respond: close LLM connection", - "post_run_step: respond failed with CancelledError()", - ] + assert _timeline(capsys)[-2:] == [CLOSED, "post_run_step: respond failed with CancelledError()"] def test_resume_after_stop(demo, capsys): @@ -84,7 +90,8 @@ def test_resume_after_stop(demo, capsys): def test_cleanup_error(demo, capsys): demo.cleanup_error_demo() - events = _timeline(capsys) - assert events[-1] == ( - "caller: close() raised ConnectionError('LLM connection did not close cleanly')" - ) + assert _timeline(capsys)[-3:] == [ + CLOSED, + "post_run_step: respond failed with ConnectionError('LLM connection did not close cleanly')", + "caller: close() raised ConnectionError('LLM connection did not close cleanly')", + ] From 370c54bf1278fe464a79f728f4f02312079714aa Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Tue, 3 Nov 2026 15:39:22 -0800 Subject: [PATCH 7/8] docs(examples): use async hooks in async demos; document disconnect behaviour Async applications shouldn't mix in sync hooks (they block the event loop), so the async demos and server.py now use async hooks and abuild(), and the README no longer recommends sync hooks for recording client disconnects. Starlette/FastAPI cancel a disconnected response through an anyio cancel scope, which also cancels Burr's await on async post_run_step hooks and persisters, so the stopped step isn't recorded. This is unchanged from main. The action is still cleaned up, and resume is unaffected: the last saved step is the one before, so the interrupted action runs again. - add cancel_task_demo (task.cancel(): hooks see CancelledError) - client_disconnect_demo now reproduces the cancel-scope case in-process, with an async persister and a resume check - add a core test pinning the cancel-scope behaviour (cleanup runs, nothing saved for the step, resume lands on it) - add anyio to the tests extra (used directly by the tests and example) --- docs/concepts/streaming-actions.rst | 8 +- examples/streaming-early-stop/README.md | 59 +++- examples/streaming-early-stop/application.py | 110 +++++-- examples/streaming-early-stop/notebook.ipynb | 289 ++++++++++++++---- examples/streaming-early-stop/server.py | 21 +- pyproject.toml | 1 + tests/core/test_application.py | 58 ++++ .../core/test_streaming_early_stop_example.py | 20 +- 8 files changed, 468 insertions(+), 98 deletions(-) diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index a11fbb7f0..0113d4c7a 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -272,7 +272,13 @@ or from another task) raises a ``RuntimeError``. A step that is stopped before it completes does not update state, and is reported to ``post_run_step`` hooks with the exception that stopped it -- ``GeneratorExit`` when it is closed, ``CancelledError`` when the consuming -task is cancelled. Persisters record it as ``failed``, so resuming re-runs that action. If the action's cleanup +task is cancelled. Persisters record it as ``failed``, so resuming re-runs that action. + +One exception: on a client disconnect, Starlette/FastAPI cancel the response through an anyio cancel scope, which +also cancels Burr's await on *async* ``post_run_step`` hooks and persisters. The action is still cleaned up, but +the stopped step isn't recorded. Resuming is unaffected: the last saved step is the one before, so the +interrupted action runs again. Don't add sync hooks to an async application to work around this; they block the +event loop. If the action's cleanup itself raises, hooks receive that error, and ``close()`` / ``aclose()`` raise it. A cleanup error never replaces an error that was already stopping the stream; it is logged instead. diff --git a/examples/streaming-early-stop/README.md b/examples/streaming-early-stop/README.md index 631e1a43d..914323601 100644 --- a/examples/streaming-early-stop/README.md +++ b/examples/streaming-early-stop/README.md @@ -26,7 +26,8 @@ the client disconnects), Burr: 1. **runs the action's cleanup right away**, not whenever the garbage collector gets to it; 2. **then reports the step as failed** to hooks, trackers and persisters, so resuming the application - re-runs it; + re-runs it (on a web-server [client disconnect](#web-servers-client-disconnects), async hooks and + persisters don't get to run, but resuming still re-runs the step); 3. **doesn't update the state**, so a half-finished response is never committed. All you have to do is stop iterating the way you would with any Python generator: @@ -37,6 +38,8 @@ All you have to do is stop iterating the way you would with any Python generator | async | `break` inside `async with contextlib.aclosing(...)` | | web server | the async pattern; the server cancels your task when the client disconnects | +Don't mix sync and async: async applications should use async hooks and persisters, as the async demos do. + Every demo in [application.py](application.py) prints a timeline of what happened, so you can see when the resource is released and what Burr reports. Run it with: @@ -101,9 +104,26 @@ post_run_step: respond failed with GeneratorExit() `await container.aclose()` after the loop does the same as `aclosing`, if you prefer it. +## Cancelling the consuming task + +**To stop a stream from somewhere else** (another task, a "stop" endpoint), cancel the task that is +consuming it (`cancel_task_demo`). The action is cleaned up and hooks receive `CancelledError`: + +``` +respond: open LLM connection +caller: got 'Burr ' +caller: got 'streams ' +caller: task.cancel() +respond: close LLM connection +post_run_step: respond failed with CancelledError() +``` + +Calling `close()` / `aclose()` while the stream is producing an item (from another task, or from inside +a hook) raises a `RuntimeError` that tells you to cancel the task instead. + ## Web servers: client disconnects -When a client disconnects, ASGI servers such as Starlette/FastAPI *cancel* the task that is +When a client disconnects, ASGI servers such as Starlette/FastAPI *cancel* the code that is producing the response. If your response generator uses the async pattern, that cancellation unwinds the Burr action immediately. See [server.py](server.py): @@ -124,12 +144,29 @@ curl -N --max-time 1 localhost:8000/chat # hang up after a second # server log: # respond: open LLM connection # respond: close LLM connection -# post_run_step: respond exception=CancelledError(...) ``` -**To stop a stream from somewhere else** (another task, a "stop" endpoint), cancel the task that is -consuming it. Calling `close()` / `aclose()` while the stream is producing an item (from another task, -or from inside a hook) raises a `RuntimeError` that tells you to do this. +There is no `post_run_step` line. Starlette cancels through an anyio *cancel scope*. Unlike +`task.cancel()`, a cancel scope keeps cancelling every `await` inside it, including Burr's await on async +`post_run_step` hooks and async persisters. So for a disconnected step: + +- the action's cleanup **does** run (but as with any code under a cancelled scope, an `await` inside its + `finally` is cancelled too; wrap such cleanup in `anyio.CancelScope(shield=True)` if it must complete); +- async hooks, trackers and persisters **don't** run, so the stopped step isn't recorded; +- **resuming is unaffected**: the last saved step is the one before, so the interrupted action runs again, + as it would if it had been saved as `failed`. + +`client_disconnect_demo` reproduces this without a server: + +``` +respond: open LLM connection +server: sent 'Burr ' +server: sent 'streams ' +server: client disconnected -> cancel the response +respond: close LLM connection +persister: saved for this run: None +resume: next action is 'respond' (it runs again) +``` ## What gets reported @@ -137,7 +174,8 @@ or from inside a hook) raises a `RuntimeError` that tells you to do this. |---|---|---|---|---| | read to the end | `None` | `completed` | yes | next action | | `break` / `close()` / `aclose()` before the end | `GeneratorExit()` | `failed` | no | **same action runs again** | -| task cancelled (e.g. client disconnect) | `CancelledError()` | `failed` | no | **same action runs again** | +| consuming task cancelled (`task.cancel()`) | `CancelledError()` | `failed` | no | **same action runs again** | +| client disconnect (Starlette/FastAPI cancel scope) | async hooks not called | nothing saved | no | **same action runs again** | | the action raised | that exception | `failed` | no | same action runs again | | stopped early, and the action's cleanup raised | the cleanup error | `failed` | no | same action runs again | @@ -162,10 +200,9 @@ Note that `GeneratorExit` and `CancelledError` derive from `BaseException`, not and Burr logs a warning. A `break` itself doesn't raise it. If you want the error in your own code, call `close()` / `aclose()` afterwards and they raise it (`cleanup_error_demo`). A cleanup error never replaces an error that was already stopping the stream (such as a bug in a hook); it is logged instead. -- **Use sync hooks/persisters to record client disconnects.** On a disconnect, Starlette/FastAPI cancels - the whole response task. Sync `post_run_step` hooks (and sync persisters) still run and see the - `CancelledError`. *Async* `post_run_step` hooks (and async persisters) are cancelled before they start, - so they never record the stopped step. +- **Client disconnects aren't recorded by async hooks or persisters** (see [above](#web-servers-client-disconnects)). + Don't work around it by adding sync hooks to an async app: they block the event loop. Resume is + correct either way. See also the [streaming actions docs](https://burr.apache.org/concepts/streaming-actions/) and the [streaming-fastapi](../streaming-fastapi) example for a full chatbot. diff --git a/examples/streaming-early-stop/application.py b/examples/streaming-early-stop/application.py index 3bd26cbac..736af79e2 100644 --- a/examples/streaming-early-stop/application.py +++ b/examples/streaming-early-stop/application.py @@ -34,10 +34,12 @@ import logging from typing import AsyncGenerator, Generator, List, Optional, Tuple +import anyio + from burr.core import Application, ApplicationBuilder, State from burr.core.action import streaming_action -from burr.core.persistence import SQLLitePersister -from burr.lifecycle import PostRunStepHook +from burr.core.persistence import AsyncInMemoryPersister, SQLLitePersister +from burr.lifecycle import PostRunStepHook, PostRunStepHookAsync TOKENS = ["Burr ", "streams ", "tokens ", "one ", "at ", "a ", "time."] @@ -57,6 +59,11 @@ def print(self, title: str): print(f" {event}") +def _report(log: Timeline, action, exception: Optional[BaseException]): + outcome = "completed" if exception is None else f"failed with {exception!r}" + log(f"post_run_step: {action.name} {outcome}") + + class ReportStep(PostRunStepHook): """Stands in for any hook/tracker/persister: shows how each step was reported.""" @@ -64,8 +71,17 @@ def __init__(self, log: Timeline): self.log = log def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): - outcome = "completed" if exception is None else f"failed with {exception!r}" - self.log(f"post_run_step: {action.name} {outcome}") + _report(self.log, action, exception) + + +class ReportStepAsync(PostRunStepHookAsync): + """The same, for async applications: don't mix sync hooks into an async app.""" + + def __init__(self, log: Timeline): + self.log = log + + async def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): + _report(self.log, action, exception) def _respond_sync(log: Timeline, cleanup_error: bool = False): @@ -103,23 +119,35 @@ async def respond(state: State) -> AsyncGenerator[Tuple[dict, Optional[State]], return respond -def application( - respond, log: Timeline, persister: Optional[SQLLitePersister] = None -) -> Application: - builder = ( +def _builder(respond, hook) -> ApplicationBuilder: + return ( ApplicationBuilder() .with_actions(respond=respond) .with_transitions() - .with_entrypoint("respond") - .with_state(State()) - .with_hooks(ReportStep(log)) + .with_hooks(hook) .with_identifiers(app_id="early-stop-demo") ) + + +def application( + respond, log: Timeline, persister: Optional[SQLLitePersister] = None +) -> Application: + builder = _builder(respond, ReportStep(log)).with_entrypoint("respond").with_state(State()) if persister is not None: builder = builder.with_state_persister(persister) return builder.build() +async def async_application( + respond, log: Timeline, persister: Optional[AsyncInMemoryPersister] = None +) -> Application: + """Async app: async hooks, async persister, built with ``abuild()``.""" + builder = _builder(respond, ReportStepAsync(log)).with_entrypoint("respond").with_state(State()) + if persister is not None: + builder = builder.with_state_persister(persister) + return await builder.abuild() + + def wants_to_stop(token: str) -> bool: """Stand-in for "the user hit stop", a moderation check, a token budget, ...""" return token == "tokens " @@ -164,7 +192,7 @@ def sync_iterate_break_demo(): async def async_aclosing_demo(): log = Timeline() - app = application(_respond_async(log), log) + app = await async_application(_respond_async(log), log) _, container = await app.astream_result(halt_after=["respond"]) async with contextlib.aclosing(container): async for item in container: @@ -177,7 +205,7 @@ async def async_aclosing_demo(): async def async_iterate_aclosing_demo(): log = Timeline() - app = application(_respond_async(log), log) + app = await async_application(_respond_async(log), log) async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: async for _, container in steps: async for item in container: @@ -191,7 +219,7 @@ async def async_iterate_aclosing_demo(): async def async_break_without_aclose_demo(): log = Timeline() - app = application(_respond_async(log), log) + app = await async_application(_respond_async(log), log) _, container = await app.astream_result(halt_after=["respond"]) async for item in container: if wants_to_stop(item["token"]): @@ -203,27 +231,64 @@ async def async_break_without_aclose_demo(): # -------------------------------------------------------------------------------------------- -# Web servers: the client disconnects -> the server cancels the task consuming the stream. +# Cancelling the consuming task: from your own code, and from a web server on client disconnect. # -------------------------------------------------------------------------------------------- -async def client_disconnect_demo(): +async def cancel_task_demo(): + """E.g. a "stop" endpoint cancelling the task that streams to the user.""" log = Timeline() - app = application(_respond_async(log, delay=0.05), log) + app = await async_application(_respond_async(log, delay=0.05), log) - async def stream_to_client(): + async def stream_to_user(): async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: async for _, container in steps: async for item in container: - log(f"server: sent {item['token']!r}") + log(f"caller: got {item['token']!r}") - task = asyncio.create_task(stream_to_client()) + task = asyncio.create_task(stream_to_user()) await asyncio.sleep(0.12) - log("server: client disconnected -> task.cancel()") + log("caller: task.cancel()") task.cancel() with contextlib.suppress(asyncio.CancelledError): await task - log.print("Web server: client disconnect cancels the consuming task") + log.print("Cancelling the consuming task: hooks see CancelledError") + + +async def client_disconnect_demo(): + """What Starlette/FastAPI do on a client disconnect: cancel an anyio cancel scope. + + Unlike ``task.cancel()``, a cancel scope keeps cancelling every ``await`` inside it, so Burr's + await on the async ``post_run_step`` hooks is cancelled too. The action is still cleaned up, and + since nothing is saved for the stopped step, resuming re-runs it. + """ + log = Timeline() + persister = AsyncInMemoryPersister() + app = await async_application(_respond_async(log, delay=0.05), log, persister=persister) + + async def stream_to_client(): + async with contextlib.aclosing(app.astream_iterate(halt_after=["respond"])) as steps: + async for _, container in steps: + async for item in container: + log(f"server: sent {item['token']!r}") + + async with anyio.create_task_group() as response: + response.start_soon(stream_to_client) + await asyncio.sleep(0.12) + log("server: client disconnected -> cancel the response") + response.cancel_scope.cancel() + + log(f"persister: saved for this run: {await persister.load(None, 'early-stop-demo')!r}") + resumed = await ( + _builder(_respond_async(log), ReportStepAsync(log)) + .initialize_from( + persister, resume_at_next_action=True, default_state={}, default_entrypoint="respond" + ) + .with_state_persister(persister) + .abuild() + ) + log(f"resume: next action is {resumed.get_next_action().name!r} (it runs again)") + log.print("Web server: client disconnect (async hooks don't run, resume is unaffected)") # -------------------------------------------------------------------------------------------- @@ -289,6 +354,7 @@ def visualize(): asyncio.run(async_aclosing_demo()) asyncio.run(async_iterate_aclosing_demo()) asyncio.run(async_break_without_aclose_demo()) + asyncio.run(cancel_task_demo()) asyncio.run(client_disconnect_demo()) resume_after_stop_demo() cleanup_error_demo() diff --git a/examples/streaming-early-stop/notebook.ipynb b/examples/streaming-early-stop/notebook.ipynb index f93d89c23..40c9275af 100644 --- a/examples/streaming-early-stop/notebook.ipynb +++ b/examples/streaming-early-stop/notebook.ipynb @@ -6,10 +6,10 @@ "id": "00ae3074", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:54.703130Z", - "iopub.status.busy": "2026-11-03T19:18:54.703026Z", - "iopub.status.idle": "2026-11-03T19:18:54.706150Z", - "shell.execute_reply": "2026-11-03T19:18:54.705826Z" + "iopub.execute_input": "2026-11-03T23:36:49.777200Z", + "iopub.status.busy": "2026-11-03T23:36:49.776974Z", + "iopub.status.idle": "2026-11-03T23:36:49.781574Z", + "shell.execute_reply": "2026-11-03T23:36:49.781176Z" } }, "outputs": [], @@ -34,17 +34,155 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "bb1abd63", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:54.707835Z", - "iopub.status.busy": "2026-11-03T19:18:54.707725Z", - "iopub.status.idle": "2026-11-03T19:18:59.113436Z", - "shell.execute_reply": "2026-11-03T19:18:59.112668Z" + "iopub.execute_input": "2026-11-03T23:36:49.783696Z", + "iopub.status.busy": "2026-11-03T23:36:49.783531Z", + "iopub.status.idle": "2026-11-03T23:36:53.206048Z", + "shell.execute_reply": "2026-11-03T23:36:53.205418Z" } }, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pyenv: pip: command not found\r\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\r\n", + "The `pip' command exists in these Python versions:\r\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 3.11.11\r\n", + " 3.11.11/envs/afscript\r\n", + " 3.11.11/envs/afscript-py311\r\n", + " 3.11.11/envs/agent_sdk-py311\r\n", + " 3.11.11/envs/agent-graph-py311\r\n", + " 3.11.11/envs/agent-graph-TAG-dsl-py311\r\n", + " 3.11.11/envs/agentforce-lifecycle-sdk-python-py311\r\n", + " 3.11.11/envs/agentgraph-dsl-repo-py311\r\n", + " 3.11.11/envs/ai-sdlc-nov-2025\r\n", + " 3.11.11/envs/atlas-barry-py311\r\n", + " 3.11.11/envs/atlas-daisy-py311\r\n", + " 3.11.11/envs/burr-docs-py311\r\n", + " 3.11.11/envs/burr-py311\r\n", + " 3.11.11/envs/burr-rc-2-py311\r\n", + " 3.11.11/envs/burr-rc3-py311\r\n", + " 3.11.11/envs/burr-ui-py311\r\n", + " 3.11.11/envs/burr-worktree-py311\r\n", + " 3.11.11/envs/deep-researcher-py311\r\n", + " 3.11.11/envs/feature-rollout-py311\r\n", + " 3.11.11/envs/hamilton-docs-py311\r\n", + " 3.11.11/envs/hamilton-py311\r\n", + " 3.11.11/envs/hamilton-qualify-py311\r\n", + " 3.11.11/envs/llm-replay-py311\r\n", + " 3.11.11/envs/module-daisy-voice-py311\r\n", + " 3.11.11/envs/module-graph-runtime-py311\r\n", + " 3.11.11/envs/openllmetry-py311\r\n", + " 3.11.11/envs/pdx-server-py311\r\n", + " 3.11.11/envs/planner-python-chat-client-py311\r\n", + " 3.11.11/envs/prototypes-py311\r\n", + " 3.11.11/envs/search-gpt-py311\r\n", + " 3.11.11/envs/temp-env-py311\r\n", + " 3.11.11/envs/temp-py311\r\n", + " 3.11.11/envs/test-burr-py311\r\n", + " 3.11.11/envs/test-burr-whl-py311\r\n", + " 3.11.11/envs/test-env-py311\r\n", + " 3.11.11/envs/test-pdx-py311\r\n", + " 3.11.11/envs/test-py311\r\n", + " 3.11.11/envs/vibe-ui-py311\r\n", + " 3.11.11/envs/workshop-4\r\n", + " 3.12.9\r\n", + " 3.12.9/envs/agent-dsl-py312\r\n", + " 3.12.9/envs/burr-rc3-verify-py312\r\n", + " 3.12.9/envs/claude-clone-py312\r\n", + " 3.12.9/envs/debug-tools-py312\r\n", + " 3.12.9/envs/hamilton-py312\r\n", + " 3.12.9/envs/llmg_prompting-py312\r\n", + " 3.12.9/envs/pr-dashboard-py312\r\n", + " 3.12.9/envs/search-gpt-py312\r\n", + " 3.12.9/envs/temp-burr-py312\r\n", + " 3.12.9/envs/test-py312\r\n", + " 3.12.9/envs/validate-hamilton-py312\r\n", + " 3.12.9/envs/ws4-py312\r\n", + " 3.13.2\r\n", + " 3.13.2/envs/hamilton-gil-py313\r\n", + " 3.13.2/envs/runtime-docker-compose-py313\r\n", + " 3.13.2/envs/test\r\n", + " 3.13.2t\r\n", + " 3.13.2t/envs/hamilton-nogil-py313t\r\n", + " 3.14.6\r\n", + " afscript\r\n", + " afscript-py311\r\n", + " agentforce-lifecycle-sdk-python-py311\r\n", + " agentgraph-dsl-repo-py311\r\n", + " agent-dsl-py312\r\n", + " agent-graph-TAG-dsl-py311\r\n", + " agent-graph-py311\r\n", + " agent_sdk-py311\r\n", + " ai-sdlc-nov-2025\r\n", + " atlas-barry-py311\r\n", + " atlas-daisy-py311\r\n", + " burr-docs-py311\r\n", + " burr-py311\r\n", + " burr-rc3-py311\r\n", + " burr-rc3-verify-py312\r\n", + " burr-rc-2-py311\r\n", + " burr-ui-py311\r\n", + " burr-worktree-py311\r\n", + " claude-clone-py312\r\n", + " debug-tools-py312\r\n", + " deep-researcher-py311\r\n", + " feature-rollout-py311\r\n", + " hamilton-docs-py311\r\n", + " hamilton-gil-py313\r\n", + " hamilton-nogil-py313t\r\n", + " hamilton-py311\r\n", + " hamilton-py312\r\n", + " hamilton-qualify-py311\r\n", + " llmg_prompting-py312\r\n", + " llm-replay-py311\r\n", + " module-daisy-voice-py311\r\n", + " module-graph-runtime-py311\r\n", + " openllmetry-py311\r\n", + " pdx-server-py311\r\n", + " planner-python-chat-client-py311\r\n", + " prototypes-py311\r\n", + " pr-dashboard-py312\r\n", + " runtime-docker-compose-py313\r\n", + " search-gpt-py311\r\n", + " search-gpt-py312\r\n", + " temp-burr-py312\r\n", + " temp-env-py311\r\n", + " temp-py311\r\n", + " test\r\n", + " test-burr-py311\r\n", + " test-burr-whl-py311\r\n", + " test-env-py311\r\n", + " test-pdx-py311\r\n", + " test-py311\r\n", + " test-py312\r\n", + " validate-hamilton-py312\r\n", + " vibe-ui-py311\r\n", + " workshop-4\r\n", + " ws4-py312\r\n", + "\r\n", + "Note: See 'pyenv help global' for tips on allowing multiple\r\n", + " Python versions to be found at the same time.\r\n" + ] + } + ], "source": [ "!pip install apache-burr" ] @@ -67,10 +205,10 @@ "id": "5203dcf7", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.116209Z", - "iopub.status.busy": "2026-11-03T19:18:59.115942Z", - "iopub.status.idle": "2026-11-03T19:18:59.713846Z", - "shell.execute_reply": "2026-11-03T19:18:59.713427Z" + "iopub.execute_input": "2026-11-03T23:36:53.208767Z", + "iopub.status.busy": "2026-11-03T23:36:53.208588Z", + "iopub.status.idle": "2026-11-03T23:36:53.653899Z", + "shell.execute_reply": "2026-11-03T23:36:53.653617Z" } }, "outputs": [], @@ -97,10 +235,10 @@ "id": "28becda8", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.716713Z", - "iopub.status.busy": "2026-11-03T19:18:59.716419Z", - "iopub.status.idle": "2026-11-03T19:18:59.719557Z", - "shell.execute_reply": "2026-11-03T19:18:59.719255Z" + "iopub.execute_input": "2026-11-03T23:36:53.655407Z", + "iopub.status.busy": "2026-11-03T23:36:53.655275Z", + "iopub.status.idle": "2026-11-03T23:36:53.657242Z", + "shell.execute_reply": "2026-11-03T23:36:53.657025Z" } }, "outputs": [ @@ -131,10 +269,10 @@ "id": "ef586cdc", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.722774Z", - "iopub.status.busy": "2026-11-03T19:18:59.722528Z", - "iopub.status.idle": "2026-11-03T19:18:59.725193Z", - "shell.execute_reply": "2026-11-03T19:18:59.724691Z" + "iopub.execute_input": "2026-11-03T23:36:53.658355Z", + "iopub.status.busy": "2026-11-03T23:36:53.658277Z", + "iopub.status.idle": "2026-11-03T23:36:53.659995Z", + "shell.execute_reply": "2026-11-03T23:36:53.659802Z" } }, "outputs": [ @@ -172,10 +310,10 @@ "id": "90a16b98", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.727245Z", - "iopub.status.busy": "2026-11-03T19:18:59.727111Z", - "iopub.status.idle": "2026-11-03T19:18:59.730054Z", - "shell.execute_reply": "2026-11-03T19:18:59.729717Z" + "iopub.execute_input": "2026-11-03T23:36:53.661116Z", + "iopub.status.busy": "2026-11-03T23:36:53.661033Z", + "iopub.status.idle": "2026-11-03T23:36:53.663232Z", + "shell.execute_reply": "2026-11-03T23:36:53.663062Z" } }, "outputs": [ @@ -205,10 +343,10 @@ "id": "9b4dd2b0", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.732184Z", - "iopub.status.busy": "2026-11-03T19:18:59.732053Z", - "iopub.status.idle": "2026-11-03T19:18:59.734460Z", - "shell.execute_reply": "2026-11-03T19:18:59.734169Z" + "iopub.execute_input": "2026-11-03T23:36:53.664257Z", + "iopub.status.busy": "2026-11-03T23:36:53.664169Z", + "iopub.status.idle": "2026-11-03T23:36:53.666109Z", + "shell.execute_reply": "2026-11-03T23:36:53.665926Z" } }, "outputs": [ @@ -246,10 +384,10 @@ "id": "d14ced01", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.735901Z", - "iopub.status.busy": "2026-11-03T19:18:59.735787Z", - "iopub.status.idle": "2026-11-03T19:18:59.749733Z", - "shell.execute_reply": "2026-11-03T19:18:59.749049Z" + "iopub.execute_input": "2026-11-03T23:36:53.667165Z", + "iopub.status.busy": "2026-11-03T23:36:53.667096Z", + "iopub.status.idle": "2026-11-03T23:36:53.679343Z", + "shell.execute_reply": "2026-11-03T23:36:53.679104Z" } }, "outputs": [ @@ -273,24 +411,66 @@ }, { "cell_type": "markdown", - "id": "e6b971fa", + "id": "d33d5c0e", "metadata": {}, "source": [ - "## Web servers: a client disconnect cancels the consuming task\n", + "## Cancelling the consuming task\n", "\n", - "The hook receives `CancelledError`. See [server.py](server.py) for the FastAPI version." + "To stop a stream from another task, cancel the task that consumes it. Hooks receive `CancelledError`." ] }, { "cell_type": "code", "execution_count": 9, + "id": "607e678f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-11-03T23:36:53.680553Z", + "iopub.status.busy": "2026-11-03T23:36:53.680489Z", + "iopub.status.idle": "2026-11-03T23:36:53.803959Z", + "shell.execute_reply": "2026-11-03T23:36:53.803617Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Cancelling the consuming task: hooks see CancelledError\n", + " respond: open LLM connection\n", + " caller: got 'Burr '\n", + " caller: got 'streams '\n", + " caller: task.cancel()\n", + " respond: close LLM connection\n", + " post_run_step: respond failed with CancelledError()\n" + ] + } + ], + "source": [ + "await demo.cancel_task_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "6ed890b8", + "metadata": {}, + "source": [ + "## Web servers: a client disconnect\n", + "\n", + "Starlette/FastAPI cancel the response through an anyio cancel scope. The action is cleaned up, but async hooks and persisters don't run for the stopped step, so it isn't recorded. Resuming still re-runs it. See [server.py](server.py) for the FastAPI version." + ] + }, + { + "cell_type": "code", + "execution_count": 10, "id": "144fdfac", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.751621Z", - "iopub.status.busy": "2026-11-03T19:18:59.751469Z", - "iopub.status.idle": "2026-11-03T19:18:59.875523Z", - "shell.execute_reply": "2026-11-03T19:18:59.875184Z" + "iopub.execute_input": "2026-11-03T23:36:53.805739Z", + "iopub.status.busy": "2026-11-03T23:36:53.805615Z", + "iopub.status.idle": "2026-11-03T23:36:53.934744Z", + "shell.execute_reply": "2026-11-03T23:36:53.934242Z" } }, "outputs": [ @@ -299,13 +479,14 @@ "output_type": "stream", "text": [ "\n", - "=== Web server: client disconnect cancels the consuming task\n", + "=== Web server: client disconnect (async hooks don't run, resume is unaffected)\n", " respond: open LLM connection\n", " server: sent 'Burr '\n", " server: sent 'streams '\n", - " server: client disconnected -> task.cancel()\n", + " server: client disconnected -> cancel the response\n", " respond: close LLM connection\n", - " post_run_step: respond failed with CancelledError()\n" + " persister: saved for this run: None\n", + " resume: next action is 'respond' (it runs again)\n" ] } ], @@ -323,14 +504,14 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 11, "id": "5a6ec507", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.877311Z", - "iopub.status.busy": "2026-11-03T19:18:59.877179Z", - "iopub.status.idle": "2026-11-03T19:18:59.880004Z", - "shell.execute_reply": "2026-11-03T19:18:59.879722Z" + "iopub.execute_input": "2026-11-03T23:36:53.937353Z", + "iopub.status.busy": "2026-11-03T23:36:53.937016Z", + "iopub.status.idle": "2026-11-03T23:36:53.941045Z", + "shell.execute_reply": "2026-11-03T23:36:53.940627Z" } }, "outputs": [ @@ -362,14 +543,14 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 12, "id": "ef2ae65c", "metadata": { "execution": { - "iopub.execute_input": "2026-11-03T19:18:59.881678Z", - "iopub.status.busy": "2026-11-03T19:18:59.881554Z", - "iopub.status.idle": "2026-11-03T19:18:59.883520Z", - "shell.execute_reply": "2026-11-03T19:18:59.883266Z" + "iopub.execute_input": "2026-11-03T23:36:53.943098Z", + "iopub.status.busy": "2026-11-03T23:36:53.942963Z", + "iopub.status.idle": "2026-11-03T23:36:53.945458Z", + "shell.execute_reply": "2026-11-03T23:36:53.945143Z" } }, "outputs": [ diff --git a/examples/streaming-early-stop/server.py b/examples/streaming-early-stop/server.py index f8eac66e4..b28a749e1 100644 --- a/examples/streaming-early-stop/server.py +++ b/examples/streaming-early-stop/server.py @@ -19,7 +19,11 @@ Run with ``uvicorn server:app``, then ``curl -N --max-time 1 localhost:8000/chat`` (curl hangs up after a second, like a user closing the tab). The server log shows the LLM connection being closed -right away and the step reported to hooks as failed with ``CancelledError``. +right away. + +The ``post_run_step`` line does *not* appear: on a disconnect, Starlette cancels the response through an +anyio cancel scope, which also cancels Burr's await on async hooks and persisters, so they don't run for +the stopped step. Nothing is saved for it, so resuming re-runs it. See the README. """ import asyncio @@ -32,7 +36,7 @@ from burr.core import ApplicationBuilder, State from burr.core.action import streaming_action -from burr.lifecycle import PostRunStepHook +from burr.lifecycle import PostRunStepHookAsync logger = logging.getLogger("uvicorn.error") app = FastAPI() @@ -53,27 +57,26 @@ async def respond(state: State) -> AsyncGenerator[Tuple[dict, Optional[State]], logger.info("respond: close LLM connection") -class LogStep(PostRunStepHook): - # A *sync* hook: on a client disconnect, async post_run_step hooks are cancelled before they run. - def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): +class LogStep(PostRunStepHookAsync): + async def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): logger.info("post_run_step: %s exception=%r", action.name, exception) -def chat_application(): - return ( +async def chat_application(): + return await ( ApplicationBuilder() .with_actions(respond=respond) .with_transitions() .with_entrypoint("respond") .with_state(State()) .with_hooks(LogStep()) - .build() + .abuild() ) @app.get("/chat") async def chat() -> StreamingResponse: - burr_app = chat_application() + burr_app = await chat_application() async def sse(): # `aclosing` makes leaving this block -- normally, on error, or because the server diff --git a/pyproject.toml b/pyproject.toml index f85c2b3f4..4a3d15843 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,6 +89,7 @@ release = [ ] tests = [ + "anyio", "pytest", "pytest-asyncio", "apache-burr[graphviz]", diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 6c337cae2..291f95a4c 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -25,6 +25,7 @@ import uuid from typing import Any, Awaitable, Callable, Dict, Generator, Literal, Optional, Tuple, Union +import anyio import pytest from burr.core import State @@ -3066,6 +3067,63 @@ async def respond_to_client(): assert persister.load("", "early-stop")["status"] == "failed" +async def test_astream_iterate_cancel_scope_cleans_up_and_resume_reruns_the_step(): + """Starlette cancels a disconnected response through an anyio cancel scope, which also cancels + Burr's await on async hooks/persisters, so they don't record the step (a known limitation). The + action must still be cleaned up, and resuming must still land on the interrupted action.""" + + @action(reads=[], writes=["x"]) + def before(state: State) -> State: + return state.update(x=0) + + log = _EarlyStopLog() + persister = AsyncInMemoryPersister() + + @streaming_action(reads=[], writes=["x"]) + async def act(state: State): + try: + for item in ["a", "b", "c"]: + await asyncio.sleep(0.05) + yield {"item": item}, None + yield {"item": "done"}, state.update(x=1) + finally: + # Recorded before awaiting: under a cancelled scope, awaits in `finally` are cancelled too. + log.events.append("cleanup") + + async def build(resume: bool) -> Application: + builder = ( + ApplicationBuilder() + .with_actions(before=before, act=act) + .with_transitions(("before", "act")) + .with_hooks(*log.hooks(is_async=True)) + .with_identifiers(app_id="early-stop") + .with_state_persister(persister) + ) + if resume: + builder = builder.initialize_from( + persister, resume_at_next_action=True, default_state={}, default_entrypoint="before" + ) + else: + builder = builder.with_entrypoint("before").with_state(State()) + return await builder.abuild() + + app = await build(resume=False) + + async def respond_to_client(scope: anyio.CancelScope): + async with contextlib.aclosing(app.astream_iterate(halt_after=["act"])) as steps: + async for step, container in steps: + async for _ in container: + if step.name == "act": + scope.cancel() + + async with anyio.create_task_group() as response: + response.start_soon(respond_to_client, response.cancel_scope) + + assert log.events == [("reported", None), "cleanup"] + assert (await persister.load(None, "early-stop"))["position"] == "before" + assert (await build(resume=True)).get_next_action().name == "act" + + @pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) async def test_stopped_step_persists_failed_and_resume_reruns_it(is_async: bool): log = _EarlyStopLog() diff --git a/tests/core/test_streaming_early_stop_example.py b/tests/core/test_streaming_early_stop_example.py index ebe4db94c..6f8047392 100644 --- a/tests/core/test_streaming_early_stop_example.py +++ b/tests/core/test_streaming_early_stop_example.py @@ -76,9 +76,27 @@ async def test_async_break_without_aclose_defers_cleanup(demo, capsys): ) +async def test_cancel_task(demo, capsys): + await demo.cancel_task_demo() + assert _timeline(capsys)[-3:] == [ + "caller: task.cancel()", + CLOSED, + "post_run_step: respond failed with CancelledError()", + ] + + async def test_client_disconnect(demo, capsys): + """A cancel scope (Starlette's disconnect) also cancels the await on async hooks/persisters: + the action is cleaned up, the step isn't recorded, and resuming still re-runs it.""" await demo.client_disconnect_demo() - assert _timeline(capsys)[-2:] == [CLOSED, "post_run_step: respond failed with CancelledError()"] + events = _timeline(capsys) + assert not any(event.startswith("post_run_step") for event in events) + assert events[-4:] == [ + "server: client disconnected -> cancel the response", + CLOSED, + "persister: saved for this run: None", + "resume: next action is 'respond' (it runs again)", + ] def test_resume_after_stop(demo, capsys): From 049ee0ab70db7bfd93b2ff8ddc701af52ad63a3b Mon Sep 17 00:00:00 2001 From: Stefan Krawczyk Date: Wed, 4 Nov 2026 09:26:03 -0800 Subject: [PATCH 8/8] fix(core): make close()/aclose() errors consistent and repeatable When stopping a stream raised, the error depended on sync vs async and on which call it was: async let a post_run_step hook error replace the cleanup error on the first aclose() and then raised the cleanup error afterwards; a hook error alone was raised once and then lost, although the step had not been recorded. Contract, identical for sync and async: - the first error wins: an error from the action's cleanup (what hooks receive), else an error from the post_run_step callback; - a callback error superseded by a cleanup error is logged, not lost; - every close()/aclose() raises that same error. Also pin and document that under an anyio cancel scope (Starlette client disconnect) an await inside the action's own cleanup is cancelled unless shielded, instead of claiming the action "is still cleaned up". Reported by @xihaian251 in #1027. --- burr/core/action.py | 79 +++++++++---- docs/concepts/streaming-actions.rst | 23 ++-- examples/streaming-early-stop/README.md | 22 +++- examples/streaming-early-stop/application.py | 7 +- examples/streaming-early-stop/server.py | 4 +- tests/core/test_application.py | 114 +++++++++++++++++-- 6 files changed, 200 insertions(+), 49 deletions(-) diff --git a/burr/core/action.py b/burr/core/action.py index 0d49e878c..42f515b62 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -1438,11 +1438,21 @@ def _resolve_close_future(future: "asyncio.Future", exc: Optional[BaseException] future.set_result(None) -def _log_abandoned_cleanup_error() -> None: +def _log_abandoned_stop_error() -> None: # Nobody is left to raise this to: the iterator is being closed by ``break`` or garbage - # collection. Hooks have already seen it, and ``close()``/``aclose()`` will raise it. + # collection. ``close()``/``aclose()`` will raise it. logger.warning( - "Cleanup of a streaming action raised after the stream was stopped early", exc_info=True + "Stopping a streaming action early raised (in its cleanup or a post_run_step hook); " + "close() / aclose() raise it again", + exc_info=True, + ) + + +def _log_superseded_callback_error(error: BaseException, callback_error: Exception) -> None: + logger.warning( + "A post_run_step hook raised while reporting a streaming action stopped by %r", + error, + exc_info=(type(callback_error), callback_error, callback_error.__traceback__), ) @@ -1521,7 +1531,7 @@ def __init__( self._result = None self._callback_realized = False self._closed = False - # Cleanup error from the first close; later ``close()`` calls raise it again. + # Error from the first close (cleanup, else the callback); later calls raise it again. self._close_error: Optional[BaseException] = None def _realize_callback(self, exc: Optional[BaseException]) -> None: @@ -1537,9 +1547,13 @@ def close(self) -> None: Closes the action's generator chain so its cleanup (``finally`` blocks, ``with`` exits) runs now rather than at garbage collection, then runs the post-step callback exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a - no-op once the stream has completed. An exception raised by the action's cleanup - propagates, and is raised again by later calls -- including after ``for ... break`` - already closed the stream. + no-op once the stream has completed. + + If stopping raises, the first error wins: an error from the action's cleanup (which is + also what the callback receives), else an error from the callback (e.g. a + ``post_run_step`` hook). A callback error superseded by a cleanup error is logged. Every + call raises that same error -- including after ``for ... break`` already closed the + stream. A stream closed before it completes is reported to the callback (and so to ``post_run_step`` hooks and persisters) with a ``GeneratorExit``, i.e. as a failed step. @@ -1557,13 +1571,21 @@ def close(self) -> None: self._closed = True # Closing before completion means the step did not finish: report it as failed. exc: Optional[BaseException] = GeneratorExit() if self._result is None else None + error: Optional[BaseException] = None try: self.streaming_result_generator.close() except BaseException as e: - exc = self._close_error = e - raise - finally: + exc = error = e + try: self._realize_callback(exc) + except Exception as e: + if error is None: + error = e + else: + _log_superseded_callback_error(error, e) + if error is not None: + self._close_error = error + raise error def __next__(self) -> StreamResultType: if self._result is not None: @@ -1590,7 +1612,7 @@ def gen_fn(): try: self.close() except Exception: - _log_abandoned_cleanup_error() + _log_abandoned_stop_error() raise finally: self._realize_callback(sys.exc_info()[1]) @@ -1679,8 +1701,12 @@ async def aclose(self) -> None: Closes the action's generator chain so its cleanup (``finally`` blocks, ``async with`` exits) runs now rather than at garbage collection, then runs the post-step callback exactly once. ``get()`` afterwards returns ``(None, initial_state)``. Idempotent, and a - no-op once the stream has completed. An exception raised by the action's cleanup - propagates, and is raised again by later calls. + no-op once the stream has completed. + + If stopping raises, the first error wins: an error from the action's cleanup (which is + also what the callback receives), else an error from the callback (e.g. a + ``post_run_step`` hook). A callback error superseded by a cleanup error is logged. Every + call raises that same error. A stream closed before it completes is reported to the callback (and so to ``post_run_step`` hooks and persisters) with a ``GeneratorExit``, i.e. as a failed step. @@ -1700,17 +1726,26 @@ async def aclose(self) -> None: self._close_future = asyncio.get_running_loop().create_future() # Closing before completion means the step did not finish: report it as failed. exc: Optional[BaseException] = GeneratorExit() if self._result is None else None - raised = False + error: Optional[BaseException] = None try: - await self.streaming_result_generator.aclose() - except BaseException as e: - exc, raised = e, True - raise - finally: + try: + await self.streaming_result_generator.aclose() + except BaseException as e: + exc = error = e try: await self._realize_callback(exc) - finally: - _resolve_close_future(self._close_future, exc if raised else None) + except Exception as e: + if error is None: + error = e + else: + _log_superseded_callback_error(error, e) + except BaseException as e: # e.g. cancelled while reporting + error = e + raise + finally: + _resolve_close_future(self._close_future, error) + if error is not None: + raise error async def __anext__(self) -> StreamResultType: """Moves to the next state in the streaming result""" @@ -1746,7 +1781,7 @@ async def gen_fn(): try: await self.aclose() except Exception: - _log_abandoned_cleanup_error() + _log_abandoned_stop_error() else: await asyncio.wait([self._close_future]) raise diff --git a/docs/concepts/streaming-actions.rst b/docs/concepts/streaming-actions.rst index 0113d4c7a..854d806ca 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -274,13 +274,22 @@ A step that is stopped before it completes does not update state, and is reporte with the exception that stopped it -- ``GeneratorExit`` when it is closed, ``CancelledError`` when the consuming task is cancelled. Persisters record it as ``failed``, so resuming re-runs that action. -One exception: on a client disconnect, Starlette/FastAPI cancel the response through an anyio cancel scope, which -also cancels Burr's await on *async* ``post_run_step`` hooks and persisters. The action is still cleaned up, but -the stopped step isn't recorded. Resuming is unaffected: the last saved step is the one before, so the -interrupted action runs again. Don't add sync hooks to an async application to work around this; they block the -event loop. If the action's cleanup -itself raises, hooks receive that error, and ``close()`` / ``aclose()`` raise it. A cleanup error never replaces -an error that was already stopping the stream; it is logged instead. +If stopping raises, the first error wins. An error from the action's cleanup is what hooks receive, and what +``close()`` / ``aclose()`` raise; failing that, an error from a ``post_run_step`` hook is raised. Every later +``close()`` / ``aclose()`` raises the same error. After a bare ``break`` there is nobody to raise it to, so it is +logged. An error raised while stopping never replaces an error that was already stopping the stream (such as a +bug in a hook); it is logged instead. + +**Client disconnects** are an exception to the reporting above. Starlette/FastAPI cancel a disconnected response through an anyio +cancel scope, which keeps cancelling every ``await`` inside it: + +* the action's cleanup starts, but an ``await`` inside it is cancelled too; wrap cleanup that must finish in + ``anyio.CancelScope(shield=True)``; +* Burr's await on *async* ``post_run_step`` hooks and persisters is cancelled, so the stopped step isn't + recorded; +* resuming is unaffected: the last saved step is the one before, so the interrupted action runs again. + +Don't add sync hooks to an async application to work around this; they block the event loop. See the `streaming-early-stop example `_ for runnable versions of these patterns, including a FastAPI endpoint that handles client disconnects. diff --git a/examples/streaming-early-stop/README.md b/examples/streaming-early-stop/README.md index 914323601..202cedc65 100644 --- a/examples/streaming-early-stop/README.md +++ b/examples/streaming-early-stop/README.md @@ -150,8 +150,16 @@ There is no `post_run_step` line. Starlette cancels through an anyio *cancel sco `task.cancel()`, a cancel scope keeps cancelling every `await` inside it, including Burr's await on async `post_run_step` hooks and async persisters. So for a disconnected step: -- the action's cleanup **does** run (but as with any code under a cancelled scope, an `await` inside its - `finally` is cancelled too; wrap such cleanup in `anyio.CancelScope(shield=True)` if it must complete); +- the action's cleanup **starts**, but an `await` inside it is cancelled too, as for any code under a + cancelled scope. Wrap cleanup that must finish (e.g. releasing a connection) in + `anyio.CancelScope(shield=True)`: + + ```python + finally: + with anyio.CancelScope(shield=True): + await connection.release() + ``` + - async hooks, trackers and persisters **don't** run, so the stopped step isn't recorded; - **resuming is unaffected**: the last saved step is the one before, so the interrupted action runs again, as it would if it had been saved as `failed`. @@ -196,10 +204,12 @@ Note that `GeneratorExit` and `CancelledError` derive from `BaseException`, not ## Things to know -- **Cleanup errors.** If the action's `finally` raises while it is being stopped, hooks receive that error, - and Burr logs a warning. A `break` itself doesn't raise it. If you want the error in your own code, call - `close()` / `aclose()` afterwards and they raise it (`cleanup_error_demo`). A cleanup error never - replaces an error that was already stopping the stream (such as a bug in a hook); it is logged instead. +- **Errors while stopping: the first one wins.** If the action's `finally` raises while it is being stopped, + hooks receive that error, and `close()` / `aclose()` raise it (`cleanup_error_demo`). Otherwise, an error + from a `post_run_step` hook is raised. Every later `close()` / `aclose()` raises the same error. A `break` + itself doesn't raise; Burr logs a warning and `close()` / `aclose()` raise it afterwards. An error raised + while stopping never replaces an error that was already stopping the stream (such as a bug in a hook); it is + logged instead. - **Client disconnects aren't recorded by async hooks or persisters** (see [above](#web-servers-client-disconnects)). Don't work around it by adding sync hooks to an async app: they block the event loop. Resume is correct either way. diff --git a/examples/streaming-early-stop/application.py b/examples/streaming-early-stop/application.py index 736af79e2..a76fcb671 100644 --- a/examples/streaming-early-stop/application.py +++ b/examples/streaming-early-stop/application.py @@ -259,8 +259,9 @@ async def client_disconnect_demo(): """What Starlette/FastAPI do on a client disconnect: cancel an anyio cancel scope. Unlike ``task.cancel()``, a cancel scope keeps cancelling every ``await`` inside it, so Burr's - await on the async ``post_run_step`` hooks is cancelled too. The action is still cleaned up, and - since nothing is saved for the stopped step, resuming re-runs it. + await on the async ``post_run_step`` hooks is cancelled too. The action's cleanup starts (an + ``await`` in it would be cancelled too, unless shielded), and since nothing is saved for the + stopped step, resuming re-runs it. """ log = Timeline() persister = AsyncInMemoryPersister() @@ -322,7 +323,7 @@ def resume_after_stop_demo(): # -------------------------------------------------------------------------------------------- -# Cleanup errors: hooks always see them; close() raises them if you want them. +# Cleanup errors: hooks always see them; close() raises them (every time) if you want them. # -------------------------------------------------------------------------------------------- diff --git a/examples/streaming-early-stop/server.py b/examples/streaming-early-stop/server.py index b28a749e1..fa5e87339 100644 --- a/examples/streaming-early-stop/server.py +++ b/examples/streaming-early-stop/server.py @@ -23,7 +23,9 @@ The ``post_run_step`` line does *not* appear: on a disconnect, Starlette cancels the response through an anyio cancel scope, which also cancels Burr's await on async hooks and persisters, so they don't run for -the stopped step. Nothing is saved for it, so resuming re-runs it. See the README. +the stopped step. Nothing is saved for it, so resuming re-runs it. An ``await`` in the action's own +cleanup would be cancelled too; shield it with ``anyio.CancelScope(shield=True)`` if it must finish. +See the README. """ import asyncio diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 291f95a4c..d47d1170d 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -3067,10 +3067,15 @@ async def respond_to_client(): assert persister.load("", "early-stop")["status"] == "failed" -async def test_astream_iterate_cancel_scope_cleans_up_and_resume_reruns_the_step(): - """Starlette cancels a disconnected response through an anyio cancel scope, which also cancels - Burr's await on async hooks/persisters, so they don't record the step (a known limitation). The - action must still be cleaned up, and resuming must still land on the interrupted action.""" +@pytest.mark.parametrize("shield", [False, True], ids=["unshielded", "shielded"]) +async def test_astream_iterate_cancel_scope_runs_cleanup_and_resume_reruns_the_step(shield: bool): + """Starlette cancels a disconnected response through an anyio cancel scope. That scope keeps + cancelling every ``await`` inside it, so (as for any code under it, not specific to Burr): + + - the action's cleanup starts, but an ``await`` in it is cancelled unless shielded; + - Burr's await on async ``post_run_step`` hooks/persisters is cancelled: the step isn't recorded; + - resuming still lands on the interrupted action (the last saved step is the one before). + """ @action(reads=[], writes=["x"]) def before(state: State) -> State: @@ -3079,16 +3084,26 @@ def before(state: State) -> State: log = _EarlyStopLog() persister = AsyncInMemoryPersister() + class Connection: + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info): + log.events.append("release started") + if shield: + with anyio.CancelScope(shield=True): + await asyncio.sleep(0) + else: + await asyncio.sleep(0) + log.events.append("release finished") + @streaming_action(reads=[], writes=["x"]) async def act(state: State): - try: + async with Connection(): for item in ["a", "b", "c"]: await asyncio.sleep(0.05) yield {"item": item}, None yield {"item": "done"}, state.update(x=1) - finally: - # Recorded before awaiting: under a cancelled scope, awaits in `finally` are cancelled too. - log.events.append("cleanup") async def build(resume: bool) -> Application: builder = ( @@ -3119,7 +3134,8 @@ async def respond_to_client(scope: anyio.CancelScope): async with anyio.create_task_group() as response: response.start_soon(respond_to_client, response.cancel_scope) - assert log.events == [("reported", None), "cleanup"] + released = ["release started", "release finished"] if shield else ["release started"] + assert log.events == [("reported", None), *released] assert (await persister.load(None, "early-stop"))["position"] == "before" assert (await build(resume=True)).get_next_action().name == "act" @@ -3195,7 +3211,7 @@ def test_stream_result_cleanup_error_is_reported_and_raised_by_close(no_cyclic_g for _ in container: break assert log.events == ["cleanup", ("reported", _EarlyStopCleanupError)] - assert "Cleanup of a streaming action raised" in caplog.text + assert "Stopping a streaming action early raised" in caplog.text with pytest.raises(_EarlyStopCleanupError): container.close() with pytest.raises(_EarlyStopCleanupError): @@ -3220,6 +3236,84 @@ async def test_astream_result_cleanup_error_is_reported_and_raised_by_aclose( assert log.events == ["cleanup", ("reported", _EarlyStopCleanupError)] +class _EarlyStopHookError(Exception): + pass + + +def _raising_post_run_step_hook(is_async: bool): + class RaisingHook(PostRunStepHook): + def post_run_step(self, **future_kwargs): + raise _EarlyStopHookError("bug in post_run_step") + + class RaisingHookAsync(PostRunStepHookAsync): + async def post_run_step(self, **future_kwargs): + raise _EarlyStopHookError("bug in post_run_step") + + return RaisingHookAsync() if is_async else RaisingHook() + + +@pytest.mark.parametrize("how", ["close", "break_then_close"]) +@pytest.mark.parametrize( + "cleanup_raises", [True, False], ids=["cleanup_and_hook_raise", "only_hook_raises"] +) +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_close_raises_the_first_error_and_the_same_one_every_time( + is_async: bool, cleanup_raises: bool, how: str, no_cyclic_gc, caplog +): + """Contract for errors while stopping a stream, identical for sync and async: + + - the first error wins: a cleanup error, else the ``post_run_step`` hook's error; + - every ``close()`` / ``aclose()`` call raises that same error (the step is not re-reported); + - hooks see the cleanup error; a hook error superseded by it is logged, not lost. + """ + log = _EarlyStopLog() + app = _early_stop_app( + _early_stop_action(log, is_async, cleanup_raises=cleanup_raises), + log, + is_async, + hooks=(_raising_post_run_step_hook(is_async),), + ) + expected = _EarlyStopCleanupError if cleanup_raises else _EarlyStopHookError + if is_async: + _, container = await app.astream_result(halt_after=["act"]) + if how == "close": + iterator = container.__aiter__() + await iterator.__anext__() + else: + async for _ in container: + break # the stop's error has nobody to go to here; aclose() raises it + await asyncio.sleep(0) + raised = [] + for _ in range(3): + with pytest.raises(expected) as e: + await container.aclose() + raised.append(e.value) + if how == "close": + await iterator.aclose() # leaving the loop afterwards doesn't raise again + else: + _, container = app.stream_result(halt_after=["act"]) + if how == "close": + iterator = iter(container) + next(iterator) + else: + for _ in container: + break + raised = [] + for _ in range(3): + with pytest.raises(expected) as e: + container.close() + raised.append(e.value) + if how == "close": + iterator.close() + assert raised[0] is raised[1] is raised[2] + assert log.events == [ + "cleanup", + ("reported", _EarlyStopCleanupError if cleanup_raises else GeneratorExit), + ] + if cleanup_raises: + assert "bug in post_run_step" in caplog.text + + class _BrokenStreamItemHook(PostStreamItemHook): def post_stream_item(self, **future_kwargs): raise KeyError("bug in a hook")