Skip to content

Commit

Permalink
fix(core): bureau shutdown routine
Browse files Browse the repository at this point in the history
  • Loading branch information
Archento committed Nov 19, 2024
1 parent 39947f7 commit f1fcfe3
Showing 1 changed file with 12 additions and 5 deletions.
17 changes: 12 additions & 5 deletions python/src/uagents/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -1539,18 +1539,25 @@ async def run_async(self):
await agent.setup()
if agent.agentverse["use_mailbox"] and agent.mailbox_client is not None:
tasks.append(agent.mailbox_client.run())
tasks.append(self._schedule_registration())
self._loop.create_task(self._schedule_registration())

try:
await asyncio.gather(*tasks, return_exceptions=True)
except (asyncio.CancelledError, KeyboardInterrupt):
pass
finally:
await asyncio.gather(
*[agent._shutdown() for agent in self._agents], return_exceptions=True
)
for agent in self._agents:
await agent._shutdown()
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
_ = [task.cancel() for task in tasks]
await asyncio.gather(*tasks)

def run(self):
"""
Run the bureau.
"""
self._loop.run_until_complete(self.run_async())
with contextlib.suppress(asyncio.CancelledError, KeyboardInterrupt):
self._loop.run_until_complete(self.run_async())
self._loop.stop()
self._loop.close()

0 comments on commit f1fcfe3

Please sign in to comment.