Skip to content
Merged
Changes from all commits
Commits
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
7 changes: 6 additions & 1 deletion src/a2a/server/events/event_consumer.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,6 @@ async def consume_all(self) -> AsyncGenerator[Event]:
logger.debug(
f'Dequeued event of type: {type(event)} in consume_all.'
)
yield event
self.queue.task_done()
logger.debug(
'Marked task as done in event queue in consume_all'
Expand All @@ -117,10 +116,16 @@ async def consume_all(self) -> AsyncGenerator[Event]:
)
)

# Make sure the yield is after the close events, otherwise
# the caller may end up in a blocked state where this
# generator isn't called again to close things out and the
# other part is waiting for an event or a closed queue.
if is_final_event:
logger.debug('Stopping event consumption in consume_all.')
await self.queue.close()
yield event
break
yield event
except TimeoutError:
# continue polling until there is a final event
continue
Expand Down