fix: ensure stream is closed on exit

This commit is contained in:
Greyson LaLonde
2026-01-06 18:52:01 -05:00
parent f53b8755da
commit c639455730
2 changed files with 55 additions and 44 deletions

View File

@@ -264,6 +264,7 @@ async def send_message_and_get_task_id(
Returns:
Task ID string if agent needs polling/waiting, or TaskStateResult if done.
"""
try:
async for event in event_stream:
if isinstance(event, Message):
new_messages.append(event)
@@ -312,3 +313,8 @@ async def send_message_and_get_task_id(
error="No task ID received from initial message",
history=new_messages,
)
finally:
aclose = getattr(event_stream, "aclose", None)
if aclose:
await aclose()

View File

@@ -133,6 +133,11 @@ class StreamingHandler:
history=new_messages,
)
finally:
aclose = getattr(event_stream, "aclose", None)
if aclose:
await aclose()
if final_result:
return final_result