Skip to content
9 changes: 7 additions & 2 deletions python/packages/core/agent_framework/_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down
12 changes: 12 additions & 0 deletions python/packages/core/agent_framework/observability.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down
Loading