diff --git a/burr/core/application.py b/burr/core/application.py index 415e0984e..5d54636b2 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -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: @@ -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( diff --git a/tests/integrations/test_burr_langfuse.py b/tests/integrations/test_burr_langfuse.py index 418a1a99c..4be979f56 100644 --- a/tests/integrations/test_burr_langfuse.py +++ b/tests/integrations/test_burr_langfuse.py @@ -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 @@ -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)