Python: [Feature]: `ResponseStream.__anext__` runs every per-update transform outside its error handler, so a raising hook or mapper skips cleanup entirely

[antsok]

Description

ResponseStream.__anext__ catches exceptions only around the pull from the underlying iterator. Every per-update transform runs outside that handler: _flat_map_update, _map_update, and _record_update() — the last being what runs the hooks registered via with_transform_hook / ChatContext.stream_transform_hooks — all sit after the except clauses, plus one _record_update() on the _pending_mapped_updates branch before the try. Any of them raising propagates straight to the consumer without _stream_error being set and without _run_cleanup_hooks() running. Only a failure raised while pulling the underlying iterator takes the full error path, which is a narrower set than "a mid-stream failure".

The code

agent_framework/_types.py, ResponseStream.__anext__:

async def __anext__(self) -> UpdateT:
    while True:
        if self._pending_mapped_updates:
            return await self._record_update(self._pending_mapped_updates.pop(0))   # outside the try

        try:
            ...
                update: UpdateT = await self._iterator.__anext__()
        except StopAsyncIteration:
            self._consumed = True
            await self._run_cleanup_hooks()
            await self.get_final_response()
            raise
        except Exception as exc:
            self._stream_error = exc
            try:
                await self._run_cleanup_hooks()
            finally:
                self._stream_error = None
            raise
        if self._flat_map_update is not None:                                       # outside the try
            mapped_updates = self._flat_map_update(update)
            ...
            continue
        if self._map_update is not None:                                            # outside the try
            update = self._map_update(update)
            ...
        return await self._record_update(update)                                    # outside the try

Every per-update transform is outside the handler — the two mappers and both _record_update call sites, _record_update being the only one that invokes transform hooks. What is inside is the pull, and nothing else.

What it costs

The instrumentation in observability.py registers _record_duration and _finalize_stream as cleanup hooks and relies on them for everything except the span's existence:

• _record_duration never runs, so the elapsed time _finalize_stream would emit is never even measured.
• _finalize_stream never runs, so no response attributes, no duration or token-usage histogram sample, and no capture_exception for the failure — even though _finalize_stream has a branch written for exactly this case (if result_stream._stream_error is not None: capture_exception(...)), which cannot fire because the transform-hook path never sets _stream_error.
• The span itself still ends, via weakref.finalize(wrapped_stream, _close_span). So the trace shows a chat span that closed with no outcome recorded on it at all, rather than a span marked with the error.

Any consumer-registered cleanup hook is silently skipped the same way, which is the more general problem: a cleanup hook is the natural place to put "release what this stream reserved", and it has no reason to expect a mid-stream termination that does not run it.

Measured on a real OpenAIChatCompletionClient over a mocked transport, with a body still streaming when the failure happens:

How iteration ends	Cleanup hooks ran?
A transform hook raises	No
A chunk the provider SDK cannot parse	Yes

The workaround is a trap, which is part of why this is worth fixing

The obvious remedy from a consumer's side — call _run_cleanup_hooks() after catching the hook's exception — is worse than doing nothing. _finalize_stream skips get_final_response() only when _stream_error is set, and the transform-hook path never sets it, so calling cleanup from outside runs the finalizer, and get_final_response() on an unconsumed stream does async for _ in self: pass. For an abort whose entire purpose is to stop reading a runaway response, that reads the runaway. In our case that is the 26 MB of streamed whitespace the abort exists to avoid.

Suggested direction

Bring the whole post-pull section inside the same handler as the iterator pull, so any transform that raises takes the identical path to a failure from the iterator: _stream_error set, cleanup hooks run, exception re-raised. That is all three transforms, not just the hooks — _flat_map_update and _map_update are in exactly the same position — and the _pending_mapped_updates branch at the top of the loop, which is the same _record_update() call reached by a different path.

A second thing worth deciding at the same time, because it is the other half of what this loop does on failure: __anext__'s error branch does not close self._iterator, on either the StopAsyncIteration path or the general one. That is invisible today for a provider whose own generator raised — its finally has already run — but it means a failure that originates here leaves the underlying generator suspended, which is what keeps the provider response open on agent_framework_openai's Chat Completions client. Fixing the handler's scope without also releasing the iterator changes the telemetry and nothing else; measured, the response is still open at the abort.

What this does not fix

Bringing the transforms inside the handler does not, by itself, release the provider HTTP response — the section above is the reason. Measured as a grid on OpenAIChatCompletionClient: the response is closed at the abort only when the core releases its iterator and the provider's _stream() wraps the SDK stream in async with; either alone leaves it to the garbage collector. The provider half is a separate defect, drafted separately. Scoped to this report alone, the gain is that an aborted call stops being invisible to telemetry.

Code Sample

Language/SDK

Python

[karthik-0306]

Thanks for the detailed writeup, Anton Sokolovskyi (@antsok) .

Eduard van Valkenburg (@eavanvalkenburg) I'd like to take this one.

Investigated on current main — confirmed both escapes described: the _pending_mapped_updates branch runs _record_update (and any transform hooks) before the try even starts, and _flat_map_update/_map_update/_record_update all execute after the try/except closes. So a raising hook or mapper propagates without _stream_error ever being set and without _run_cleanup_hooks() running.

Proposed fix, per the suggested direction: move the pending-mapped-updates return and all three per-update transforms inside the same try block as the iterator pull, so any exception from a hook or mapper takes the identical path as an iterator-pull failure — _stream_error set, cleanup hooks run, exception re-raised.

A few things worth noting from checking this against the actual code:

• _run_cleanup_hooks() guards itself against double-invocation via _cleanup_run, so there's no double-cleanup risk between the StopAsyncIteration and Exception branches.
• .map()/.flat_map() derived streams each own independent _stream_error/cleanup state, so this fix is local to whichever stream's callback actually raised — an outer mapper's failure doesn't touch the inner stream's cleanup.
• _stream_error is cleared in a finally block after _run_cleanup_hooks() returns, so it is set (and visible to cleanup hooks) during cleanup but reads as None after the async for exits. Calling it out since it affects how you'd write assertions against this; the fix doesn't change this behavior.
• Edge case worth flagging: a hook or mapper that raises StopAsyncIteration (not intended, but currently unconstrained) would still end iteration either way — but before this fix, it did so without running cleanup or setting _consumed; after, it goes through the same completion path as a normal end-of-stream, cleanup included. Not adding an explicit guard against hooks raising this in this PR.

Scoping to the transform/cleanup-hook gap only — leaving the provider iterator-close issue alone as a separate concern.

Will open a PR with this restructure plus regression tests covering a transform hook raising and a mapper raising mid-stream.