Skip to content

Python: Close delegated provider streams on public stream release #9100

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Eduard van Valkenburg (eavanvalkenburg) merged 3 commits into from
Oct 8, 2026
Merged

Python: Close delegated provider streams on public stream release #9100

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
57a72d4
Python: Close delegated provider streams on public stream release
quifox Oct 6, 2026
fc985f5
Python: Keep stream-close test consumption in one task
quifox Oct 7, 2026
f96aabc
Python: preserve stream failures during delegated cleanup
quifox Oct 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 45 additions & 27 deletions python/packages/ag-ui/agent_framework_ag_ui/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,37 +503,55 @@ async def _streaming_impl(

available_interrupts = options.get("available_interrupts", options.get("availableInterrupts"))

async for event in self._http_service.post_run(
events = self._http_service.post_run(
thread_id=thread_id,
run_id=run_id,
messages=agui_messages,
state=state,
tools=agui_tools,
available_interrupts=_serialize_available_interrupts(cast(Sequence[Any] | None, available_interrupts)),
resume=_serialize_resume(options.get("resume")),
):
logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}")
update = converter.convert_event(event)
if update is not None:
logger.debug(
"[AGUIChatClient] Converted update",
extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]},
)
# Distinguish client vs server tools
for i, content in enumerate(update.contents):
if content.type == "function_call":
logger.debug(
f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}"
)
if content.name in client_tool_set:
# Client tool - let function invocation execute it
if not content.additional_properties:
content.additional_properties = {}
content.additional_properties["agui_thread_id"] = thread_id
else:
# Server tool - wrap so function invocation ignores it
logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}")
self._register_server_tool_placeholder(content.name) # type: ignore[arg-type]
update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore

yield update
)
body_failed = False
try:
async for event in events:
logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}")
update = converter.convert_event(event)
if update is not None:
logger.debug(
"[AGUIChatClient] Converted update",
extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]},
)
# Distinguish client vs server tools
for i, content in enumerate(update.contents):
if content.type == "function_call":
logger.debug(
f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}"
)
if content.name in client_tool_set:
# Client tool - let function invocation execute it
if not content.additional_properties:
content.additional_properties = {}
content.additional_properties["agui_thread_id"] = thread_id
else:
# Server tool - wrap so function invocation ignores it
logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}")
self._register_server_tool_placeholder(content.name) # type: ignore[arg-type]
update.contents[i] = Content(type="server_function_call", function_call=content) # type: ignore

yield update
except GeneratorExit:
# An explicit close() has no body failure to take precedence over cleanup.
raise
except BaseException:
body_failed = True
raise
finally:
close = getattr(events, "aclose", None)
if close is not None:
try:
await close()
except Exception:
if not body_failed:
raise
logger.warning("AG-UI event cleanup failed while handling an existing exception.", exc_info=True)
29 changes: 23 additions & 6 deletions python/packages/ag-ui/agent_framework_ag_ui/_http_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@
import json
import logging
from collections.abc import AsyncIterable, Mapping, Sequence
from contextlib import AsyncExitStack
from http.cookiejar import CookieJar, DefaultCookiePolicy
from types import TracebackType
from typing import Any, cast

import httpx
Expand Down Expand Up @@ -238,12 +240,27 @@ async def post_run(
)

# Stream the response using SSE
async with self.http_client.stream(
"POST",
self.endpoint,
json=request_data,
headers={"Accept": "text/event-stream"},
) as response:
async with AsyncExitStack() as stack:
stream_context = self.http_client.stream(
"POST",
self.endpoint,
json=request_data,
headers={"Accept": "text/event-stream"},
)
response = await stream_context.__aenter__()

async def exit_response(
exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> bool | None:
try:
return await stream_context.__aexit__(exc_type, exc, tb)
except Exception:
if exc is None or isinstance(exc, GeneratorExit):
raise
logger.warning("HTTP response cleanup failed while handling an existing exception.", exc_info=True)
return False

stack.push_async_exit(exit_response)
try:
response.raise_for_status()
except httpx.HTTPStatusError as e:
Expand Down
Loading
Loading