xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-01 19:30:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Ensure asyncio event queue is drained prior to CLI termination

Browse Source
This commit is contained in:
MarkRx
2024-08-26 15:14:45 -04:00
parent 53a974c282
commit 93773f3c08
1 changed files with 16 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
pr_agent/cli.py
View File

@ -4,7 +4,7 @@ import os
from pr_agent.agent.pr_agent import PRAgent, commands
from pr_agent.config_loader import get_settings
from pr_agent.log import setup_logger
from pr_agent.log import setup_logger, get_logger
log_level = os.environ.get("LOG_LEVEL", "INFO")
setup_logger(log_level)
@ -71,10 +71,21 @@ def run(inargs=None, args=None):
command = args.command.lower()
get_settings().set("CONFIG.CLI_MODE", True)
if args.issue_url:
result = asyncio.run(PRAgent().handle_request(args.issue_url, [command] + args.rest))
else:
result = asyncio.run(PRAgent().handle_request(args.pr_url, [command] + args.rest))
async def inner():
if args.issue_url:
result = await asyncio.create_task(PRAgent().handle_request(args.issue_url, [command] + args.rest))
else:
result = await asyncio.create_task(PRAgent().handle_request(args.pr_url, [command] + args.rest))
if get_settings().litellm.get("enable_callbacks", False):
# There may be additional events on the event queue from the run above. If there are give them time to complete.
get_logger().debug("Waiting for event queue to complete")
await asyncio.wait([task for task in asyncio.all_tasks() if task is not asyncio.current_task()])
return result
result = asyncio.run(inner())
if not result:
parser.print_help()