diff --git a/burr/core/action.py b/burr/core/action.py index a69db06c6..37cd46a52 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -1466,7 +1466,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. @@ -1557,7 +1557,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], ], ): @@ -1628,7 +1628,7 @@ async def just_results() -> AsyncGeneratorReturnType: yield results, final_state async def empty_callback( - result: Optional[StreamResultType], state: State, exc: Optional[Exception] + result: Optional[StreamResultType], state: State, exc: Optional[BaseException] ): pass diff --git a/burr/core/application.py b/burr/core/application.py index 25bce4a10..28662a761 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -691,16 +691,20 @@ def __init__(self, method: ExecuteMethod): def call_pre(self, app) -> bool: if should_run_hooks := (app.uid not in _run_call_var.get({})): _run_call_var.set({**_run_call_var.get({}), **{app.uid: self.method}}) - app._adapter_set.call_all_lifecycle_hooks_sync( - "pre_run_execute_call", - app_id=app._uid, - partition_key=app._partition_key, - state=app.state, - method=self.method, - ) + try: + app._adapter_set.call_all_lifecycle_hooks_sync( + "pre_run_execute_call", + app_id=app._uid, + partition_key=app._partition_key, + state=app.state, + method=self.method, + ) + except BaseException: + _run_call_var.set({k: v for k, v in _run_call_var.get().items() if k != app.uid}) + raise return should_run_hooks - def call_post(self, app, exc) -> bool: + def call_post(self, app, exc: Optional[BaseException]) -> bool: if should_run_hooks := ( app.uid in _run_call_var.get(dict) and _run_call_var.get()[app.uid] == self.method ): @@ -718,16 +722,20 @@ def call_post(self, app, exc) -> bool: async def acall_pre(self, app) -> bool: if should_run_hooks := (app.uid not in _run_call_var.get({})): _run_call_var.set({**_run_call_var.get({}), **{app.uid: self.method}}) - await app._adapter_set.call_all_lifecycle_hooks_sync_and_async( - "pre_run_execute_call", - app_id=app._uid, - partition_key=app._partition_key, - state=app.state, - method=self.method, - ) + try: + await app._adapter_set.call_all_lifecycle_hooks_sync_and_async( + "pre_run_execute_call", + app_id=app._uid, + partition_key=app._partition_key, + state=app.state, + method=self.method, + ) + except BaseException: + _run_call_var.set({k: v for k, v in _run_call_var.get().items() if k != app.uid}) + raise return should_run_hooks - async def acall_post(self, app, exc) -> bool: + async def acall_post(self, app, exc: Optional[BaseException]) -> bool: if should_run_hooks := ( app.uid in _run_call_var.get(dict) and _run_call_var.get()[app.uid] == self.method ): @@ -746,23 +754,74 @@ def __call__(self, fn: CallableT) -> CallableT: @functools.wraps(fn) async def wrapper_async(app_self, *args, **kwargs): # We only run at the top level, so we decorate it if we're there - await self.acall_pre(app_self) - exc = None + should_run_hooks = False + exc: Optional[BaseException] = None try: + should_run_hooks = await self.acall_pre(app_self) return await fn(app_self, *args, **kwargs) + except BaseException as e: + exc = e + raise finally: - await self.acall_post(app_self, exc) + if should_run_hooks: + await self.acall_post(app_self, exc) @functools.wraps(fn) def wrapper_sync(app_self, *args, **kwargs): - self.call_pre(app_self) - exc = None + should_run_hooks = False + exc: Optional[BaseException] = None try: + should_run_hooks = self.call_pre(app_self) return fn(app_self, *args, **kwargs) + except BaseException as e: + exc = e + raise finally: - self.call_post(app_self, exc) + if should_run_hooks: + self.call_post(app_self, exc) - return wrapper_async if inspect.iscoroutinefunction(fn) else wrapper_sync + @functools.wraps(fn) + def wrapper_generator(app_self, *args, **kwargs): + def generator(): + should_run_hooks = False + exc: Optional[BaseException] = None + try: + should_run_hooks = self.call_pre(app_self) + return (yield from fn(app_self, *args, **kwargs)) + except BaseException as e: + exc = e + raise + finally: + if should_run_hooks: + self.call_post(app_self, exc) + + return generator() + + @functools.wraps(fn) + def wrapper_async_generator(app_self, *args, **kwargs): + async def generator(): + should_run_hooks = False + exc: Optional[BaseException] = None + try: + should_run_hooks = await self.acall_pre(app_self) + async for item in fn(app_self, *args, **kwargs): + yield item + except BaseException as e: + exc = e + raise + finally: + if should_run_hooks: + await self.acall_post(app_self, exc) + + return generator() + + if inspect.iscoroutinefunction(fn): + return wrapper_async + if inspect.isasyncgenfunction(fn): + return wrapper_async_generator + if inspect.isgeneratorfunction(fn): + return wrapper_generator + return wrapper_sync class TracerFactoryContextHook(PreRunStepHook, PostRunStepHook): @@ -1497,43 +1556,47 @@ def streaming_response(state: State, prompt: str) -> Generator[dict, None, Tuple result, state = output.get() print(format(result['response'], color)) """ - self.validate_correct_async_use() call_execute_method_wrapper = _call_execute_method_pre_post(ExecuteMethod.stream_result) call_execute_method_wrapper.call_pre(self) - halt_before, halt_after, inputs = self._process_control_flow_params( - halt_before, halt_after, inputs - ) - self._validate_halt_conditions(halt_before, halt_after) - next_action = self.get_next_action() - if next_action is None: - raise ValueError( - f"Cannot stream result -- no next action found! Prior action was: {self._state[PRIOR_STEP]}" - ) - if next_action.name not in halt_after: - # fast forward until we get to the action - # run already handles incrementing sequence IDs, nothing to worry about here - next_action, results, state = self.run( - halt_before=halt_after + halt_before, inputs=inputs + try: + self.validate_correct_async_use() + halt_before, halt_after, inputs = self._process_control_flow_params( + halt_before, halt_after, inputs ) - # In this case, we are ready to halt and return an empty generator - # The results will be None, and the state will be the final state - # For context, this is specifically for the case in which you want to have - # multiple terminal points with a unified API, where some are streaming, and some are not. - if next_action.name in halt_before and next_action.name not in halt_after: - call_execute_method_wrapper.call_post(self, None) - return next_action, StreamingResultContainer.pass_through( - results=results, final_state=state + self._validate_halt_conditions(halt_before, halt_after) + next_action = self.get_next_action() + if next_action is None: + raise ValueError( + f"Cannot stream result -- no next action found! Prior action was: {self._state[PRIOR_STEP]}" ) - self._increment_sequence_id() - self._adapter_set.call_all_lifecycle_hooks_sync( - "pre_run_step", - action=next_action, - state=self._state, - inputs=inputs, - sequence_id=self.sequence_id, - app_id=self._uid, - partition_key=self._partition_key, - ) + if next_action.name not in halt_after: + # fast forward until we get to the action + # run already handles incrementing sequence IDs, nothing to worry about here + next_action, results, state = self.run( + halt_before=halt_after + halt_before, inputs=inputs + ) + # In this case, we are ready to halt and return an empty generator + # The results will be None, and the state will be the final state + # For context, this is specifically for the case in which you want to have + # multiple terminal points with a unified API, where some are streaming, and some are not. + if next_action.name in halt_before and next_action.name not in halt_after: + call_execute_method_wrapper.call_post(self, None) + return next_action, StreamingResultContainer.pass_through( + results=results, final_state=state + ) + self._increment_sequence_id() + self._adapter_set.call_all_lifecycle_hooks_sync( + "pre_run_step", + action=next_action, + state=self._state, + inputs=inputs, + sequence_id=self.sequence_id, + app_id=self._uid, + partition_key=self._partition_key, + ) + except BaseException as e: + call_execute_method_wrapper.call_post(self, e) + raise # we need to track if there's any exceptions that occur during this try: @@ -1548,7 +1611,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", @@ -1622,7 +1685,7 @@ def callback( partition_key=self._partition_key, lifecycle_adapters=self._adapter_set, ) - except Exception as e: + except BaseException as e: # We only want to raise this in the case of an exception # otherwise, this will get delegated to the finally # block of the streaming result container @@ -1751,42 +1814,46 @@ async def streaming_response(state: State, prompt: str) -> Generator[dict, None, result, state = await output.get() print(format(result['response'], color)) """ - call_execute_method_wrapper = _call_execute_method_pre_post(ExecuteMethod.stream_result) + call_execute_method_wrapper = _call_execute_method_pre_post(ExecuteMethod.astream_result) await call_execute_method_wrapper.acall_pre(self) - halt_before, halt_after, inputs = self._process_control_flow_params( - halt_before, halt_after, inputs - ) - self._validate_halt_conditions(halt_before, halt_after) - next_action = self.get_next_action() - if next_action is None: - raise ValueError( - f"Cannot stream result -- no next action found! Prior action was: {self._state[PRIOR_STEP]}" - ) - if next_action.name not in halt_after: - # fast forward until we get to the action - # run already handles incrementing sequence IDs, nothing to worry about here - next_action, results, state = await self.arun( - halt_before=halt_after + halt_before, inputs=inputs + try: + halt_before, halt_after, inputs = self._process_control_flow_params( + halt_before, halt_after, inputs ) - # In this case, we are ready to halt and return an empty generator - # The results will be None, and the state will be the final state - # For context, this is specifically for the case in which you want to have - # multiple terminal points with a unified API, where some are streaming, and some are not. - if next_action.name in halt_before and next_action.name not in halt_after: - await call_execute_method_wrapper.acall_post(self, None) - return next_action, AsyncStreamingResultContainer.pass_through( - results=results, final_state=state + self._validate_halt_conditions(halt_before, halt_after) + next_action = self.get_next_action() + if next_action is None: + raise ValueError( + f"Cannot stream result -- no next action found! Prior action was: {self._state[PRIOR_STEP]}" ) - self._increment_sequence_id() - await self._adapter_set.call_all_lifecycle_hooks_sync_and_async( - "pre_run_step", - action=next_action, - state=self._state, - inputs=inputs, - sequence_id=self.sequence_id, - app_id=self._uid, - partition_key=self._partition_key, - ) + if next_action.name not in halt_after: + # fast forward until we get to the action + # run already handles incrementing sequence IDs, nothing to worry about here + next_action, results, state = await self.arun( + halt_before=halt_after + halt_before, inputs=inputs + ) + # In this case, we are ready to halt and return an empty generator + # The results will be None, and the state will be the final state + # For context, this is specifically for the case in which you want to have + # multiple terminal points with a unified API, where some are streaming, and some are not. + if next_action.name in halt_before and next_action.name not in halt_after: + await call_execute_method_wrapper.acall_post(self, None) + return next_action, AsyncStreamingResultContainer.pass_through( + results=results, final_state=state + ) + self._increment_sequence_id() + await self._adapter_set.call_all_lifecycle_hooks_sync_and_async( + "pre_run_step", + action=next_action, + state=self._state, + inputs=inputs, + sequence_id=self.sequence_id, + app_id=self._uid, + partition_key=self._partition_key, + ) + except BaseException as e: + await call_execute_method_wrapper.acall_post(self, e) + raise try: def process_result( @@ -1799,7 +1866,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", @@ -1887,7 +1954,7 @@ async def callback( partition_key=self._partition_key, lifecycle_adapters=self._adapter_set, ) - except Exception as e: + except BaseException as e: # We only want to raise this in the case of an exception # otherwise, this will get delegated to the finally # block of the streaming result container diff --git a/burr/integrations/langfuse.py b/burr/integrations/langfuse.py index 20e506f68..d6fbe2efa 100644 --- a/burr/integrations/langfuse.py +++ b/burr/integrations/langfuse.py @@ -312,7 +312,7 @@ def post_run_execute_call( self, *, state: "State", - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs: Any, ): try: diff --git a/burr/integrations/opentelemetry.py b/burr/integrations/opentelemetry.py index 7a6970e88..8bed994d1 100644 --- a/burr/integrations/opentelemetry.py +++ b/burr/integrations/opentelemetry.py @@ -108,7 +108,7 @@ def convert_to_otel_attribute(attr: Any): return str(attr) -def _exit_span(exc: Optional[Exception] = None): +def _exit_span(exc: Optional[BaseException] = None): """Ditto with _enter_span, but for exiting the span. Pops the token off the stack and detaches the context.""" stack = token_stack.get()[:] token, span = stack.pop() @@ -254,7 +254,7 @@ def post_run_step( def post_run_execute_call( self, *, - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs, ): _exit_span(exception) diff --git a/burr/lifecycle/base.py b/burr/lifecycle/base.py index 66d8bd7e6..a384902b3 100644 --- a/burr/lifecycle/base.py +++ b/burr/lifecycle/base.py @@ -354,7 +354,7 @@ def post_run_execute_call( partition_key: str, state: "State", method: ExecuteMethod, - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs, ): pass @@ -372,7 +372,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/tests/core/test_application.py b/tests/core/test_application.py index 9313cefc9..cf9e45919 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -310,7 +310,7 @@ def post_run_execute_call( partition_key: str, state: "State", method: ExecuteMethod, - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs, ): self.post_run_execute_calls.append(("post", locals())) @@ -341,7 +341,7 @@ async def post_run_execute_call( partition_key: str, state: "State", method: ExecuteMethod, - exception: Optional[Exception], + exception: Optional[BaseException], **future_kwargs, ): self.post_run_execute_calls.append(("post", locals())) @@ -2654,6 +2654,10 @@ def test_stream_iterate(exhaust_intermediate_generators: bool): assert len(stream_event_tracker.post_end_stream_calls) == 2 assert len(stream_event_tracker.post_stream_item_calls) == 20 assert len(stream_event_tracker.post_stream_item_calls) == 20 + assert len(action_tracker.pre_run_execute_calls) == 1 + assert len(action_tracker.post_run_execute_calls) == 1 + assert action_tracker.post_run_execute_calls[0][1]["method"] == ExecuteMethod.stream_iterate + assert action_tracker.post_run_execute_calls[0][1]["exception"] is None @pytest.mark.asyncio @@ -2712,6 +2716,10 @@ async def test_astream_iterate(exhaust_intermediate_generators: bool): assert len(stream_event_tracker.pre_start_stream_calls) == 2 assert len(stream_event_tracker.post_end_stream_calls) == 2 assert len(stream_event_tracker.post_stream_item_calls) == 20 + assert len(action_tracker.pre_run_execute_calls) == 1 + assert len(action_tracker.post_run_execute_calls) == 1 + assert action_tracker.post_run_execute_calls[0][1]["method"] == ExecuteMethod.astream_iterate + assert action_tracker.post_run_execute_calls[0][1]["exception"] is None async def test_astream_result_halt_after_run_through_streaming(): @@ -3345,6 +3353,414 @@ async def test_application_run_step_runs_hooks(): assert len(hooks[0].pre_run_execute_calls) == 1 +@pytest.mark.parametrize( + "method, execute_method", + [("step", ExecuteMethod.step), ("run", ExecuteMethod.run)], +) +def test_sync_execute_hook_reports_action_exception(method, execute_method): + tracker = CallCaptureTracker() + broken_action = base_broken_action.with_name("broken") + app = Application( + state=State({}), + entrypoint="broken", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[broken_action], transitions=[]), + ) + + with pytest.raises(BrokenStepException): + if method == "step": + app.step() + else: + app.run(halt_after=["broken"]) + + assert len(tracker.pre_run_execute_calls) == 1 + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == execute_method + assert isinstance(post_kwargs["exception"], BrokenStepException) + + +@pytest.mark.parametrize( + "method, execute_method", + [("astep", ExecuteMethod.astep), ("arun", ExecuteMethod.arun)], +) +async def test_async_execute_hook_reports_action_exception(method, execute_method): + sync_tracker = CallCaptureTracker() + async_tracker = ExecuteMethodTrackerAsync() + broken_action = base_broken_action_async.with_name("broken") + app = Application( + state=State({}), + entrypoint="broken", + adapter_set=internal.LifecycleAdapterSet(sync_tracker, async_tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[broken_action], transitions=[]), + ) + + with pytest.raises(BrokenStepException): + if method == "astep": + await app.astep() + else: + await app.arun(halt_after=["broken"]) + + assert len(sync_tracker.pre_run_execute_calls) == 1 + assert len(sync_tracker.post_run_execute_calls) == 1 + assert len(async_tracker.pre_run_execute_calls) == 1 + assert len(async_tracker.post_run_execute_calls) == 1 + sync_post_kwargs = sync_tracker.post_run_execute_calls[0][1] + async_post_kwargs = async_tracker.post_run_execute_calls[0][1] + assert sync_post_kwargs["method"] == execute_method + assert async_post_kwargs["method"] == execute_method + assert isinstance(sync_post_kwargs["exception"], BrokenStepException) + assert isinstance(async_post_kwargs["exception"], BrokenStepException) + + +async def test_async_execute_hook_reports_cancellation(): + tracker = ExecuteMethodTrackerAsync() + + @action(reads=[], writes=[]) + async def cancelled_action(state: State) -> State: + raise asyncio.CancelledError() + + app = ( + ApplicationBuilder() + .with_actions(cancelled=cancelled_action) + .with_transitions() + .with_entrypoint("cancelled") + .with_state() + .with_hooks(tracker) + .build() + ) + + with pytest.raises(asyncio.CancelledError): + await app.astep() + + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.astep + assert isinstance(post_kwargs["exception"], asyncio.CancelledError) + + +def test_execute_hook_pre_failure_does_not_leak_sentinel(): + class RaisesOnce(PreApplicationExecuteCallHook, PostApplicationExecuteCallHook): + def __init__(self): + self.pre_calls = 0 + self.post_calls = 0 + + def pre_run_execute_call(self, **future_kwargs): + self.pre_calls += 1 + if self.pre_calls == 1: + raise RuntimeError("pre hook failure") + + def post_run_execute_call(self, **future_kwargs): + self.post_calls += 1 + + tracker = RaisesOnce() + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + with pytest.raises(RuntimeError, match="pre hook failure"): + app.step() + app.step() + + assert tracker.pre_calls == 2 + assert tracker.post_calls == 1 + + +def test_stream_result_setup_failure_reports_exception_and_clears_sentinel(): + tracker = CallCaptureTracker() + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + with pytest.raises(ValueError, match="not registered actions"): + app.stream_result(halt_after=["missing"]) + app.step() + + assert len(tracker.pre_run_execute_calls) == 2 + assert len(tracker.post_run_execute_calls) == 2 + stream_post_kwargs = tracker.post_run_execute_calls[0][1] + step_post_kwargs = tracker.post_run_execute_calls[1][1] + assert stream_post_kwargs["method"] == ExecuteMethod.stream_result + assert isinstance(stream_post_kwargs["exception"], ValueError) + assert step_post_kwargs["method"] == ExecuteMethod.step + assert step_post_kwargs["exception"] is None + + +async def test_astream_result_setup_failure_reports_exception_and_clears_sentinel(): + tracker = ExecuteMethodTrackerAsync() + counter_action = base_counter_action_async.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + with pytest.raises(ValueError, match="not registered actions"): + await app.astream_result(halt_after=["missing"]) + await app.astep() + + assert len(tracker.pre_run_execute_calls) == 2 + assert len(tracker.post_run_execute_calls) == 2 + stream_post_kwargs = tracker.post_run_execute_calls[0][1] + step_post_kwargs = tracker.post_run_execute_calls[1][1] + assert stream_post_kwargs["method"] == ExecuteMethod.astream_result + assert isinstance(stream_post_kwargs["exception"], ValueError) + assert step_post_kwargs["method"] == ExecuteMethod.astep + assert step_post_kwargs["exception"] is None + + +def test_stream_result_reports_stream_exception_to_execute_hook(): + tracker = CallCaptureTracker() + broken_action = MultiStepStreamingCounterWithExceptionNoResult().with_name("broken") + app = Application( + state=State({"count": 0, "tracker": []}), + entrypoint="broken", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[broken_action], transitions=[]), + ) + + _, stream = app.stream_result(halt_after=["broken"]) + with pytest.raises(RuntimeError, match="simulated failure with no result"): + list(stream) + + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.stream_result + assert isinstance(post_kwargs["exception"], RuntimeError) + + +async def test_astream_result_reports_stream_exception_to_execute_hook(): + tracker = ExecuteMethodTrackerAsync() + broken_action = MultiStepStreamingCounterWithExceptionNoResultAsync().with_name("broken") + app = Application( + state=State({"count": 0, "tracker": []}), + entrypoint="broken", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[broken_action], transitions=[]), + ) + + _, stream = await app.astream_result(halt_after=["broken"]) + with pytest.raises(RuntimeError, match="simulated failure with no result"): + async for _ in stream: + pass + + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.astream_result + assert isinstance(post_kwargs["exception"], RuntimeError) + + +async def test_astream_result_reports_cancellation_to_execute_hook(): + class CancelledStreamingAction(AsyncStreamingAction): + async def stream_run(self, state: State, **run_kwargs) -> AsyncGenerator[dict, None]: + raise asyncio.CancelledError() + yield # Make this an async generator. + + @property + def reads(self) -> list[str]: + return [] + + @property + def writes(self) -> list[str]: + return [] + + def update(self, result: dict, state: State) -> State: + return state + + tracker = ExecuteMethodTrackerAsync() + app = Application( + state=State({}), + entrypoint="cancelled", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[CancelledStreamingAction().with_name("cancelled")], transitions=[]), + ) + + _, stream = await app.astream_result(halt_after=["cancelled"]) + with pytest.raises(asyncio.CancelledError): + async for _ in stream: + pass + + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.astream_result + assert isinstance(post_kwargs["exception"], asyncio.CancelledError) + + +def test_iterate_execute_hook_reports_action_exception_after_iteration_starts(): + tracker = CallCaptureTracker() + broken_action = base_broken_action.with_name("broken") + app = Application( + state=State({}), + entrypoint="broken", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[broken_action], transitions=[]), + ) + + iterator = app.iterate(halt_after=["broken"]) + assert not tracker.pre_run_execute_calls + assert not tracker.post_run_execute_calls + with pytest.raises(BrokenStepException): + next(iterator) + + assert len(tracker.pre_run_execute_calls) == 1 + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.iterate + assert isinstance(post_kwargs["exception"], BrokenStepException) + + +def test_iterate_execute_hook_finishes_after_iteration_ends(): + tracker = CallCaptureTracker() + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + iterator = app.iterate(halt_after=["counter"]) + assert not tracker.pre_run_execute_calls + assert not tracker.post_run_execute_calls + next(iterator) + assert len(tracker.pre_run_execute_calls) == 1 + assert not tracker.post_run_execute_calls + with pytest.raises(StopIteration): + next(iterator) + + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.iterate + assert post_kwargs["exception"] is None + + +async def test_aiterate_execute_hook_reports_action_exception_after_iteration_starts(): + tracker = ExecuteMethodTrackerAsync() + broken_action = base_broken_action_async.with_name("broken") + app = Application( + state=State({}), + entrypoint="broken", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[broken_action], transitions=[]), + ) + + iterator = app.aiterate(halt_after=["broken"]) + assert not tracker.pre_run_execute_calls + assert not tracker.post_run_execute_calls + with pytest.raises(BrokenStepException): + await anext(iterator) + + assert len(tracker.pre_run_execute_calls) == 1 + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.aiterate + assert isinstance(post_kwargs["exception"], BrokenStepException) + + +async def test_aiterate_execute_hook_finishes_after_iteration_ends(): + tracker = ExecuteMethodTrackerAsync() + counter_action = base_counter_action_async.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + iterator = app.aiterate(halt_after=["counter"]) + assert not tracker.pre_run_execute_calls + assert not tracker.post_run_execute_calls + await anext(iterator) + assert len(tracker.pre_run_execute_calls) == 1 + assert not tracker.post_run_execute_calls + with pytest.raises(StopAsyncIteration): + await anext(iterator) + + assert len(tracker.post_run_execute_calls) == 1 + post_kwargs = tracker.post_run_execute_calls[0][1] + assert post_kwargs["method"] == ExecuteMethod.aiterate + assert post_kwargs["exception"] is None + + +async def test_async_execute_hook_pre_failure_does_not_leak_sentinel(): + class RaisesOnce(PreApplicationExecuteCallHookAsync, PostApplicationExecuteCallHookAsync): + def __init__(self): + self.pre_calls = 0 + self.post_calls = 0 + + async def pre_run_execute_call(self, **future_kwargs): + self.pre_calls += 1 + if self.pre_calls == 1: + raise RuntimeError("pre hook failure") + + async def post_run_execute_call(self, **future_kwargs): + self.post_calls += 1 + + tracker = RaisesOnce() + counter_action = base_counter_action_async.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + adapter_set=internal.LifecycleAdapterSet(tracker), + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + with pytest.raises(RuntimeError, match="pre hook failure"): + await app.astep() + await app.astep() + + assert tracker.pre_calls == 2 + assert tracker.post_calls == 1 + + def test_application_post_application_create_hook(): class PostApplicationCreateTracker(PostApplicationCreateHook): def __init__(self):