diff --git a/burr/core/action.py b/burr/core/action.py index b71940d47..42f515b62 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -17,9 +17,12 @@ import abc import ast +import asyncio import builtins +import contextlib import copy import inspect +import logging import sys import textwrap import types @@ -54,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: @@ -152,8 +157,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 @@ -1423,6 +1429,40 @@ 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_stop_error() -> None: + # Nobody is left to raise this to: the iterator is being closed by ``break`` or garbage + # collection. ``close()``/``aclose()`` will raise it. + logger.warning( + "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__), + ) + + +_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: @@ -1472,7 +1512,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. @@ -1490,6 +1530,62 @@ def __init__( self._initial_state = initial_state self._result = None self._callback_realized = False + self._closed = False + # 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: + 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. + + 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. + + 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 + error: Optional[BaseException] = None + try: + self.streaming_result_generator.close() + except BaseException as e: + 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: @@ -1509,13 +1605,17 @@ 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_stop_error() + raise 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 @@ -1563,7 +1663,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], ], ): @@ -1582,13 +1682,81 @@ def __init__( self._callback = callback self._result = None self._callback_realized = False + # 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: + 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. + + 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. + + 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 + error: Optional[BaseException] = None + try: + try: + await self.streaming_result_generator.aclose() + except BaseException as e: + exc = error = e + try: + await self._realize_callback(exc) + 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""" 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 @@ -1605,13 +1773,20 @@ 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_stop_error() + else: + await asyncio.wait([self._close_future]) + raise 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() @@ -1634,7 +1809,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 @@ -1741,8 +1916,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..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, @@ -387,6 +426,8 @@ def _run_single_step_streaming_action( e, exc_info=True, ) + finally: + _close_action_generator(generator, action.name, sys.exc_info()[1]) if state_update is None: raise ValueError( @@ -456,6 +497,8 @@ async def _arun_single_step_streaming_action( e, exc_info=True, ) + finally: + await _aclose_action_generator(generator, action.name, sys.exc_info()[1]) if state_update is None: raise ValueError( @@ -527,6 +570,8 @@ def _run_multi_step_streaming_action( e, exc_info=True, ) + finally: + _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) @@ -585,6 +630,8 @@ async def _arun_multi_step_streaming_action( e, exc_info=True, ) + finally: + 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) @@ -1439,7 +1486,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: @@ -1566,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", @@ -1694,7 +1742,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: @@ -1821,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", @@ -1975,7 +2024,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): @@ -2012,7 +2066,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..e0457b233 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 @@ -346,23 +347,25 @@ 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 ) -> 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/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 89b126935..854d806ca 100644 --- a/docs/concepts/streaming-actions.rst +++ b/docs/concepts/streaming-actions.rst @@ -224,6 +224,76 @@ 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 +~~~~~~~~~~~~~~~~~~~~~~~ + +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 + + action, streaming_result = application.stream_result( + halt_after='streaming_response', inputs={"prompt": prompt} + ) + for result in streaming_result: + if should_stop(result): + 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 + + action, async_streaming_result = await application.astream_result( + halt_after='streaming_response', inputs={"prompt": prompt} + ) + 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: + ... + +``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. Persisters record it as ``failed``, so resuming re-runs that action. + +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. + -------------- 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..202cedc65 --- /dev/null +++ b/examples/streaming-early-stop/README.md @@ -0,0 +1,218 @@ + + +# 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), 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 (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: + +| | 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 | + +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: + +```bash +python application.py +``` + +![statemachine](statemachine.png) + +## 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 # the action's cleanup runs here, then the step is reported + +result, state = container.get() # (None, state without the update) +``` + +``` +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 +``` + +Breaking out of `app.stream_iterate(...)` works the same way. + +## Async: `break` inside `contextlib.aclosing` + +```python +import contextlib + +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() +``` + +`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 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): + +```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 +``` + +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 **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`. + +`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 + +| How the stream ended | `post_run_step(exception=...)` | persister status | state updated? | on resume | +|---|---|---|---|---| +| read to the end | `None` | `completed` | yes | next action | +| `break` / `close()` / `aclose()` before the end | `GeneratorExit()` | `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 | + +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): + 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 + +- **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. + +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..a76fcb671 --- /dev/null +++ b/examples/streaming-early-stop/application.py @@ -0,0 +1,361 @@ +# 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 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 + +import anyio + +from burr.core import Application, ApplicationBuilder, State +from burr.core.action import streaming_action +from burr.core.persistence import AsyncInMemoryPersister, SQLLitePersister +from burr.lifecycle import PostRunStepHook, PostRunStepHookAsync + +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}") + + +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.""" + + def __init__(self, log: Timeline): + self.log = log + + def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): + _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): + @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 _builder(respond, hook) -> ApplicationBuilder: + return ( + ApplicationBuilder() + .with_actions(respond=respond) + .with_transitions() + .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 " + + +# -------------------------------------------------------------------------------------------- +# Sync: just break. +# -------------------------------------------------------------------------------------------- + + +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 # 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("Sync: break") + + +def sync_iterate_break_demo(): + log = Timeline() + 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") + + +# -------------------------------------------------------------------------------------------- +# Async: break inside `async with contextlib.aclosing(...)`. +# -------------------------------------------------------------------------------------------- + + +async def async_aclosing_demo(): + log = Timeline() + 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: + 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 = 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: + log(f"caller: got {item['token']!r}") + if wants_to_stop(item["token"]): + break + break + 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 = 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"]): + 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)") + + +# -------------------------------------------------------------------------------------------- +# Cancelling the consuming task: from your own code, and from a web server on client disconnect. +# -------------------------------------------------------------------------------------------- + + +async def cancel_task_demo(): + """E.g. a "stop" endpoint cancelling the task that streams to the user.""" + log = Timeline() + app = await async_application(_respond_async(log, delay=0.05), log) + + 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"caller: got {item['token']!r}") + + task = asyncio.create_task(stream_to_user()) + await asyncio.sleep(0.12) + log("caller: task.cancel()") + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await 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'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() + 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)") + + +# -------------------------------------------------------------------------------------------- +# 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 + 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: hooks always see them; close() raises them (every time) if you want them. +# -------------------------------------------------------------------------------------------- + + +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: + if wants_to_stop(item["token"]): + 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: reported to hooks, and raised by close()") + + +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__": + 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(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 new file mode 100644 index 000000000..40c9275af --- /dev/null +++ b/examples/streaming-early-stop/notebook.ipynb @@ -0,0 +1,596 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "00ae3074", + "metadata": { + "execution": { + "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": [], + "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": 2, + "id": "bb1abd63", + "metadata": { + "execution": { + "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": [ + { + "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" + ] + }, + { + "cell_type": "markdown", + "id": "6e825efb", + "metadata": {}, + "source": [ + "# Stopping a streaming action early\n", + "\n", + "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": "5203dcf7", + "metadata": { + "execution": { + "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": [], + "source": [ + "import asyncio\n", + "import logging\n", + "\n", + "import application as demo\n", + "\n", + "logging.basicConfig(level=logging.ERROR)" + ] + }, + { + "cell_type": "markdown", + "id": "2bc86470", + "metadata": {}, + "source": [ + "## Sync: just `break`" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "28becda8", + "metadata": { + "execution": { + "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": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Sync: break\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: after the loop\n", + " caller: get() -> result=None, 'response' in state=False\n" + ] + } + ], + "source": [ + "demo.sync_break_demo()" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ef586cdc", + "metadata": { + "execution": { + "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": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\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: after the loop\n" + ] + } + ], + "source": [ + "demo.sync_iterate_break_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "3ce77bb1", + "metadata": {}, + "source": [ + "## Async: `break` inside `contextlib.aclosing`" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "90a16b98", + "metadata": { + "execution": { + "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": [ + { + "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-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": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\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-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": [ + { + "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", + " post_run_step: respond failed with GeneratorExit()\n", + " caller: some time later\n" + ] + } + ], + "source": [ + "await demo.async_break_without_aclose_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "d33d5c0e", + "metadata": {}, + "source": [ + "## Cancelling the consuming task\n", + "\n", + "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-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": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\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 -> cancel the response\n", + " respond: close LLM connection\n", + " persister: saved for this run: None\n", + " resume: next action is 'respond' (it runs again)\n" + ] + } + ], + "source": [ + "await demo.client_disconnect_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "df4f3273", + "metadata": {}, + "source": [ + "## Persistence: a stopped step is saved as `failed`, and resuming re-runs it" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "5a6ec507", + "metadata": { + "execution": { + "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": [ + { + "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", + " 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" + ] + } + ], + "source": [ + "demo.resume_after_stop_demo()" + ] + }, + { + "cell_type": "markdown", + "id": "94a900f7", + "metadata": {}, + "source": [ + "## Cleanup errors: reported to hooks, and raised by `close()`" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "ef2ae65c", + "metadata": { + "execution": { + "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": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "=== Cleanup errors: reported to hooks, and raised by close()\n", + " respond: open LLM connection\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" + ] + } + ], + "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..fa5e87339 --- /dev/null +++ b/examples/streaming-early-stop/server.py @@ -0,0 +1,91 @@ +# 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. + +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. 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 +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 PostRunStepHookAsync + +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(PostRunStepHookAsync): + async def post_run_step(self, *, action, exception: Optional[BaseException], **future_kwargs): + logger.info("post_run_step: %s exception=%r", action.name, exception) + + +async def chat_application(): + return await ( + ApplicationBuilder() + .with_actions(respond=respond) + .with_transitions() + .with_entrypoint("respond") + .with_state(State()) + .with_hooks(LogStep()) + .abuild() + ) + + +@app.get("/chat") +async def chat() -> StreamingResponse: + burr_app = await 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 000000000..1b88069e3 Binary files /dev/null and b/examples/streaming-early-stop/statemachine.png differ diff --git a/pyproject.toml b/pyproject.toml index 07be284f1..4a3d15843 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,6 +89,7 @@ release = [ ] tests = [ + "anyio", "pytest", "pytest-asyncio", "apache-burr[graphviz]", @@ -330,6 +331,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_application.py b/tests/core/test_application.py index efb7d205b..d47d1170d 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -17,12 +17,15 @@ import asyncio import collections +import contextlib import datetime +import gc import logging import typing import uuid from typing import Any, Awaitable, Callable, Dict, Generator, Literal, Optional, Tuple, Union +import anyio import pytest from burr.core import State @@ -40,6 +43,7 @@ action, default, expr, + streaming_action, ) from burr.core.application import ( PRIOR_STEP, @@ -2770,6 +2774,620 @@ async def test_astream_iterate(exhaust_intermediate_generators: bool): assert len(stream_event_tracker.post_stream_item_calls) == 20 +# --- 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. + + +class _EarlyStopCleanupError(Exception): + pass + + +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.""" + + def __init__(self): + self.events: list = [] + self.post_end_stream_calls = 0 + + def hooks(self, is_async: bool) -> list: + log = self + + def record(exception): + log.events.append(("reported", type(exception) if exception is not None else None)) + + class SyncHook(PostRunStepHook, PostEndStreamHook): + def post_run_step(self, *, exception, **future_kwargs): + record(exception) + + 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) + + def post_end_stream(self, **future_kwargs): + log.post_end_stream_calls += 1 + + return [AsyncHook() if is_async else SyncHook()] + + +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") + + if is_async: + + @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() + + else: + + @streaming_action(reads=[], writes=["x"]) + def act(state: State): + try: + for item in ["a", "b", "c"]: + yield {"item": item}, None + yield {"item": "done"}, state.update(x=1) + finally: + cleanup() + + return act + + +def _early_stop_multi_step_action(log: _EarlyStopLog, is_async: bool) -> Action: + if is_async: + + 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") + + def update(self, result: dict, state: State) -> State: + return state.update(x=result["i"]) + + else: + + class MultiStepAct(StreamingAction): + reads = [] + writes = ["x"] + + def stream_run(self, state: State, **run_kwargs): + try: + for i in range(3): + yield {"i": i} + finally: + log.events.append("cleanup") + + def update(self, result: dict, state: State) -> State: + return state.update(x=result["i"]) + + return MultiStepAct() + + +@action(reads=["x"], writes=["y"]) +def _early_stop_after(state: State) -> State: + return state.update(y=state["x"] + 1) + + +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 _early_stop_persister() -> SQLLitePersister: + persister = SQLLitePersister(db_path=":memory:", table_name="early_stop") + persister.initialize() + return persister + + +STOPPED = [("reported", GeneratorExit)] + + +@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() + + +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 + 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_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] + + +@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] + + +@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: + if consume_first_item: + for _ in container: + break + break + # 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 + + +@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 + + +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() + + 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() + + 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" + + +@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: + return state.update(x=0) + + 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): + 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) + + 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) + + 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" + + +@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 + ) + 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 + assert log.events == ["cleanup", ("reported", _EarlyStopCleanupError)] + assert "Stopping a streaming action early raised" in caplog.text + with pytest.raises(_EarlyStopCleanupError): + container.close() + with pytest.raises(_EarlyStopCleanupError): + container.close() + + +@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 + for _ in range(awaits_before_aclose): + await asyncio.sleep(0) + with pytest.raises(_EarlyStopCleanupError): + await container.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") + + +@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 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(),) + ) + _, 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(): action_tracker = CallCaptureTracker() stream_event_tracker = StreamEventCaptureTrackerAsync() 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..6f8047392 --- /dev/null +++ b/tests/core/test_streaming_early_stop_example.py @@ -0,0 +1,115 @@ +# 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" + +CLOSED = "respond: close LLM connection" +STOPPED = "post_run_step: respond failed with GeneratorExit()" + + +@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 _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) + _cleanup_then_report_then(events, "caller: after the loop") + assert events[-1] == "caller: get() -> result=None, 'response' in state=False" + + +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") + + +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("caller: after the loop -- nothing has been cleaned up yet") < ( + events.index(CLOSED) + ) + + +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() + 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): + 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() + 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')", + ] diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 662ecb514..d0b32bd53 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/tests/integrations/test_burr_pydantic.py @@ -824,3 +824,82 @@ 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"] + + +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"]