Ensure asyncio event queue is drained prior to CLI termination

[MarkRx]

User description

The LiteLLM library may add events to the asyncio event queue from a LLM execution (such as to run event handlers from callbacks). If the CLI terminates before these events run the agent may crash with asyncio.exceptions.CancelledError.

This change gives the CLI a little extra time to ensure the event queue is completed before terminating. The timeout is hard coded to reduce configuration sprawl but could be a configuration option if desired. If there are no extra events then the CLI terminates immediately.

See BerriAI/litellm#5293 (comment)

---

PR Type

Bug fix, Enhancement

---

Description

• Added a mechanism to ensure the asyncio event queue is drained before CLI termination
• Implemented a 30-second timeout for event queue completion
• Wrapped the main PRAgent().handle_request call in asyncio.create_task for better task management
• Improved error handling by allowing time for potential lingering events to complete
• Enhanced debugging capabilities by adding a debug log message for the waiting period

---

Changes walkthrough 📝

	Relevant files
Error handling	cli.py Implement asyncio event queue draining mechanism                  pr_agent/cli.py Imported get_logger function from pr_agent.log Replaced direct asyncio.run call with a new async function inner() Added a timeout mechanism to wait for event queue completion Wrapped PRAgent().handle_request in asyncio.create_task +17/-5

---

💡 PR-Agent usage:
Comment /help on the PR to get a list of all available PR-Agent tools and their descriptions

[codiumai-pr-agent-pro]

PR-Agent was enabled for this repository. To continue using it, please link your git user with your CodiumAI identity here.

PR Reviewer Guide 🔍

⏱️ Estimated effort to review: 2 🔵🔵⚪⚪⚪

🏅 Score: 95

🧪 No relevant tests

🔒 No security concerns identified

🔀 No multiple PR themes

⚡ Key issues to review Potential Performance Impact The added while loop with a sleep of 0.1 seconds might cause unnecessary delays if there are no pending tasks.

[codiumai-pr-agent-pro]

PR-Agent was enabled for this repository. To continue using it, please link your git user with your CodiumAI identity here.

PR Code Suggestions ✨

Category	Suggestion	Score
Performance	Use a more efficient method to wait for asynchronous tasks to complete Replace the while loop with asyncio.wait() to more efficiently wait for all tasks to complete. pr_agent/cli.py [83-85] async with asyncio.timeout(30): - while len(asyncio.all_tasks()) > 1: - await asyncio.sleep(0.1) + await asyncio.wait([task for task in asyncio.all_tasks() if task is not asyncio.current_task()]) Apply this suggestion Suggestion importance[1-10]: 9 Why: Replacing the while loop with asyncio.wait() is more efficient and idiomatic for waiting on multiple tasks to complete in asyncio.	9
Best practice	Use a more appropriate asyncio function for task handling Consider using asyncio.gather() instead of asyncio.create_task() to ensure proper task handling and avoid potential race conditions. pr_agent/cli.py [75-79] async def inner(): if args.issue_url: - result = await asyncio.create_task(PRAgent().handle_request(args.issue_url, [command] + args.rest)) + result = await asyncio.gather(PRAgent().handle_request(args.issue_url, [command] + args.rest)) else: - result = await asyncio.create_task(PRAgent().handle_request(args.pr_url, [command] + args.rest)) + result = await asyncio.gather(PRAgent().handle_request(args.pr_url, [command] + args.rest)) Apply this suggestion Suggestion importance[1-10]: 8 Why: Using asyncio.gather() instead of asyncio.create_task() is a better practice for handling multiple tasks concurrently and ensures proper task execution.	8
Error handling	✅ Add error handling for potential timeout scenarios Consider handling potential TimeoutError that could be raised by the asyncio.timeout() context manager. pr_agent/cli.py [83-85] -async with asyncio.timeout(30): - while len(asyncio.all_tasks()) > 1: - await asyncio.sleep(0.1) +try: + async with asyncio.timeout(30): + while len(asyncio.all_tasks()) > 1: + await asyncio.sleep(0.1) +except asyncio.TimeoutError: + get_logger().warning("Timeout reached while waiting for event queue to complete") Apply this suggestion Suggestion importance[1-10]: 7 Why: Handling the potential TimeoutError improves the robustness of the code, although the current implementation may already handle this implicitly.	7
Maintainability	Improve code organization by extracting a reusable function Consider extracting the event queue draining logic into a separate function for better code organization and reusability. pr_agent/cli.py [81-85] -# There may be additional events on the event queue from the run above. If there are give them up to 30 seconds to complete. -get_logger().debug("Waiting up to 30 seconds for event queue to complete") -async with asyncio.timeout(30): - while len(asyncio.all_tasks()) > 1: - await asyncio.sleep(0.1) +async def drain_event_queue(timeout=30): + get_logger().debug(f"Waiting up to {timeout} seconds for event queue to complete") + try: + async with asyncio.timeout(timeout): + await asyncio.wait([task for task in asyncio.all_tasks() if task is not asyncio.current_task()]) + except asyncio.TimeoutError: + get_logger().warning("Timeout reached while waiting for event queue to complete") +# ... in the inner() function: +await drain_event_queue() + Apply this suggestion Suggestion importance[1-10]: 6 Why: Extracting the event queue draining logic into a separate function improves code organization and reusability, but it's a minor improvement in this context.	6

[mrT23]

I don't want to wait 30 additional seconds each time I run the CLI.
I think its hacky, and not a good tradeoff

[MarkRx]

It doesn't take an additional 30 seconds to run the CLI - it terminates as soon as the extra tasks complete. The 30 second timeout was a defensive guard to prevent the CLI from hanging in case something went wrong.

[mrT23]

@MarkRx
the while loop is ugly. And this is a major file.

i suggest a one-liner:

and also add if:

if litellm.enable_callbacks:
...