vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

add wait for async task (#344)

Browse Source
This commit is contained in:
LawyZheng
2024-05-20 12:59:53 +08:00
committed by GitHub
parent 6908ca5844
commit b6d25dff8c
1 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
skyvern/forge/async_operations.py
View File

@@ -96,6 +96,16 @@ class AsyncOperationPool:
def get_aio_task(self, task_id: str, operation_type: str) -> asyncio.Task | None:
return self._aio_tasks.get(task_id, {}).get(operation_type, None)
async def wait_for_task(self, task_id: str, operation_type: str, timeout: float | None = None) -> None:
running_task = self.get_aio_task(task_id=task_id, operation_type=operation_type)
if running_task and not running_task.done():
LOG.info(
"wait for the running aio task to be done",
task_id=task_id,
operation_type=operation_type,
)
await asyncio.wait_for(running_task, timeout)
def run_operation(self, task_id: str, agent_phase: AgentPhase) -> None:
# get the operation from the pool
operation = self._get_operation(task_id, agent_phase)