fix: report execute hook failures

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
 

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

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:

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)