fix(core): propagate early close through streaming actions

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)