Description

After consuming the first text update, closing a public Python response stream can return before the delegated provider stream has released its HTTP response. The same happens when leaving async with stream after breaking out of iteration.

For AG-UI this reproduces with AGUIChatClient.get_response(..., stream=True) and Agent.run(..., stream=True), both with default function invocation and with function invocation disabled. OpenAI Chat Completions also reproduces through the default function-invocation layer; the direct provider path with that layer disabled closes correctly.

Expected: once await stream.close() or the stream context exit completes, the active delegated response is closed. A caller-supplied shared HTTP client should remain usable.

Actual: the public close completes while the response/body remains open. This is a deterministic cleanup/ownership gap; it does not depend on a live model or garbage collection.

Code Sample

The following uses HTTPX MockTransport and makes no network requests. Keep the response and body referenced, consume only the first chunk, then check cleanup immediately after public close:

import asyncio
import json

import httpx
from agent_framework import Message
from agent_framework_ag_ui import AGUIChatClient


class Body(httpx.AsyncByteStream):
    def __init__(self):
        self.closed = False
        self.chunks_read = 0

    async def __aiter__(self):
        for text in ("first", "later"):
            self.chunks_read += 1
            event = {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": text}
            yield f"data: {json.dumps(event)}\n\n".encode()

    async def aclose(self):
        self.closed = True


async def main():
    body = Body()
    responses = []

    async def handler(request):
        await request.aread()
        response = httpx.Response(
            200, headers={"content-type": "text/event-stream"}, stream=body, request=request
        )
        responses.append(response)
        return response

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
        client = AGUIChatClient(endpoint="https://agui.example.test/run", http_client=http_client)
        stream = client.get_response([Message(role="user", contents=["Question"])], stream=True)
        try:
            assert (await anext(stream)).text == "first"
            assert body.chunks_read == 1
            await stream.close()
            print({"response_closed": responses[0].is_closed, "body_closed": body.closed})
            assert responses[0].is_closed and body.closed
            assert not http_client.is_closed
        finally:
            await stream.close()
            await responses[0].aclose()
            await client.close()


asyncio.run(main())

On the current implementation, both closure flags are False and the assertion fails. With the forwarding cleanup fixed, both values are True.

Package Versions

Source checkout at b9d24c8fb484c8330abe8bb9e7500ca3c3bbf46c: agent-framework-core 1.20.0 and agent-framework-ag-ui 1.5.0.

Python Version

Python 3.12.12, Linux x86_64.

Additional Context

The provider-specific OpenAI cleanup in #8773 fixed #8762. This report concerns the remaining forwarding layers: function invocation owns an inner ResponseStream, while the AG-UI conversion generator owns the post_run event stream. Neither forwarding scope currently awaits inner closure when the consumer stops early.

Full consumption and directly closing the AG-UI HTTP-service generator pass as controls. Loopback TCP/SSE checks confirm the early-close gap with real HTTPX socket transport as well.

I will submit a focused fix using the existing ResponseStream context manager and awaited AG-UI event-stream closure, with regression tests for public client and agent entry points.