diff --git a/skyvern/forge/sdk/core/retry.py b/skyvern/forge/sdk/core/retry.py new file mode 100644 index 00000000..2579024f --- /dev/null +++ b/skyvern/forge/sdk/core/retry.py @@ -0,0 +1,65 @@ +import asyncio +import time +from functools import wraps +from typing import Any, Callable, Type, TypeVar, cast + +import structlog + +LOG = structlog.get_logger() + +# Type for any callable (sync or async) +F = TypeVar("F", bound=Callable[..., Any]) + + +def retry( + exceptions: Type[Exception] | tuple[Type[Exception], ...] | None = None, + tries: int = 3, + delay: int = 3, + backoff: int = 2, +) -> Callable[[F], F]: + """ + Decorator to retry a function a specified number of times with a delay between attempts. + Works with both async and sync functions. + + Args: + exceptions: A tuple of exceptions to catch and retry on. + tries: The total number attempts to make. + delay: The initial delay between attempts. + backoff: The factor by which the delay increases after each attempt. + """ + if exceptions is None: + exceptions = Exception + + def retry_decorator(func: F) -> F: + @wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + remaining_tries, current_delay = tries, delay + while remaining_tries > 1: + try: + return await func(*args, **kwargs) + except exceptions as e: + LOG.warning(f"Retrying {func.__name__} in {current_delay} seconds... {e}") + await asyncio.sleep(current_delay) + remaining_tries -= 1 + current_delay *= backoff + return await func(*args, **kwargs) + + @wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: + remaining_tries, current_delay = tries, delay + while remaining_tries > 1: + try: + return func(*args, **kwargs) + except exceptions as e: + LOG.warning(f"Retrying {func.__name__} in {current_delay} seconds... {e}") + time.sleep(current_delay) + remaining_tries -= 1 + current_delay *= backoff + return func(*args, **kwargs) + + # Return async wrapper if function is async, otherwise sync wrapper + if asyncio.iscoroutinefunction(func): + return cast(F, async_wrapper) + return cast(F, sync_wrapper) + + return retry_decorator