Skip to content

fix: run post execute hooks when stream_result fails before streaming #968

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
elijahbenizzy merged 1 commit into from
Oct 17, 2026
Merged

fix: run post execute hooks when stream_result fails before streaming #968

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 72 additions & 64 deletions burr/core/application.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1503,40 +1503,44 @@ def streaming_response(state: State, prompt: str) -> Generator[dict, None, Tuple
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:
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 Exception 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:
Expand Down Expand Up @@ -1756,40 +1760,44 @@ async def streaming_response(state: State, prompt: str) -> Generator[dict, None,
"""
call_execute_method_wrapper = _call_execute_method_pre_post(ExecuteMethod.stream_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 Exception as e:
await call_execute_method_wrapper.acall_post(self, e)
raise
try:

def process_result(
Expand Down
81 changes: 81 additions & 0 deletions tests/integrations/test_burr_langfuse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.sdk.util.instrumentation import InstrumentationScope
from opentelemetry.trace import get_current_span

from burr.core import ApplicationBuilder, State, action
from burr.integrations.langfuse import BURR_TRACER_NAME, LangfuseBridge, burr_span_export_filter
Expand Down Expand Up @@ -276,3 +277,83 @@ def test_burr_span_export_filter():
assert burr_span_export_filter(
_make_readable_span("some.other.scope", {"gen_ai.system": "openai"})
)


@action(reads=[], writes=["x"])
def maybe_failing_action(state: State, fail: bool = False) -> Tuple[dict, State]:
if fail:
raise RuntimeError("boom")
result = {"x": 1}
return result, state.update(**result)


@pytest.fixture
def propagation_attributes(monkeypatch):
attributes = ContextVar("propagation_attributes", default=None)

@contextmanager
def capture_propagation_attributes(**kwargs):
token = attributes.set(kwargs)
try:
yield
finally:
attributes.reset(token)

monkeypatch.setattr(
"burr.integrations.langfuse.propagate_attributes", capture_propagation_attributes
)
return attributes


def _build_stream_app(bridge: LangfuseBridge):
return (
ApplicationBuilder()
.with_actions(a=maybe_failing_action, b=finish_action)
.with_transitions(("a", "b"))
.with_entrypoint("a")
.with_identifiers(app_id="test-app-id", partition_key="test-user")
.with_hooks(bridge)
.build()
)


def _assert_no_context_left_and_next_run_traced(app, exporter, propagation_attributes):
assert not get_current_span().is_recording()
assert propagation_attributes.get() is None
exporter.clear()
app.run(halt_after=["b"])
spans_by_name = {span.name: span for span in exporter.get_finished_spans()}
assert set(spans_by_name) == {"run", "a", "b"}
assert spans_by_name["run"].parent is None


@pytest.mark.parametrize(
"halt_after,inputs,error",
[(["b"], {"fail": True}, RuntimeError), (["nope"], {}, ValueError)],
ids=["failing_fast_forward", "unknown_halt_after"],
)
def test_langfuse_bridge_stream_result_failure_closes_trace(
span_capture, propagation_attributes, halt_after, inputs, error
):
exporter, tracer = span_capture
app = _build_stream_app(LangfuseBridge(langfuse_client=Mock(), tracer=tracer))
with pytest.raises(error):
app.stream_result(halt_after=halt_after, inputs=inputs)

root = {span.name: span for span in exporter.get_finished_spans()}["stream_result"]
assert not root.status.is_ok
_assert_no_context_left_and_next_run_traced(app, exporter, propagation_attributes)


@pytest.mark.asyncio
async def test_langfuse_bridge_astream_result_failure_closes_trace(
span_capture, propagation_attributes
):
exporter, tracer = span_capture
app = _build_stream_app(LangfuseBridge(langfuse_client=Mock(), tracer=tracer))
with pytest.raises(RuntimeError):
await app.astream_result(halt_after=["b"], inputs={"fail": True})

root = {span.name: span for span in exporter.get_finished_spans()}["stream_result"]
assert not root.status.is_ok
_assert_no_context_left_and_next_run_traced(app, exporter, propagation_attributes)
Loading