fix: Client hangs when implementing `AgentExecutor` and `await`ing tw… · a2aproject/a2a-python@c147a83 · GitHub
Skip to content

Commit c147a83

Browse files
meoow113严骏驰
andauthored
fix: Client hangs when implementing AgentExecutor and awaiting twice in execute method (#379)
Fixes #367 - Add clear_events method for immediate queue cleanup - Integrate with event_consumer to prevent hanging - Support child queue clearing for complete cleanup --------- Co-authored-by: 严骏驰 <yanjunchi@chinamobile.com>
1 parent db82a65 commit c147a83

3 files changed

Lines changed: 199 additions & 9 deletions

File tree

‎src/a2a/server/events/event_consumer.py‎

Lines changed: 1 addition & 1 deletion

‎src/a2a/server/events/event_queue.py‎

Lines changed: 71 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -132,25 +132,38 @@ def tap(self) -> 'EventQueue':
132132
self._children.append(queue)
133133
return queue
134134

135-
async def close(self) -> None:
136-
"""Closes the queue for future push events.
135+
async def close(self, immediate: bool = False) -> None:
136+
"""Closes the queue for future push events and also closes all child queues.
137+
138+
Once closed, no new events can be enqueued. For Python 3.13+, this will trigger
139+
`asyncio.QueueShutDown` when the queue is empty and a consumer tries to dequeue.
140+
For lower versions, the queue will be marked as closed and optionally cleared.
141+
142+
Args:
143+
immediate (bool):
144+
- True: Immediately closes the queue and clears all unprocessed events without waiting for them to be consumed. This is suitable for scenarios where you need to forcefully interrupt and quickly release resources.
145+
- False (default): Gracefully closes the queue, waiting for all queued events to be processed (i.e., the queue is drained) before closing. This is suitable when you want to ensure all events are handled.
137146
138-
Once closed, `dequeue_event` will eventually raise `asyncio.QueueShutDown`
139-
when the queue is empty. Also closes all child queues.
140147
"""
141148
logger.debug('Closing EventQueue.')
142149
async with self._lock:
143150
# If already closed, just return.
144-
if self._is_closed:
151+
if self._is_closed and not immediate:
145152
return
146-
self._is_closed = True
153+
if not self._is_closed:
154+
self._is_closed = True
147155
# If using python 3.13 or higher, use the shutdown method
148156
if sys.version_info >= (3, 13):
149-
self.queue.shutdown()
157+
self.queue.shutdown(immediate)
150158
for child in self._children:
151-
await child.close()
159+
await child.close(immediate)
152160
# Otherwise, join the queue
153161
else:
162+
if immediate:
163+
await self.clear_events(True)
164+
for child in self._children:
165+
await child.close(immediate)
166+
return
154167
tasks = [asyncio.create_task(self.queue.join())]
155168
tasks.extend(
156169
asyncio.create_task(child.close()) for child in self._children
@@ -160,3 +173,53 @@ async def close(self) -> None:
160173
def is_closed(self) -> bool:
161174
"""Checks if the queue is closed."""
162175
return self._is_closed
176+
177+
async def clear_events(self, clear_child_queues: bool = True) -> None:
178+
"""Clears all events from the current queue and optionally all child queues.
179+
180+
This method removes all pending events from the queue without processing them.
181+
Child queues can be optionally cleared based on the clear_child_queues parameter.
182+
183+
Args:
184+
clear_child_queues: If True (default), clear all child queues as well.
185+
If False, only clear the current queue, leaving child queues untouched.
186+
"""
187+
logger.debug('Clearing all events from EventQueue and child queues.')
188+
189+
# Clear all events from the queue, even if closed
190+
cleared_count = 0
191+
async with self._lock:
192+
try:
193+
while True:
194+
event = self.queue.get_nowait()
195+
logger.debug(
196+
f'Discarding unprocessed event of type: {type(event)}, content: {event}'
197+
)
198+
self.queue.task_done()
199+
cleared_count += 1
200+
except asyncio.QueueEmpty:
201+
pass
202+
except Exception as e:
203+
# Handle Python 3.13+ QueueShutDown
204+
if (
205+
sys.version_info >= (3, 13)
206+
and type(e).__name__ == 'QueueShutDown'
207+
):
208+
pass
209+
else:
210+
raise
211+
212+
if cleared_count > 0:
213+
logger.debug(
214+
f'Cleared {cleared_count} unprocessed events from EventQueue.'
215+
)
216+
217+
# Clear all child queues (lock released before awaiting child tasks)
218+
if clear_child_queues and self._children:
219+
child_tasks = [
220+
asyncio.create_task(child.clear_events())
221+
for child in self._children
222+
]
223+
224+
if child_tasks:
225+
await asyncio.gather(*child_tasks, return_exceptions=True)

‎tests/server/events/test_event_queue.py‎

Lines changed: 127 additions & 0 deletions

0 commit comments

Comments
 (0)