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

kash/utils/api_utils/gather_limited.py

@@ -16,6 +16,8 @@ from kash.utils.api_utils.api_retries import (
     RetryExhaustedException,
     RetrySettings,
     calculate_backoff,
+    extract_http_status_code,
+    is_http_status_retriable,
 )
 from kash.utils.api_utils.progress_protocol import Labeler, ProgressTracker, TaskState
 
@@ -23,17 +25,33 @@ 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:
+    """
+    Rate limiting configuration with max RPS and max concurrency.
+    """
+
+    rps: float = DEFAULT_MAX_RPS
+    concurrency: int = DEFAULT_MAX_CONCURRENT
+
 
 @dataclass(frozen=True)
 class FuncTask:
     """
     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.
     """
 
     func: Callable[..., Any]
     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
@@ -44,9 +62,30 @@ SyncSpec: TypeAlias = Callable[[], T] | FuncTask
 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 +114,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,
+    global_limit: Limit = Limit(),
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = False,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -87,8 +126,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,
+    global_limit: Limit = Limit(),
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -98,21 +137,26 @@ 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,
+    global_limit: Limit = Limit(),
+    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
+
     Can optionally display live progress with retry indicators using TaskStatus.
 
     Accepts:
@@ -126,7 +170,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 +179,31 @@ 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"),
+            global_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
+            }
+        )
+
         ```
 
     Args:
         *coro_specs: Callables or coroutines to execute
-        max_concurrent: Maximum number of concurrent executions
-        max_rps: Maximum requests per second
+        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
@@ -159,9 +216,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",
+        global_limit.concurrency,
+        global_limit.rps,
         retry_settings,
     )
     if not coro_specs:
@@ -179,8 +236,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(global_limit.concurrency)
+    global_rate_limiter = AsyncLimiter(global_limit.rps, 1.0)
+
+    # 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 +257,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):
@@ -206,8 +283,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 +313,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,
+    global_limit: Limit = Limit(),
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = False,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -248,8 +327,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,
+    global_limit: Limit = Limit(),
+    bucket_limits: dict[str, Limit] | None = None,
     return_exceptions: bool = True,
     retry_settings: RetrySettings | None = DEFAULT_RETRIES,
     status: ProgressTracker | None = None,
@@ -261,9 +340,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,
+    global_limit: Limit = Limit(),
+    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,19 +350,25 @@ 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.
+    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.
 
     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
 
+    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
+
     Supports cooperative cancellation and graceful thread termination on interruption.
 
     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
+        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
@@ -296,29 +381,31 @@ async def gather_limited_sync(
 
     Example:
         ```python
-        # Without cooperative cancellation
+        # Without bucket limits (backward compatible)
         results = await gather_limited_sync(
             lambda: some_sync_function(arg1),
             lambda: another_sync_function(arg2),
-            max_concurrent=3,
-            max_rps=2.0,
+            global_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"),
+            global_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
+            }
         )
         ```
     """
     log.info(
-        "Executing with concurrency %s at %s rps, %s",
-        max_concurrent,
-        max_rps,
+        "Executing with global limits: concurrency %s at %s rps, %s",
+        global_limit.concurrency,
+        global_limit.rps,
         retry_settings,
     )
     if not sync_specs:
@@ -326,8 +413,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(global_limit.concurrency)
+    global_rate_limiter = AsyncLimiter(global_limit.rps, 1.0)
+
+    # 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 +434,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 +468,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 +543,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,8 +572,10 @@ 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,
+    global_rate_limiter: AsyncLimiter,
+    bucket_semaphore: asyncio.Semaphore | None,
+    bucket_rate_limiter: AsyncLimiter | None,
     global_retry_counter: RetryCounter,
     status: ProgressTracker | None = None,
     task_id: Any | None = None,
@@ -473,9 +585,22 @@ async def _execute_with_retry(
     start_time = time.time()
     last_exception: Exception | None = None
 
+    # 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)
+
+        # Not an HTTP error, use the main is_retriable result
+        return True
+
     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:
+        if attempt > 0 and last_exception:
             # Try to increment global retry counter
             if not await global_retry_counter.try_increment():
                 log.error(
@@ -493,7 +618,7 @@ 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} "
@@ -508,8 +633,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"
             )
@@ -526,12 +654,20 @@ async def _execute_with_retry(
             await asyncio.sleep(backoff_time)
 
         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()
+            # 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()
+                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()
+
         except Exception as e:
             last_exception = e  # Always store the exception
 
@@ -549,439 +685,19 @@ 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):
+            # 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)
                 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, 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())