diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 87799d0848..584c1f0110 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2816,6 +2816,7 @@ def __init__( cleanup_hooks if cleanup_hooks is not None else [] ) self._cleanup_run: bool = False + self._stream_error: Exception | None = None self._inner_stream: ResponseStream[Any, Any] | None = None self._inner_stream_source: ResponseStream[Any, Any] | Awaitable[ResponseStream[Any, Any]] | None = None self._wrap_inner: bool = False @@ -2948,8 +2949,12 @@ async def __anext__(self) -> UpdateT: await self._run_cleanup_hooks() await self.get_final_response() raise - except Exception: - await self._run_cleanup_hooks() + except Exception as exc: + self._stream_error = exc + try: + await self._run_cleanup_hooks() + finally: + self._stream_error = None raise if self._map_update is not None: update = self._map_update(update) # type: ignore[assignment] diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 8d2eb05136..6998e5994f 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -1323,6 +1323,12 @@ async def _finalize_stream() -> None: from ._types import ChatResponse try: + if result_stream._stream_error is not None: # pyright: ignore[reportPrivateUsage] + # Stream errored; skip get_final_response() to avoid firing + # result hooks such as after_run context providers on error + # paths. Capture the error on the span before returning. + capture_exception(span=span, exception=result_stream._stream_error, timestamp=time_ns()) # pyright: ignore[reportPrivateUsage] + return response: ChatResponse[Any] = await result_stream.get_final_response() duration = duration_state.get("duration") response_attributes = _get_response_attributes(attributes, response) @@ -1579,6 +1585,12 @@ async def _finalize_stream() -> None: from ._types import AgentResponse try: + if result_stream._stream_error is not None: # pyright: ignore[reportPrivateUsage] + # Stream errored; skip get_final_response() to avoid firing + # result hooks such as after_run context providers on error + # paths. Capture the error on the span before returning. + capture_exception(span=span, exception=result_stream._stream_error, timestamp=time_ns()) # pyright: ignore[reportPrivateUsage] + return response: AgentResponse[Any] = await result_stream.get_final_response() duration = duration_state.get("duration") response_attributes = _get_response_attributes(