kash-shell 0.3.23__py3-none-any.whl → 0.3.24__py3-none-any.whl

kash/utils/api_utils/gather_limited.py

@@ -4,9 +4,10 @@ import asyncio
 import inspect
 import logging
 import threading
+import time
 from collections.abc import Callable, Coroutine
 from dataclasses import dataclass, field
-from typing import Any, TypeAlias, TypeVar, cast, overload
+from typing import Any, Generic, TypeAlias, TypeVar, cast, overload
 
 from aiolimiter import AsyncLimiter
 
@@ -17,7 +18,6 @@ from kash.utils.api_utils.api_retries import (
     RetrySettings,
     calculate_backoff,
     extract_http_status_code,
-    is_http_status_retriable,
 )
 from kash.utils.api_utils.progress_protocol import Labeler, ProgressTracker, TaskState
 
@@ -25,10 +25,6 @@ T = TypeVar("T")
 
 log = logging.getLogger(__name__)
 
-DEFAULT_MAX_CONCURRENT: int = 5
-DEFAULT_MAX_RPS: float = 5.0
-DEFAULT_CANCEL_TIMEOUT: float = 1.0
-
 
 @dataclass(frozen=True)
 class Limit:
@@ -36,27 +32,48 @@ class Limit:
     Rate limiting configuration with max RPS and max concurrency.
     """
 
-    rps: float = DEFAULT_MAX_RPS
-    concurrency: int = DEFAULT_MAX_CONCURRENT
+    rps: float
+    concurrency: int
+
+
+DEFAULT_CANCEL_TIMEOUT: float = 1.0
 
 
 @dataclass(frozen=True)
-class FuncTask:
+class TaskResult(Generic[T]):
+    """
+    Optional wrapper for task results that can signal rate limiting behavior.
+    Use this to wrap results that should bypass rate limiting (e.g., cache hits).
+    """
+
+    value: T
+    disable_limits: bool = False
+
+
+@dataclass(frozen=True)
+class FuncTask(Generic[T]):
     """
     A task described as an unevaluated function with args and kwargs.
     This task format allows you to use args and kwargs in the Labeler.
     It also allows specifying a bucket for rate limiting.
+
+    For async functions: The function should return a coroutine that yields either `T` or `TaskResult[T]`.
+    For sync functions: The function should return either `T` or `TaskResult[T]` directly.
+
+    Using `TaskResult[T]` allows controlling rate limiting behavior (e.g., cache hits can bypass rate limits).
     """
 
-    func: Callable[..., Any]
+    func: Callable[..., Any]  # Keep as Any since it can be sync or async
     args: tuple[Any, ...] = ()
     kwargs: dict[str, Any] = field(default_factory=dict)
     bucket: str = "default"
 
 
 # Type aliases for coroutine and sync specifications, including unevaluated function specs
-CoroSpec: TypeAlias = Callable[[], Coroutine[None, None, T]] | Coroutine[None, None, T] | FuncTask
-SyncSpec: TypeAlias = Callable[[], T] | FuncTask
+CoroSpec: TypeAlias = (
+    Callable[[], Coroutine[None, None, T]] | Coroutine[None, None, T] | FuncTask[T]
+)
+SyncSpec: TypeAlias = Callable[[], T] | FuncTask[T]
 
 # Specific labeler types using the generic Labeler pattern
 CoroLabeler: TypeAlias = Labeler[CoroSpec[T]]
@@ -114,7 +131,7 @@ class RetryCounter:
 @overload
 async def gather_limited_async(
     *coro_specs: CoroSpec[T],
-    global_limit: Limit = Limit(),
+    limit: Limit | None,
     bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = False,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
@@ -126,7 +143,7 @@ async def gather_limited_async(
 @overload
 async def gather_limited_async(
     *coro_specs: CoroSpec[T],
-    global_limit: Limit = Limit(),
+    limit: Limit | None,
     bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
@@ -137,7 +154,7 @@ async def gather_limited_async(
 
 async def gather_limited_async(
     *coro_specs: CoroSpec[T],
-    global_limit: Limit = Limit(),
+    limit: Limit | None,
     bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,  # Default to True for resilient batch operations
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
@@ -157,12 +174,22 @@ async def gather_limited_async(
     - Configurable handling of conditional status codes like 403 Forbidden
     - Defaults to return_exceptions=True for resilient batch operations
 
+    Prevents API flooding during rate limit backoffs:
+    - Semaphore slots are held during entire retry cycles, including backoff periods
+    - New tasks cannot start while existing tasks are backing off from rate limits
+    - Rate limiters are only acquired during actual execution attempts
+
     Can optionally display live progress with retry indicators using TaskStatus.
 
     Accepts:
     - Callables that return coroutines: `lambda: some_async_func(arg)` (recommended for retries)
     - Coroutines directly: `some_async_func(arg)` (only if retries disabled)
-    - FuncSpec objects: `FuncSpec(some_async_func, (arg1, arg2), {"kwarg": value})` (args accessible to labeler)
+    - FuncTask objects: `FuncTask(some_async_func, (arg1, arg2), {"kwarg": value})` (args accessible to labeler)
+
+    Functions can return `TaskResult[T]` to bypass rate limiting for specific results:
+    - `TaskResult(value, disable_limits=True)`: bypass rate limiting (e.g., cache hits)
+    - `TaskResult(value, disable_limits=False)`: apply normal rate limiting
+    - `value` directly: apply normal rate limiting
 
     Examples:
         ```python
@@ -190,18 +217,34 @@ async def gather_limited_async(
             FuncTask(fetch_api, ("data1",), bucket="api1"),
             FuncTask(fetch_api, ("data2",), bucket="api2"),
             FuncTask(fetch_api, ("data3",), bucket="api3"),
-            global_limit=Limit(rps=100, concurrency=50),
+            limit=Limit(rps=100, concurrency=50),
             bucket_limits={
                 "api1": Limit(rps=20, concurrency=10),  # Specific limit for api1
                 "*": Limit(rps=15, concurrency=8),      # Fallback for api2, api3, and others
             }
         )
 
+        # With cache bypass using TaskResult:
+        async def fetch_with_cache(url: str) -> TaskResult[dict] | dict:
+            cached_data = await cache.get(url)
+            if cached_data:
+                return TaskResult(cached_data, disable_limits=True)  # Cache hit: no rate limit
+
+            data = await fetch_api(url)  # Will be rate limited
+            await cache.set(url, data)
+            return data
+
+        await gather_limited_async(
+            lambda: fetch_with_cache("http://api.com/data1"),
+            lambda: fetch_with_cache("http://api.com/data2"),
+            limit=Limit(rps=10, concurrency=5)  # Cache hits bypass these limits
+        )
+
         ```
 
     Args:
         *coro_specs: Callables or coroutines to execute
-        global_limit: Global limits applied to all tasks regardless of bucket
+        limit: Global limits applied to all tasks regardless of bucket
         bucket_limits: Optional per-bucket limits. Tasks use their bucket field to determine limits.
                       Use "*" as a fallback limit for buckets without specific limits.
         return_exceptions: If True, exceptions are returned as results
@@ -217,8 +260,8 @@ async def gather_limited_async(
     """
     log.info(
         "Executing with global limits: concurrency %s at %s rps, %s",
-        global_limit.concurrency,
-        global_limit.rps,
+        limit.concurrency if limit else "None",
+        limit.rps if limit else "None",
         retry_settings,
     )
     if not coro_specs:
@@ -237,8 +280,8 @@ async def gather_limited_async(
                 )
 
     # Global limits (apply to all tasks regardless of bucket)
-    global_semaphore = asyncio.Semaphore(global_limit.concurrency)
-    global_rate_limiter = AsyncLimiter(global_limit.rps, 1.0)
+    global_semaphore = asyncio.Semaphore(limit.concurrency) if limit else None
+    global_rate_limiter = AsyncLimiter(limit.rps, 1.0) if limit else None
 
     # Per-bucket limits (if bucket_limits provided)
     bucket_semaphores: dict[str, asyncio.Semaphore] = {}
@@ -275,7 +318,7 @@ async def gather_limited_async(
             elif callable(coro_spec):
                 coro = coro_spec()
             else:
-                # Direct coroutine - only valid if retries disabled
+                # Direct coroutine: only valid if retries disabled
                 coro = coro_spec
             return await coro
 
@@ -313,7 +356,7 @@ async def gather_limited_async(
 @overload
 async def gather_limited_sync(
     *sync_specs: SyncSpec[T],
-    global_limit: Limit = Limit(),
+    limit: Limit | None,
     bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = False,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
@@ -327,7 +370,7 @@ async def gather_limited_sync(
 @overload
 async def gather_limited_sync(
     *sync_specs: SyncSpec[T],
-    global_limit: Limit = Limit(),
+    limit: Limit | None,
     bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
@@ -340,7 +383,7 @@ async def gather_limited_sync(
 
 async def gather_limited_sync(
     *sync_specs: SyncSpec[T],
-    global_limit: Limit = Limit(),
+    limit: Limit | None,
     bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,  # Default to True for resilient batch operations
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
@@ -350,42 +393,38 @@ async def gather_limited_sync(
     cancel_timeout: float = DEFAULT_CANCEL_TIMEOUT,
 ) -> list[T] | list[T | BaseException]:
     """
-    Rate-limited version of `asyncio.gather()` for sync functions with HTTP-aware retry logic.
-    Handles the asyncio.to_thread() boundary correctly for consistent exception propagation.
+    Sync version of `gather_limited_async()` that executes synchronous functions with the same
+    rate limiting, retry logic, and progress tracking capabilities.
+    See `gather_limited_async()` for details.
 
-    Supports two levels of retry limits:
-    - Per-task retries: max_task_retries attempts per individual task
-    - Global retries: max_total_retries attempts across all tasks
+    Sync-specific differences:
 
-    Features HTTP-aware retry classification:
-    - Automatically detects HTTP status codes (403, 429, 500, etc.) and applies appropriate retry behavior
-    - Configurable handling of conditional status codes like 403 Forbidden
-    - Defaults to return_exceptions=True for resilient batch operations
+    **Function Execution**: Runs sync functions via `asyncio.to_thread()` instead of executing
+    coroutines directly. Validates that callables don't accidentally return coroutines.
+
+    **Cancellation Support**: Provides cooperative cancellation for long-running sync functions
+    through optional `cancel_event` and graceful thread termination.
 
-    Supports cooperative cancellation and graceful thread termination on interruption.
+    **Input Validation**: Ensures sync functions don't return coroutines, which would indicate
+    incorrect usage (async functions should use `gather_limited_async()`).
 
     Args:
-        *sync_specs: Callables that return values (not coroutines) or FuncTask objects
-        global_limit: Global limits applied to all tasks regardless of bucket
-        bucket_limits: Optional per-bucket limits. Tasks use their bucket field to determine limits.
-                      Use "*" as a fallback limit for buckets without specific limits.
-        return_exceptions: If True, exceptions are returned as results
-        retry_settings: Configuration for retry behavior, or None to disable retries
-        status: Optional ProgressTracker instance for progress display
-        labeler: Optional function to generate labels: labeler(index, spec) -> str
+        *sync_specs: Callables that return values (not coroutines) or FuncTask objects.
+                    Functions can return `TaskResult[T]` to control rate limiting behavior.
         cancel_event: Optional threading.Event that will be set on cancellation
         cancel_timeout: Max seconds to wait for threads to terminate on cancellation
+        (All other args identical to gather_limited_async())
 
     Returns:
         List of results in the same order as input specifications
 
     Example:
         ```python
-        # Without bucket limits (backward compatible)
+        # Basic usage with sync functions
         results = await gather_limited_sync(
             lambda: some_sync_function(arg1),
             lambda: another_sync_function(arg2),
-            global_limit=Limit(rps=2.0, concurrency=3),
+            limit=Limit(rps=2.0, concurrency=3),
             retry_settings=RetrySettings(max_task_retries=3, max_total_retries=25)
         )
 
@@ -394,18 +433,34 @@ async def gather_limited_sync(
             FuncTask(fetch_from_api, ("data1",), bucket="api1"),
             FuncTask(fetch_from_api, ("data2",), bucket="api2"),
             FuncTask(fetch_from_api, ("data3",), bucket="api3"),
-            global_limit=Limit(rps=100, concurrency=50),
+            limit=Limit(rps=100, concurrency=50),
             bucket_limits={
                 "api1": Limit(rps=20, concurrency=10),  # Specific limit for api1
                 "*": Limit(rps=15, concurrency=8),      # Fallback for api2, api3, and others
             }
         )
+
+        # With cache bypass using TaskResult
+        def sync_fetch_with_cache(url: str) -> TaskResult[dict] | dict:
+            cached_data = cache.get_sync(url)
+            if cached_data:
+                return TaskResult(cached_data, disable_limits=True)  # Cache hit: no rate limit
+
+            data = requests.get(url).json()  # Will be rate limited
+            cache.set_sync(url, data)
+            return data
+
+        results = await gather_limited_sync(
+            lambda: sync_fetch_with_cache("http://api.com/data1"),
+            lambda: sync_fetch_with_cache("http://api.com/data2"),
+            limit=Limit(rps=5, concurrency=3)  # Cache hits bypass these limits
+        )
         ```
     """
     log.info(
         "Executing with global limits: concurrency %s at %s rps, %s",
-        global_limit.concurrency,
-        global_limit.rps,
+        limit.concurrency if limit else "None",
+        limit.rps if limit else "None",
         retry_settings,
     )
     if not sync_specs:
@@ -414,8 +469,8 @@ async def gather_limited_sync(
     retry_settings = retry_settings or NO_RETRIES
 
     # Global limits (apply to all tasks regardless of bucket)
-    global_semaphore = asyncio.Semaphore(global_limit.concurrency)
-    global_rate_limiter = AsyncLimiter(global_limit.rps, 1.0)
+    global_semaphore = asyncio.Semaphore(limit.concurrency) if limit else None
+    global_rate_limiter = AsyncLimiter(limit.rps, 1.0) if limit else None
 
     # Per-bucket limits (if bucket_limits provided)
     bucket_semaphores: dict[str, asyncio.Semaphore] = {}
@@ -572,34 +627,94 @@ async def _gather_with_interrupt_handling(
 async def _execute_with_retry(
     executor: Callable[[], Coroutine[None, None, T]],
     retry_settings: RetrySettings,
-    global_semaphore: asyncio.Semaphore,
-    global_rate_limiter: AsyncLimiter,
+    global_semaphore: asyncio.Semaphore | None,
+    global_rate_limiter: AsyncLimiter | None,
     bucket_semaphore: asyncio.Semaphore | None,
     bucket_rate_limiter: AsyncLimiter | None,
     global_retry_counter: RetryCounter,
     status: ProgressTracker | None = None,
     task_id: Any | None = None,
 ) -> T:
-    import time
+    """
+    Execute a task with retry logic, holding semaphores for the entire retry cycle.
 
-    start_time = time.time()
-    last_exception: Exception | None = None
+    Semaphores are held during backoff periods to prevent API flooding when tasks
+    are already backing off from rate limits. Only rate limiters are acquired/released
+    for each execution attempt.
+    """
+    # Acquire semaphores once for the entire retry cycle (including backoff periods)
+    if global_semaphore is not None:
+        async with global_semaphore:
+            if bucket_semaphore is not None:
+                async with bucket_semaphore:
+                    return await _execute_with_retry_inner(
+                        executor,
+                        retry_settings,
+                        global_rate_limiter,
+                        bucket_rate_limiter,
+                        global_retry_counter,
+                        status,
+                        task_id,
+                    )
+            else:
+                return await _execute_with_retry_inner(
+                    executor,
+                    retry_settings,
+                    global_rate_limiter,
+                    None,
+                    global_retry_counter,
+                    status,
+                    task_id,
+                )
+    else:
+        # No global semaphore, check bucket semaphore
+        if bucket_semaphore is not None:
+            async with bucket_semaphore:
+                return await _execute_with_retry_inner(
+                    executor,
+                    retry_settings,
+                    global_rate_limiter,
+                    bucket_rate_limiter,
+                    global_retry_counter,
+                    status,
+                    task_id,
+                )
+        else:
+            # No semaphores at all, go straight to running
+            return await _execute_with_retry_inner(
+                executor,
+                retry_settings,
+                global_rate_limiter,
+                None,
+                global_retry_counter,
+                status,
+                task_id,
+            )
 
-    # Create HTTP-aware is_retriable function based on settings
-    def is_retriable_for_task(exception: Exception) -> bool:
-        # First check the main is_retriable function
-        if not retry_settings.is_retriable(exception):
-            return False
 
-        status_code = extract_http_status_code(exception)
-        if status_code:
-            return is_http_status_retriable(status_code, retry_settings.http_retry_map)
+async def _execute_with_retry_inner(
+    executor: Callable[[], Coroutine[None, None, T]],
+    retry_settings: RetrySettings,
+    global_rate_limiter: AsyncLimiter | None,
+    bucket_rate_limiter: AsyncLimiter | None,
+    global_retry_counter: RetryCounter,
+    status: ProgressTracker | None = None,
+    task_id: Any | None = None,
+) -> T:
+    """
+    Inner retry logic that handles rate limiting and backoff.
+
+    This function assumes semaphores are already held by the caller and only
+    manages rate limiters for each execution attempt.
+    """
+    if status and task_id:
+        await status.start(task_id)
 
-        # Not an HTTP error, use the main is_retriable result
-        return True
+    start_time = time.time()
+    last_exception: Exception | None = None
 
     for attempt in range(retry_settings.max_task_retries + 1):
-        # Handle backoff before acquiring any resources
+        # Handle backoff before acquiring rate limiters (semaphores remain held)
         if attempt > 0 and last_exception:
             # Try to increment global retry counter
             if not await global_retry_counter.try_increment():
@@ -624,7 +739,7 @@ async def _execute_with_retry(
                     f"Attempt {attempt}/{retry_settings.max_task_retries} "
                     f"(waiting {backoff_time:.1f}s): {type(last_exception).__name__}: {last_exception}"
                 )
-                await status.update(task_id, error_msg=retry_info)
+                await status.update(task_id, TaskState.WAITING, error_msg=retry_info)
 
                 # Use debug level for Rich trackers, warning/info for console trackers
                 use_debug_level = status.suppress_logs
@@ -651,22 +766,46 @@ async def _execute_with_retry(
             else:
                 log.warning(rate_limit_msg)
                 log.info(exception_msg)
+
+            # Sleep during backoff while holding semaphore slots
             await asyncio.sleep(backoff_time)
 
+            # Set state back to running before next attempt
+            if status and task_id:
+                await status.update(task_id, TaskState.RUNNING)
+
         try:
-            # Acquire global limits first, then bucket-specific limits if present
-            async with global_semaphore, global_rate_limiter:
-                if bucket_semaphore is not None and bucket_rate_limiter is not None:
-                    async with bucket_semaphore, bucket_rate_limiter:
-                        # Mark task as started now that we've passed rate limiting
-                        if status and task_id is not None and attempt == 0:
-                            await status.start(task_id)
-                        return await executor()
+            # Execute task to check for potential rate limit bypass
+            raw_result = await executor()
+
+            # Check if result indicates limits should be disabled (e.g., cache hit)
+            if isinstance(raw_result, TaskResult):
+                if raw_result.disable_limits:
+                    # Bypass rate limiting and return immediately
+                    return cast(T, raw_result.value)
                 else:
-                    # Mark task as started now that we've passed rate limiting
-                    if status and task_id is not None and attempt == 0:
-                        await status.start(task_id)
-                    return await executor()
+                    # Wrapped but limits enabled: extract value for rate limiting
+                    result_value = cast(T, raw_result.value)
+            else:
+                # Unwrapped result: apply normal rate limiting
+                result_value = cast(T, raw_result)
+
+            # Apply rate limiting for non-bypassed results
+            if global_rate_limiter is not None:
+                async with global_rate_limiter:
+                    if bucket_rate_limiter is not None:
+                        async with bucket_rate_limiter:
+                            return result_value
+                    else:
+                        return result_value
+            else:
+                # No global rate limiter, check bucket rate limiter
+                if bucket_rate_limiter is not None:
+                    async with bucket_rate_limiter:
+                        return result_value
+                else:
+                    # No rate limiting at all
+                    return result_value
 
         except Exception as e:
             last_exception = e  # Always store the exception
@@ -674,10 +813,10 @@ async def _execute_with_retry(
             if attempt == retry_settings.max_task_retries:
                 # Final attempt failed
                 if retry_settings.max_task_retries == 0:
-                    # No retries configured - raise original exception directly
+                    # No retries configured: raise original exception directly
                     raise
                 else:
-                    # Retries were attempted but exhausted - wrap with context
+                    # Retries were attempted but exhausted: wrap with context
                     total_time = time.time() - start_time
                     log.error(
                         f"Max task retries ({retry_settings.max_task_retries}) exhausted after {total_time:.1f}s. "
@@ -685,19 +824,18 @@ async def _execute_with_retry(
                     )
                     raise RetryExhaustedException(e, retry_settings.max_task_retries, total_time)
 
-            # Check if this is a retriable exception using our HTTP-aware function
-            if is_retriable_for_task(e):
-                # Continue to next retry attempt (global limits will be checked at top of loop)
+            # Check if this is a retriable exception using the centralized logic
+            if retry_settings.should_retry(e):
+                # Continue to next retry attempt (semaphores remain held for backoff)
                 continue
             else:
                 # Non-retriable exception, log and re-raise immediately
                 status_code = extract_http_status_code(e)
                 status_info = f" (HTTP {status_code})" if status_code else ""
 
-                log.warning(
-                    "Non-retriable exception%s (not retrying): %s", status_info, e, exc_info=True
-                )
+                log.warning("Non-retriable exception%s (not retrying): %s", status_info, e)
+                log.debug("Exception traceback:", exc_info=True)
                 raise
 
     # This should never be reached, but satisfy type checker
-    raise RuntimeError("Unexpected code path in _execute_with_retry")
+    raise RuntimeError("Unexpected code path in _execute_with_retry_inner")

kash/utils/api_utils/http_utils.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+
+def extract_http_status_code(exception: Exception) -> int | None:
+    """
+    Extract HTTP status code from various exception types.
+
+    Args:
+        exception: The exception to extract status code from
+
+    Returns:
+        HTTP status code or None if not found
+    """
+    # Check for httpx.HTTPStatusError and requests.HTTPError
+    if hasattr(exception, "response"):
+        response = getattr(exception, "response", None)
+        if response and hasattr(response, "status_code"):
+            return getattr(response, "status_code", None)
+
+    # Check for aiohttp errors
+    if hasattr(exception, "status"):
+        return getattr(exception, "status", None)
+
+    # Parse from exception message as fallback
+    exception_str = str(exception)
+
+    # Try to find status code patterns in the message
+    import re
+
+    # Pattern for "403 Forbidden", "HTTP 429", etc.
+    status_patterns = [
+        r"\b(\d{3})\s+(?:Forbidden|Unauthorized|Not Found|Too Many Requests|Internal Server Error|Bad Gateway|Service Unavailable|Gateway Timeout)\b",
+        r"\bHTTP\s+(\d{3})\b",
+        r"\b(\d{3})\s+error\b",
+        r"status\s*(?:code)?:\s*(\d{3})\b",
+    ]
+
+    for pattern in status_patterns:
+        match = re.search(pattern, exception_str, re.IGNORECASE)
+        if match:
+            try:
+                return int(match.group(1))
+            except (ValueError, IndexError):
+                continue
+
+    return None