kash-shell 0.3.22__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
 
@@ -16,6 +17,7 @@ from kash.utils.api_utils.api_retries import (
     RetryExhaustedException,
     RetrySettings,
     calculate_backoff,
+    extract_http_status_code,
 )
 from kash.utils.api_utils.progress_protocol import Labeler, ProgressTracker, TaskState
 
@@ -25,28 +27,82 @@ log = logging.getLogger(__name__)
 
 
 @dataclass(frozen=True)
-class FuncTask:
+class Limit:
+    """
+    Rate limiting configuration with max RPS and max concurrency.
+    """
+
+    rps: float
+    concurrency: int
+
+
+DEFAULT_CANCEL_TIMEOUT: float = 1.0
+
+
+@dataclass(frozen=True)
+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]]
 SyncLabeler: TypeAlias = Labeler[SyncSpec[T]]
 
-DEFAULT_MAX_CONCURRENT: int = 5
-DEFAULT_MAX_RPS: float = 5.0
-DEFAULT_CANCEL_TIMEOUT: float = 1.0
+
+def _get_bucket_limits(
+    bucket: str,
+    bucket_semaphores: dict[str, asyncio.Semaphore],
+    bucket_rate_limiters: dict[str, AsyncLimiter],
+) -> tuple[asyncio.Semaphore | None, AsyncLimiter | None]:
+    """
+    Get bucket-specific limits with fallback to "*" wildcard.
+
+    Checks for exact bucket match first, then falls back to "*" if available.
+    Returns (None, None) if neither exact match nor "*" fallback exists.
+    """
+    # Try exact bucket match first
+    bucket_semaphore = bucket_semaphores.get(bucket)
+    bucket_rate_limiter = bucket_rate_limiters.get(bucket)
+
+    if bucket_semaphore is not None and bucket_rate_limiter is not None:
+        return bucket_semaphore, bucket_rate_limiter
+
+    # Fall back to "*" wildcard if available
+    bucket_semaphore = bucket_semaphores.get("*")
+    bucket_rate_limiter = bucket_rate_limiters.get("*")
+
+    return bucket_semaphore, bucket_rate_limiter
 
 
 class RetryCounter:
@@ -75,8 +131,8 @@ class RetryCounter:
 @overload
 async def gather_limited_async(
     *coro_specs: CoroSpec[T],
-    max_concurrent: int = DEFAULT_MAX_CONCURRENT,
-    max_rps: float = DEFAULT_MAX_RPS,
+    limit: Limit | None,
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = False,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -87,8 +143,8 @@ async def gather_limited_async(
 @overload
 async def gather_limited_async(
     *coro_specs: CoroSpec[T],
-    max_concurrent: int = DEFAULT_MAX_CONCURRENT,
-    max_rps: float = DEFAULT_MAX_RPS,
+    limit: Limit | None,
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -98,27 +154,42 @@ async def gather_limited_async(
 
 async def gather_limited_async(
     *coro_specs: CoroSpec[T],
-    max_concurrent: int = DEFAULT_MAX_CONCURRENT,
-    max_rps: float = DEFAULT_MAX_RPS,
-    return_exceptions: bool = False,
+    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,
     status: ProgressTracker | None = None,
     labeler: CoroLabeler[T] | None = None,
 ) -> list[T] | list[T | BaseException]:
     """
-    Rate-limited version of `asyncio.gather()` with retry logic and optional progress display.
+    Rate-limited version of `asyncio.gather()` with HTTP-aware retry logic and optional progress display.
     Uses the aiolimiter leaky-bucket algorithm with exponential backoff on failures.
 
     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 (prevents cascade failures)
 
+    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
+
+    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
@@ -126,7 +197,7 @@ async def gather_limited_async(
         from kash.utils.rich_custom.task_status import TaskStatus
 
         async with TaskStatus() as status:
-            await gather_limited(
+            await gather_limited_async(
                 lambda: fetch_url("http://example.com"),
                 lambda: process_data(data),
                 status=status,
@@ -135,18 +206,47 @@ async def gather_limited_async(
             )
 
         # Without progress display:
-        await gather_limited(
+        await gather_limited_async(
             lambda: fetch_url("http://example.com"),
             lambda: process_data(data),
             retry_settings=RetrySettings(max_task_retries=3, max_total_retries=25)
         )
 
+        # With bucket-specific limits and "*" fallback:
+        await gather_limited_async(
+            FuncTask(fetch_api, ("data1",), bucket="api1"),
+            FuncTask(fetch_api, ("data2",), bucket="api2"),
+            FuncTask(fetch_api, ("data3",), bucket="api3"),
+            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
-        max_concurrent: Maximum number of concurrent executions
-        max_rps: Maximum requests per second
+        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
@@ -159,9 +259,9 @@ async def gather_limited_async(
         ValueError: If coroutines are passed when retries are enabled
     """
     log.info(
-        "Executing with concurrency %s at %s rps, %s",
-        max_concurrent,
-        max_rps,
+        "Executing with global limits: concurrency %s at %s rps, %s",
+        limit.concurrency if limit else "None",
+        limit.rps if limit else "None",
         retry_settings,
     )
     if not coro_specs:
@@ -179,8 +279,18 @@ async def gather_limited_async(
                     f"lambda: your_async_func(args) instead of your_async_func(args)"
                 )
 
-    semaphore = asyncio.Semaphore(max_concurrent)
-    rate_limiter = AsyncLimiter(max_rps, 1.0)
+    # Global limits (apply to all tasks regardless of bucket)
+    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] = {}
+    bucket_rate_limiters: dict[str, AsyncLimiter] = {}
+
+    if bucket_limits:
+        for bucket_name, limit in bucket_limits.items():
+            bucket_semaphores[bucket_name] = asyncio.Semaphore(limit.concurrency)
+            bucket_rate_limiters[bucket_name] = AsyncLimiter(limit.rps, 1.0)
 
     # Global retry counter (shared across all tasks)
     global_retry_counter = RetryCounter(retry_settings.max_total_retries)
@@ -190,6 +300,16 @@ async def gather_limited_async(
         label = labeler(i, coro_spec) if labeler else f"task:{i}"
         task_id = await status.add(label) if status else None
 
+        # Determine bucket and get appropriate limits
+        bucket = "default"
+        if isinstance(coro_spec, FuncTask):
+            bucket = coro_spec.bucket
+
+        # Get bucket-specific limits if available
+        bucket_semaphore, bucket_rate_limiter = _get_bucket_limits(
+            bucket, bucket_semaphores, bucket_rate_limiters
+        )
+
         async def executor() -> T:
             # Create a fresh coroutine for each attempt
             if isinstance(coro_spec, FuncTask):
@@ -198,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
 
@@ -206,8 +326,10 @@ async def gather_limited_async(
             result = await _execute_with_retry(
                 executor,
                 retry_settings,
-                semaphore,
-                rate_limiter,
+                global_semaphore,
+                global_rate_limiter,
+                bucket_semaphore,
+                bucket_rate_limiter,
                 global_retry_counter,
                 status,
                 task_id,
@@ -234,8 +356,8 @@ async def gather_limited_async(
 @overload
 async def gather_limited_sync(
     *sync_specs: SyncSpec[T],
-    max_concurrent: int = DEFAULT_MAX_CONCURRENT,
-    max_rps: float = DEFAULT_MAX_RPS,
+    limit: Limit | None,
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = False,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -248,8 +370,8 @@ async def gather_limited_sync(
 @overload
 async def gather_limited_sync(
     *sync_specs: SyncSpec[T],
-    max_concurrent: int = DEFAULT_MAX_CONCURRENT,
-    max_rps: float = DEFAULT_MAX_RPS,
+    limit: Limit | None,
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -261,9 +383,9 @@ async def gather_limited_sync(
 
 async def gather_limited_sync(
     *sync_specs: SyncSpec[T],
-    max_concurrent: int = DEFAULT_MAX_CONCURRENT,
-    max_rps: float = DEFAULT_MAX_RPS,
-    return_exceptions: bool = False,
+    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,
     status: ProgressTracker | None = None,
     labeler: SyncLabeler[T] | None = None,
@@ -271,54 +393,74 @@ 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 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:
+
+    **Function Execution**: Runs sync functions via `asyncio.to_thread()` instead of executing
+    coroutines directly. Validates that callables don't accidentally return coroutines.
 
-    Supports cooperative cancellation and graceful thread termination on interruption.
+    **Cancellation Support**: Provides cooperative cancellation for long-running sync functions
+    through optional `cancel_event` and graceful thread termination.
+
+    **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
-        max_concurrent: Maximum number of concurrent executions
-        max_rps: Maximum requests per second
-        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 cooperative cancellation
+        # Basic usage with sync functions
         results = await gather_limited_sync(
             lambda: some_sync_function(arg1),
             lambda: another_sync_function(arg2),
-            max_concurrent=3,
-            max_rps=2.0,
+            limit=Limit(rps=2.0, concurrency=3),
             retry_settings=RetrySettings(max_task_retries=3, max_total_retries=25)
         )
 
-        # With cooperative cancellation
-        cancel_event = threading.Event()
+        # With bucket-specific limits and "*" fallback
         results = await gather_limited_sync(
-            lambda: cancellable_sync_function(cancel_event, arg1),
-            lambda: another_cancellable_function(cancel_event, arg2),
-            cancel_event=cancel_event,
-            cancel_timeout=5.0,
+            FuncTask(fetch_from_api, ("data1",), bucket="api1"),
+            FuncTask(fetch_from_api, ("data2",), bucket="api2"),
+            FuncTask(fetch_from_api, ("data3",), bucket="api3"),
+            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 concurrency %s at %s rps, %s",
-        max_concurrent,
-        max_rps,
+        "Executing with global limits: concurrency %s at %s rps, %s",
+        limit.concurrency if limit else "None",
+        limit.rps if limit else "None",
         retry_settings,
     )
     if not sync_specs:
@@ -326,8 +468,18 @@ async def gather_limited_sync(
 
     retry_settings = retry_settings or NO_RETRIES
 
-    semaphore = asyncio.Semaphore(max_concurrent)
-    rate_limiter = AsyncLimiter(max_rps, 1.0)
+    # Global limits (apply to all tasks regardless of bucket)
+    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] = {}
+    bucket_rate_limiters: dict[str, AsyncLimiter] = {}
+
+    if bucket_limits:
+        for bucket_name, limit in bucket_limits.items():
+            bucket_semaphores[bucket_name] = asyncio.Semaphore(limit.concurrency)
+            bucket_rate_limiters[bucket_name] = AsyncLimiter(limit.rps, 1.0)
 
     # Global retry counter (shared across all tasks)
     global_retry_counter = RetryCounter(retry_settings.max_total_retries)
@@ -337,6 +489,16 @@ async def gather_limited_sync(
         label = labeler(i, sync_spec) if labeler else f"task:{i}"
         task_id = await status.add(label) if status else None
 
+        # Determine bucket and get appropriate limits
+        bucket = "default"
+        if isinstance(sync_spec, FuncTask):
+            bucket = sync_spec.bucket
+
+        # Get bucket-specific limits if available
+        bucket_semaphore, bucket_rate_limiter = _get_bucket_limits(
+            bucket, bucket_semaphores, bucket_rate_limiters
+        )
+
         async def executor() -> T:
             # Call sync function via asyncio.to_thread, handling retry at this level
             if isinstance(sync_spec, FuncTask):
@@ -361,8 +523,10 @@ async def gather_limited_sync(
             result = await _execute_with_retry(
                 executor,
                 retry_settings,
-                semaphore,
-                rate_limiter,
+                global_semaphore,
+                global_rate_limiter,
+                bucket_semaphore,
+                bucket_rate_limiter,
                 global_retry_counter,
                 status,
                 task_id,
@@ -434,7 +598,8 @@ async def _gather_with_interrupt_handling(
         if cancelled_count > 0:
             try:
                 await asyncio.wait_for(
-                    asyncio.gather(*async_tasks, return_exceptions=True), timeout=cancel_timeout
+                    asyncio.gather(*async_tasks, return_exceptions=True),
+                    timeout=cancel_timeout,
                 )
             except (TimeoutError, asyncio.CancelledError):
                 log.warning("Some tasks did not cancel within timeout")
@@ -462,20 +627,95 @@ async def _gather_with_interrupt_handling(
 async def _execute_with_retry(
     executor: Callable[[], Coroutine[None, None, T]],
     retry_settings: RetrySettings,
-    semaphore: asyncio.Semaphore,
-    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.
+
+    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,
+            )
+
+
+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)
 
     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
-        if attempt > 0 and last_exception is not None:
+        # 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():
                 log.error(
@@ -493,13 +733,13 @@ async def _execute_with_retry(
             )
 
             # Record retry in status display and log appropriately
-            if status and task_id is not None:
+            if status and task_id:
                 # Include retry attempt info and backoff time in the status display
                 retry_info = (
                     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
@@ -508,8 +748,11 @@ async def _execute_with_retry(
                 use_debug_level = False
 
             # Log retry information at appropriate level
+            status_code = extract_http_status_code(last_exception)
+            status_info = f" (HTTP {status_code})" if status_code else ""
+
             rate_limit_msg = (
-                f"Rate limit hit (attempt {attempt}/{retry_settings.max_task_retries} "
+                f"Rate limit hit{status_info} (attempt {attempt}/{retry_settings.max_task_retries} "
                 f"{global_retry_counter.count}/{global_retry_counter.max_total_retries or '∞'} total) "
                 f"backing off for {backoff_time:.2f}s"
             )
@@ -523,25 +766,57 @@ 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 semaphore and rate limiter right before making the call
-            async with semaphore, 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:
+                    # 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
 
             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. "
@@ -549,439 +824,18 @@ async def _execute_with_retry(
                     )
                     raise RetryExhaustedException(e, retry_settings.max_task_retries, total_time)
 
-            # Check if this is a retriable exception
-            if retry_settings.is_retriable(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
-                log.warning("Non-retriable exception (not retrying): %s", e, exc_info=True)
+                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)
+                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")
-
-
-## Tests
-
-
-def test_gather_limited_sync():
-    """Test gather_limited_sync with sync functions."""
-    import asyncio
-    import time
-
-    async def run_test():
-        def sync_func(value: int) -> int:
-            """Simple sync function for testing."""
-            time.sleep(0.1)  # Simulate some work
-            return value * 2
-
-        # Test basic functionality
-        results = await gather_limited_sync(
-            lambda: sync_func(1),
-            lambda: sync_func(2),
-            lambda: sync_func(3),
-            max_concurrent=2,
-            max_rps=10.0,
-            retry_settings=NO_RETRIES,
-        )
-
-        assert results == [2, 4, 6]
-
-    # Run the async test
-    asyncio.run(run_test())
-
-
-def test_gather_limited_sync_with_retries():
-    """Test that sync functions can be retried on retriable exceptions."""
-    import asyncio
-
-    async def run_test():
-        call_count = 0
-
-        def flaky_sync_func() -> str:
-            """Sync function that fails first time, succeeds second time."""
-            nonlocal call_count
-            call_count += 1
-            if call_count == 1:
-                raise Exception("Rate limit exceeded")  # Retriable
-            return "success"
-
-        # Should succeed after retry
-        results = await gather_limited_sync(
-            lambda: flaky_sync_func(),
-            retry_settings=RetrySettings(
-                max_task_retries=2,
-                initial_backoff=0.1,
-                max_backoff=1.0,
-                backoff_factor=2.0,
-            ),
-        )
-
-        assert results == ["success"]
-        assert call_count == 2  # Called twice (failed once, succeeded once)
-
-    # Run the async test
-    asyncio.run(run_test())
-
-
-def test_gather_limited_async_basic():
-    """Test gather_limited with async functions using callables."""
-    import asyncio
-
-    async def run_test():
-        async def async_func(value: int) -> int:
-            """Simple async function for testing."""
-            await asyncio.sleep(0.05)  # Simulate async work
-            return value * 3
-
-        # Test with callables (recommended pattern)
-        results = await gather_limited_async(
-            lambda: async_func(1),
-            lambda: async_func(2),
-            lambda: async_func(3),
-            max_concurrent=2,
-            max_rps=10.0,
-            retry_settings=NO_RETRIES,
-        )
-
-        assert results == [3, 6, 9]
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_direct_coroutines():
-    """Test gather_limited with direct coroutines when retries disabled."""
-    import asyncio
-
-    async def run_test():
-        async def async_func(value: int) -> int:
-            await asyncio.sleep(0.05)
-            return value * 4
-
-        # Test with direct coroutines (only works when retries disabled)
-        results = await gather_limited_async(
-            async_func(1),
-            async_func(2),
-            async_func(3),
-            retry_settings=NO_RETRIES,  # Required for direct coroutines
-        )
-
-        assert results == [4, 8, 12]
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_coroutine_retry_validation():
-    """Test that passing coroutines with retries enabled raises ValueError."""
-    import asyncio
-
-    async def run_test():
-        async def async_func(value: int) -> int:
-            return value
-
-        coro = async_func(1)  # Direct coroutine
-
-        # Should raise ValueError when trying to use coroutines with retries
-        try:
-            await gather_limited_async(
-                coro,  # Direct coroutine
-                lambda: async_func(2),  # Callable
-                retry_settings=RetrySettings(
-                    max_task_retries=1,
-                    initial_backoff=0.1,
-                    max_backoff=1.0,
-                    backoff_factor=2.0,
-                ),
-            )
-            raise AssertionError("Expected ValueError")
-        except ValueError as e:
-            coro.close()  # Close the unused coroutine to prevent RuntimeWarning
-            assert "position 0" in str(e)
-            assert "cannot be retried" in str(e)
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_async_with_retries():
-    """Test that async functions can be retried when using callables."""
-    import asyncio
-
-    async def run_test():
-        call_count = 0
-
-        async def flaky_async_func() -> str:
-            """Async function that fails first time, succeeds second time."""
-            nonlocal call_count
-            call_count += 1
-            if call_count == 1:
-                raise Exception("Rate limit exceeded")  # Retriable
-            return "async_success"
-
-        # Should succeed after retry using callable
-        results = await gather_limited_async(
-            lambda: flaky_async_func(),
-            retry_settings=RetrySettings(
-                max_task_retries=2,
-                initial_backoff=0.1,
-                max_backoff=1.0,
-                backoff_factor=2.0,
-            ),
-        )
-
-        assert results == ["async_success"]
-        assert call_count == 2  # Called twice (failed once, succeeded once)
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_sync_coroutine_validation():
-    """Test that passing async function callables to sync version raises ValueError."""
-    import asyncio
-
-    async def run_test():
-        async def async_func(value: int) -> int:
-            return value
-
-        # Should raise ValueError when trying to use async functions in sync version
-        try:
-            await gather_limited_sync(
-                lambda: async_func(1),  # Returns coroutine - should be rejected
-                retry_settings=NO_RETRIES,
-            )
-            raise AssertionError("Expected ValueError")
-        except ValueError as e:
-            assert "returned a coroutine" in str(e)
-            assert "gather_limited_sync() is for synchronous functions only" in str(e)
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_retry_exhaustion():
-    """Test that retry exhaustion produces clear error messages."""
-    import asyncio
-
-    async def run_test():
-        call_count = 0
-
-        def always_fails() -> str:
-            """Function that always raises retriable exceptions."""
-            nonlocal call_count
-            call_count += 1
-            raise Exception("Rate limit exceeded")  # Always retriable
-
-        # Should exhaust retries and raise RetryExhaustedException
-        try:
-            await gather_limited_sync(
-                lambda: always_fails(),
-                retry_settings=RetrySettings(
-                    max_task_retries=2,
-                    initial_backoff=0.01,
-                    max_backoff=0.1,
-                    backoff_factor=2.0,
-                ),
-            )
-            raise AssertionError("Expected RetryExhaustedException")
-        except RetryExhaustedException as e:
-            assert "Max retries (2) exhausted" in str(e)
-            assert "Rate limit exceeded" in str(e)
-            assert isinstance(e.original_exception, Exception)
-            assert call_count == 3  # Initial attempt + 2 retries
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_return_exceptions():
-    """Test return_exceptions=True behavior for both functions."""
-    import asyncio
-
-    async def run_test():
-        def failing_sync() -> str:
-            raise ValueError("sync error")
-
-        async def failing_async() -> str:
-            raise ValueError("async error")
-
-        # Test sync version with exceptions returned
-        sync_results = await gather_limited_sync(
-            lambda: "success",
-            lambda: failing_sync(),
-            return_exceptions=True,
-            retry_settings=NO_RETRIES,
-        )
-
-        assert len(sync_results) == 2
-        assert sync_results[0] == "success"
-        assert isinstance(sync_results[1], ValueError)
-        assert str(sync_results[1]) == "sync error"
-
-        async def success_async() -> str:
-            return "async_success"
-
-        # Test async version with exceptions returned
-        async_results = await gather_limited_async(
-            lambda: success_async(),
-            lambda: failing_async(),
-            return_exceptions=True,
-            retry_settings=NO_RETRIES,
-        )
-
-        assert len(async_results) == 2
-        assert async_results[0] == "async_success"
-        assert isinstance(async_results[1], ValueError)
-        assert str(async_results[1]) == "async error"
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_global_retry_limit():
-    """Test that global retry limits are enforced across all tasks."""
-    import asyncio
-
-    async def run_test():
-        retry_counts = {"task1": 0, "task2": 0}
-
-        def flaky_task(task_name: str) -> str:
-            """Tasks that always fail but track retry counts."""
-            retry_counts[task_name] += 1
-            raise Exception(f"Rate limit exceeded in {task_name}")
-
-        # Test with very low global retry limit
-        try:
-            await gather_limited_sync(
-                lambda: flaky_task("task1"),
-                lambda: flaky_task("task2"),
-                retry_settings=RetrySettings(
-                    max_task_retries=5,  # Each task could retry up to 5 times
-                    max_total_retries=3,  # But only 3 total retries across all tasks
-                    initial_backoff=0.01,
-                    max_backoff=0.1,
-                    backoff_factor=2.0,
-                ),
-                return_exceptions=True,
-            )
-        except Exception:
-            pass  # Expected to fail due to rate limits
-
-        # Verify that total retries across both tasks doesn't exceed global limit
-        total_retries = (retry_counts["task1"] - 1) + (
-            retry_counts["task2"] - 1
-        )  # -1 for initial attempts
-        assert total_retries <= 3, f"Total retries {total_retries} exceeded global limit of 3"
-
-        # Verify that both tasks were attempted at least once
-        assert retry_counts["task1"] >= 1
-        assert retry_counts["task2"] >= 1
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_funcspec_format():
-    """Test gather_limited with FuncSpec format and custom labeler accessing args."""
-    import asyncio
-
-    async def run_test():
-        def sync_func(name: str, value: int, multiplier: int = 2) -> str:
-            """Sync function that takes args and kwargs."""
-            return f"{name}: {value * multiplier}"
-
-        async def async_func(name: str, value: int, multiplier: int = 2) -> str:
-            """Async function that takes args and kwargs."""
-            await asyncio.sleep(0.01)
-            return f"{name}: {value * multiplier}"
-
-        captured_labels = []
-
-        def custom_labeler(i: int, spec: Any) -> str:
-            if isinstance(spec, FuncTask):
-                # Extract meaningful info from args for labeling
-                if spec.args and len(spec.args) > 0:
-                    label = f"Processing {spec.args[0]}"
-                else:
-                    label = f"Task {i}"
-            else:
-                label = f"Task {i}"
-            captured_labels.append(label)
-            return label
-
-        # Test sync version with FuncSpec format and custom labeler
-        sync_results = await gather_limited_sync(
-            FuncTask(sync_func, ("user1", 100), {"multiplier": 3}),  # user1: 300
-            FuncTask(sync_func, ("user2", 50)),  # user2: 100 (default multiplier)
-            labeler=custom_labeler,
-            retry_settings=NO_RETRIES,
-        )
-
-        assert sync_results == ["user1: 300", "user2: 100"]
-        assert captured_labels == ["Processing user1", "Processing user2"]
-
-        # Reset labels for async test
-        captured_labels.clear()
-
-        # Test async version with FuncSpec format and custom labeler
-        async_results = await gather_limited_async(
-            FuncTask(async_func, ("api_call", 10), {"multiplier": 4}),  # api_call: 40
-            FuncTask(async_func, ("data_fetch", 5)),  # data_fetch: 10 (default multiplier)
-            labeler=custom_labeler,
-            retry_settings=NO_RETRIES,
-        )
-
-        assert async_results == ["api_call: 40", "data_fetch: 10"]
-        assert captured_labels == ["Processing api_call", "Processing data_fetch"]
-
-    asyncio.run(run_test())
-
-
-def test_gather_limited_sync_cooperative_cancellation():
-    """Test gather_limited_sync with cooperative cancellation via threading.Event."""
-    import asyncio
-    import time
-
-    async def run_test():
-        cancel_event = threading.Event()
-        call_counts = {"task1": 0, "task2": 0}
-
-        def cancellable_sync_func(task_name: str, work_duration: float) -> str:
-            """Sync function that checks cancellation event periodically."""
-            call_counts[task_name] += 1
-            start_time = time.time()
-
-            while time.time() - start_time < work_duration:
-                if cancel_event.is_set():
-                    return f"{task_name}: cancelled"
-                time.sleep(0.01)  # Small sleep to allow cancellation check
-
-            return f"{task_name}: completed"
-
-        # Test cooperative cancellation - tasks should respect the cancel_event
-        results = await gather_limited_sync(
-            lambda: cancellable_sync_func("task1", 0.1),  # Short duration
-            lambda: cancellable_sync_func("task2", 0.1),  # Short duration
-            cancel_event=cancel_event,
-            cancel_timeout=1.0,
-            retry_settings=NO_RETRIES,
-        )
-
-        # Should complete normally since cancel_event wasn't set
-        assert results == ["task1: completed", "task2: completed"]
-        assert call_counts["task1"] == 1
-        assert call_counts["task2"] == 1
-
-        # Test that cancel_event can be used independently
-        cancel_event.set()  # Set cancellation signal
-
-        results2 = await gather_limited_sync(
-            lambda: cancellable_sync_func("task1", 1.0),  # Would take long if not cancelled
-            lambda: cancellable_sync_func("task2", 1.0),  # Would take long if not cancelled
-            cancel_event=cancel_event,
-            cancel_timeout=1.0,
-            retry_settings=NO_RETRIES,
-        )
-
-        # Should be cancelled quickly since cancel_event is already set
-        assert results2 == ["task1: cancelled", "task2: cancelled"]
-        # Call counts should increment
-        assert call_counts["task1"] == 2
-        assert call_counts["task2"] == 2
-
-    asyncio.run(run_test())
+    raise RuntimeError("Unexpected code path in _execute_with_retry_inner")