Python: Failed or cancelled streaming run with agent middleware still calls after_run #9474

Observed Behavior

When an Agent has any agent middleware and a streaming run fails in the outer stream (a transform hook raising, or the consumer being cancelled), every ContextProvider.after_run still fires with the partial response, and InMemoryHistoryProvider saves the partial output as a completed turn. Without middleware the same failure does not call after_run and nothing is saved.

Expected Behavior

A streaming run that ends with an exception or a cancellation should not run after_run providers or persist the partial response, with or without middleware, the same as the no-middleware path does today.

Steps to Reproduce

1. Create an Agent over a chat client that streams a few updates, with InMemoryHistoryProvider and a provider that records after_run calls.
2. Add a pass-through @agent_middleware.
3. Call agent.run("hi", stream=True, session=session).with_transform_hook(hook) where hook raises RuntimeError (or asyncio.CancelledError) on the first update, and iterate it.
4. The exception propagates, but after_run was called with "part0 " and the session history holds ["hi", "part0 "].

Minimal Reproduction

"""An error in an outer ResponseStream closes the inner stream as if it succeeded.

An Agent with agent middleware wraps the telemetry stream in an outer
ResponseStream. When a transform hook on that outer stream raises (or the
consumer is cancelled), the inner stream is closed with no error recorded, its
cleanup hook calls get_final_response(), and every ContextProvider.after_run
fires with the partial response. The history provider saves it as a completed
turn. Without middleware the same failure does not call after_run.
"""

import asyncio
from collections.abc import Sequence

from agent_framework import (
    Agent,
    BaseChatClient,
    ChatResponse,
    ChatResponseUpdate,
    Content,
    ContextProvider,
    FunctionInvocationLayer,
    InMemoryHistoryProvider,
    ResponseStream,
    agent_middleware,
)


class FakeStreamingClient(FunctionInvocationLayer, BaseChatClient):
    def _inner_get_response(self, *, messages, stream, options, **kwargs):
        async def _stream():
            for i in range(3):
                yield ChatResponseUpdate(contents=[Content.from_text(f"part{i} ")], role="assistant")
                await asyncio.sleep(0)

        def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
            return ChatResponse.from_updates(updates)

        return ResponseStream(_stream(), finalizer=_finalize)


class RecordingProvider(ContextProvider):
    def __init__(self) -> None:
        super().__init__("recorder")
        self.after_run_texts: list[str] = []

    async def after_run(self, *, agent, session, context, state) -> None:
        self.after_run_texts.append(context.response.text if context.response else None)


@agent_middleware
async def passthrough(context, call_next):
    await call_next()


async def run_case(with_middleware: bool, exc: BaseException) -> None:
    recorder = RecordingProvider()
    history = InMemoryHistoryProvider()
    agent = Agent(
        client=FakeStreamingClient(),
        context_providers=[history, recorder],
        middleware=[passthrough] if with_middleware else None,
    )
    session = agent.create_session()

    def failing_hook(update):
        raise exc

    stream = agent.run("hi", stream=True, session=session).with_transform_hook(failing_hook)
    try:
        async for _ in stream:
            pass
    except BaseException as e:  # noqa: BLE001
        raised = type(e).__name__
    else:
        raised = "nothing"

    saved = await history.get_messages(session.session_id, state=session.state.get(history.source_id, {}))
    label = "with middleware   " if with_middleware else "without middleware"
    print(
        f"{label} {type(exc).__name__:<14} raised={raised:<14} "
        f"after_run calls={recorder.after_run_texts!r} history={[m.text for m in saved]!r}"
    )


async def main() -> None:
    for exc in (RuntimeError("hook failed"), asyncio.CancelledError()):
        await run_case(with_middleware=False, exc=exc)
        await run_case(with_middleware=True, exc=exc)
    print("expected: after_run calls=[] and history=[] in every case, since the run failed")


asyncio.run(main())

Output on 1.22.0:

without middleware RuntimeError   raised=RuntimeError   after_run calls=[] history=[]
with middleware    RuntimeError   raised=RuntimeError   after_run calls=['part0 '] history=['hi', 'part0 ']
without middleware CancelledError raised=CancelledError after_run calls=[] history=[]
with middleware    CancelledError raised=CancelledError after_run calls=['part0 '] history=['hi', 'part0 ']
expected: after_run calls=[] and history=[] in every case, since the run failed

The output is the same with the current main (44fe223) core sources. With the python-1.19.0 core sources every line shows after_run calls=[] history=[]. With ENABLE_INSTRUMENTATION=false on 1.22.0 every line is clean as well.

Error Messages and Stack Traces

No unexpected error. The RuntimeError / CancelledError from the hook propagates to the caller as it should; the problem is the after_run call and the saved history.

Package Versions

agent-framework-core: 1.22.0 (python-1.22.0)

Python Version

Python 3.12.3

Operating System

Linux

Regression

Yes

Additional Context

The error handler in ResponseStream.__anext__ (python/packages/core/agent_framework/_types.py lines 4769-4786 at 44fe223) records the exception only on the stream it ran in, via _handle_stream_error, and then calls self.close(). close() closes a nested ResponseStream at line 4798 without passing the error along, so the inner stream's _stream_error is still None when its cleanup hooks run. With agent middleware the stream I iterate is the ResponseStream.from_awaitable(_execute_stream()) wrapper from _middleware.py (lines 2100-2103), and the inner stream is the agent telemetry stream. Its _finalize_stream hook (observability.py 2444-2458) checks _stream_error, sees None, treats the run as successful and calls get_final_response(), which runs the agent's _post_hook result hook and so every after_run, including HistoryProvider.after_run (_sessions.py 1104-1128) saving the partial messages. The _stream_error guard in _finalize_stream exists to stop exactly this. The close() call in the error handler came in with #8773 (python-1.20.0); before that the inner stream's cleanup did not run on an outer error, which is why 1.19.0 is clean. Without middleware the outer stream is the telemetry stream itself, so it sees the error.

The common real-world trigger is a server cancelling the consuming task when the client disconnects: the cancelled turn is then stored in the session as if it had completed.

Proposed fix: propagate the terminal error into nested streams before closing them, for example give close() an optional exception argument that it hands to a nested ResponseStream (setting its _stream_error or calling its _handle_stream_error(exc) before closing it), and pass the exception from the __anext__ error handler. A regression test with an agent with middleware and a transform hook that raises or cancels, asserting that after_run is not called, would cover it. I'm happy to open a PR for this.

Acknowledgements

• I searched existing issues and did not find a duplicate.
• I personally verified this behavior and the reproduction details are authentic.
• I will wait for explicit maintainer agreement before starting implementation of a non-trivial change.