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

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:

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