Repository navigation
Conversation
When a ResponseStream failed or was cancelled, close() shut nested ResponseStreams without the error, so their cleanup hooks treated the run as successful. With agent middleware this made the agent telemetry stream call get_final_response(), which ran every ContextProvider.after_run and saved the partial turn to history. close() now takes the terminal error and records it on itself and on nested streams while their cleanup hooks run.
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
🟡 Changes recommended
Buffered streams still omit the error during nested closure, allowing failed partial responses to be persisted.
1 open issue
What changed in this PR
Propagates streaming errors through nested Python ResponseStream wrappers to prevent failed agent runs from saving partial history.
Changes:
- Adds optional exception propagation to stream closure.
- Adds regression coverage for middleware-wrapped transform failures and cancellation.
| File | Description |
|---|---|
| python/packages/core/tests/core/test_types.py | Tests that failed streams skip after_run and history persistence. |
| python/packages/core/agent_framework/_types.py | Passes terminal errors into nested stream cleanup. |
Review effort: Balanced
| if iterator is not None: | ||
| if isinstance(iterator, ResponseStream): | ||
| await cast(ResponseStream[UpdateT, Any], iterator).close() | ||
| await cast(ResponseStream[UpdateT, Any], iterator).close(exc) |
There was a problem hiding this comment.
Good catch, this was right: with .buffer_updates() a transform failure or cancellation reached _abort_buffered_materialization, which called close() without the error, so the nested stream still finalized the partial response and ran after_run. The new commit passes the recorded error to close(exc) there, the same way the streaming path does. test_outer_stream_error_skips_after_run_with_agent_middleware is now parametrized over buffered and unbuffered streams; both buffered cases failed before the change and all four pass now.
A transform failure or cancellation in a buffered outer stream went through _abort_buffered_materialization, which closed the stream without the error, so the nested stream still ran its cleanup as a normal completion. It now passes the recorded error to close(), and the regression test gains a buffered variant.
@microsoft-github-policy-service agree |
Somay kousis (Somay-kousis)
left a comment
Somay kousis (Somay-kousis)
left a comment
There was a problem hiding this comment.
Checked the fix, and it does what it says for errors raised inside __anext__ (transform hook, buffered or not).
What I ran:
- Your 4 test cases against main's
_types.py: all 4 fail. On your head: all 4 pass.
One gap, worth considering, not blocking. The description says a failed or cancelled streaming run no longer calls after_run, with or without middleware. That still happens when the error or the cancel lands in the consumer's own code inside async with. __aexit__ calls close() with no error, so the telemetry stream's cleanup finalizes the partial response:
async with agent.run("hi", stream=True, session=session) as s:
async for _ in s:
await asyncio.sleep(10) # task gets cancelled hereProbe (fake streaming client, InMemoryHistoryProvider plus a provider that counts after_run), same output on main and on your head:
async-with, body raises | mw=False: raised=RuntimeError after_run=1 saved_msgs=2 | mw=True: raised=RuntimeError after_run=1 saved_msgs=2
async-with, body cancelled | mw=False: raised=CancelledError after_run=1 saved_msgs=2 | mw=True: raised=CancelledError after_run=1 saved_msgs=2
It's not a regression and it's the same with and without middleware, so it may be out of scope here. But with your new close(exc) it's one line in __aexit__:
await self.close(exc if exc is not None and not isinstance(exc, GeneratorExit) else None)With that on your branch, both rows above go to after_run=0 saved_msgs=0, a plain break and an explicit close() still finalize as before. Whether a consumer-side exception should drop the turn is a product call, so I'd either add it here with a test or narrow the description to errors raised inside the stream.
|
Thanks for running the cases. You're right: an exception or cancellation raised in the consumer's own code inside |
Motivation & Context
closes #9474
When an
Agenthas agent middleware and a streaming run fails in the outer stream (a transform hook raising, or the consumer being cancelled), everyContextProvider.after_runstill runs with the partial response andInMemoryHistoryProvidersaves it as a completed turn. The error handler inResponseStream.__anext__recorded the error only on the stream it ran in and then calledclose(), which closed the nestedResponseStreamwithout the error, so the agent telemetry stream's cleanup hook saw no error and calledget_final_response().Description & Review Guide
ResponseStream.close()takes an optionalexc. It passes it to a nestedResponseStreamwhen closing it, and when it is set the stream runs its cleanup hooks through_handle_stream_error(exc), so_stream_erroris set while they run. The__anext__error handler now callsclose(exc). A plainclose()behaves as before.after_runor saves the partial response, with or without agent middleware. An exception raised in the consumer's own code insideasync withis unchanged and still finalizes the response on block exit. This matches the path without middleware, where the outer stream is the telemetry stream and already sees the error._handle_stream_errorinclose()is the right place to record the error for nested streams.The new test
test_outer_stream_error_skips_after_run_with_agent_middlewareintest_types.pybuilds anAgentwith a pass-through agent middleware, anInMemoryHistoryProviderand a provider that countsafter_runcalls, and iterates a streaming run with a transform hook that raisesRuntimeErrororasyncio.CancelledError. It checks that the error propagates, thatafter_runis not called and that the session history is empty. Both cases fail without the change and pass with it.Related Issue
closes #9474
AI Assistance
AI assistance details:
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix); a workflow keeps the label and title prefix in sync automatically.