keycycle 0.3.2__tar.gz → 0.4.0__tar.gz

{keycycle-0.3.2 → keycycle-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keycycle
-Version: 0.3.2
+Version: 0.4.0
 Summary: A thread-safe key rotation and rate limiting manager for API keys.
 Author-email: William Jin <jinwilliam.jin@gmail.com>
 Classifier: Programming Language :: Python :: 3

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/adapters/generic_adapter.py

@@ -35,7 +35,7 @@ from ..config.constants import (
 from ..core.utils import (
     is_rate_limit_error,
     is_temporary_rate_limit_error,
-    is_payment_required_error,
+    is_dead_key_error,
     get_key_suffix,
 )
 from ..core.backoff import ExponentialBackoff, BackoffConfig
@@ -305,9 +305,9 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
         limits = self.limit_resolver(model_id, None)
 
         for attempt in range(self.config.max_retries + 1):
-            key_usage = self.manager.get_key(model_id, limits, self.config.estimated_tokens)
+            key_usage = self.manager.acquire_key(model_id, limits, self.config.estimated_tokens)
             if not key_usage:
-                raise RuntimeError(f"No available keys for {model_id}")
+                raise RuntimeError(self.manager.unavailable_message(model_id))
 
             # Create backoff for temporary rate limits
             temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -348,11 +348,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
                         time.sleep(delay)
                         continue  # Retry with SAME key
 
-                    # Payment required (quota/credits exhausted) - key is dead
+                    # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
                     # for this process, rotate to next key
-                    if is_payment_required_error(e) and attempt < self.config.max_retries:
+                    if is_dead_key_error(e) and attempt < self.config.max_retries:
                         logger.warning(
-                            "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
+                            "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
                             model_id, get_key_suffix(key_usage.api_key),
                             attempt + 1, self.config.max_retries + 1
                         )
@@ -401,9 +401,9 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
                     final_tokens = chunk_tokens
                 yield chunk
         except Exception as e:
-            if is_payment_required_error(e):
+            if is_dead_key_error(e):
                 logger.warning(
-                    "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
+                    "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
                     model_id, get_key_suffix(key_usage.api_key)
                 )
                 key_usage.mark_dead()
@@ -446,9 +446,9 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
         limits = self.limit_resolver(model_id, None)
 
         for attempt in range(self.config.max_retries + 1):
-            key_usage = self.manager.get_key(model_id, limits, self.config.estimated_tokens)
+            key_usage = await self.manager.acquire_key_async(model_id, limits, self.config.estimated_tokens)
             if not key_usage:
-                raise RuntimeError(f"No available keys for {model_id}")
+                raise RuntimeError(self.manager.unavailable_message(model_id))
 
             # Create backoff for temporary rate limits
             temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -487,11 +487,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
                         await asyncio.sleep(delay)
                         continue  # Retry with SAME key
 
-                    # Payment required (quota/credits exhausted) - key is dead
+                    # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
                     # for this process, rotate to next key
-                    if is_payment_required_error(e) and attempt < self.config.max_retries:
+                    if is_dead_key_error(e) and attempt < self.config.max_retries:
                         logger.warning(
-                            "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
+                            "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
                             model_id, get_key_suffix(key_usage.api_key),
                             attempt + 1, self.config.max_retries + 1
                         )
@@ -540,9 +540,9 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
                     final_tokens = chunk_tokens
                 yield chunk
         except Exception as e:
-            if is_payment_required_error(e):
+            if is_dead_key_error(e):
                 logger.warning(
-                    "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
+                    "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
                     model_id, get_key_suffix(key_usage.api_key)
                 )
                 key_usage.mark_dead()

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/adapters/openai_adapter.py

@@ -14,7 +14,7 @@ from ..config.constants import (
 from ..core.utils import (
     is_rate_limit_error,
     is_temporary_rate_limit_error,
-    is_payment_required_error,
+    is_dead_key_error,
     get_key_suffix,
 )
 from ..core.backoff import ExponentialBackoff, BackoffConfig
@@ -132,9 +132,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
             kwargs['stream_options'] = {"include_usage": True}
 
         for attempt in range(self.max_retries + 1):
-            key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
+            key_usage = self.manager.acquire_key(model_id, limits, self.estimated_tokens)
             if not key_usage:
-                raise RuntimeError(f"No available keys for {model_id}")
+                raise RuntimeError(self.manager.unavailable_message(model_id))
 
             # Create backoff for temporary rate limits
             temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -171,11 +171,11 @@ class RotatingOpenAIClient(BaseRotatingClient):
                         time.sleep(delay)
                         continue  # Retry with SAME key
 
-                    # Payment required (quota/credits exhausted) - key is dead
+                    # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
                     # for this process, rotate to next key
-                    if is_payment_required_error(e) and attempt < self.max_retries:
+                    if is_dead_key_error(e) and attempt < self.max_retries:
                         logger.warning(
-                            "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
+                            "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
                             model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
                         )
                         key_usage.mark_dead()
@@ -217,9 +217,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
                     final_tokens = chunk.usage.total_tokens
                 yield chunk
         except Exception as e:
-            if is_payment_required_error(e):
+            if is_dead_key_error(e):
                 logger.warning(
-                    "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
+                    "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
                     model_id, get_key_suffix(key_usage.api_key)
                 )
                 key_usage.mark_dead()
@@ -266,9 +266,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
             kwargs['stream_options'] = {"include_usage": True}
 
         for attempt in range(self.max_retries + 1):
-            key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
+            key_usage = await self.manager.acquire_key_async(model_id, limits, self.estimated_tokens)
             if not key_usage:
-                raise RuntimeError(f"No available keys for {model_id}")
+                raise RuntimeError(self.manager.unavailable_message(model_id))
 
             # Create backoff for temporary rate limits
             temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -305,11 +305,11 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
                         await asyncio.sleep(delay)
                         continue  # Retry with SAME key
 
-                    # Payment required (quota/credits exhausted) - key is dead
+                    # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
                     # for this process, rotate to next key
-                    if is_payment_required_error(e) and attempt < self.max_retries:
+                    if is_dead_key_error(e) and attempt < self.max_retries:
                         logger.warning(
-                            "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
+                            "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
                             model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
                         )
                         key_usage.mark_dead()
@@ -351,9 +351,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
                     final_tokens = chunk.usage.total_tokens
                 yield chunk
         except Exception as e:
-            if is_payment_required_error(e):
+            if is_dead_key_error(e):
                 logger.warning(
-                    "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
+                    "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
                     model_id, get_key_suffix(key_usage.api_key)
                 )
                 key_usage.mark_dead()

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/config/constants.py

@@ -31,5 +31,12 @@ TEMP_RATE_LIMIT_MULTIPLIER = 2.0
 # Key rotation delay (seconds to wait after rotating to a new key)
 KEY_ROTATION_DELAY_SECONDS = 0.5
 
+# How long acquire_key() polls for a live key to leave its rate window /
+# cooldown before giving up. A single production key at its per-minute cap
+# is the common case (Cohere, 2026-09-07): the right answer is to wait for
+# the window, not to abort the caller's whole run.
+KEY_WAIT_MAX_SECONDS = 120.0
+KEY_WAIT_POLL_SECONDS = 1.0
+
 # Usage logging batch size
 USAGE_LOG_BATCH_SIZE = 50

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/core/__init__.py

@@ -14,6 +14,7 @@ from .utils import (
     get_key_suffix,
     is_rate_limit_error,
     is_auth_error,
+    is_dead_key_error,
     validate_api_key,
 )
 from .backoff import ExponentialBackoff, BackoffConfig
@@ -33,6 +34,7 @@ __all__ = [
     "get_key_suffix",
     "is_rate_limit_error",
     "is_auth_error",
+    "is_dead_key_error",
     "validate_api_key",
     # Backoff
     "ExponentialBackoff",

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/core/utils.py

@@ -241,6 +241,29 @@ def is_auth_error(e: Exception) -> bool:
     return False
 
 
+# A key that answers 401/402/403 is unusable for good (revoked, unpaid,
+# or its project denied) — no cooldown will fix it, so rotate off it.
+DEAD_KEY_STATUS_CODES = AUTH_STATUS_CODES | PAYMENT_REQUIRED_STATUS_CODES
+
+
+def is_dead_key_error(e: Exception) -> bool:
+    """
+    Detect errors that mean THIS KEY is dead, not merely busy: HTTP 401/403
+    (invalid, revoked, or project-denied key — e.g. Gemini's
+    "Your project has been denied access") and 402 (credits exhausted).
+
+    A status_code, when present, is authoritative: 429s are never dead keys
+    whatever their message says. Without one, fall back to the string
+    heuristics of is_payment_required_error / is_auth_error.
+    """
+    status = getattr(e, "status_code", None)
+    if status in DEAD_KEY_STATUS_CODES:
+        return True
+    if status == 429:
+        return False
+    return is_payment_required_error(e) or is_auth_error(e)
+
+
 def validate_api_key(api_key: str) -> bool:
     """
     Basic validation of API key format.

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/key_rotation/rotating_mixin.py

@@ -15,7 +15,7 @@ from ..config.constants import (
 from ..core.utils import (
     is_rate_limit_error,
     is_temporary_rate_limit_error,
-    is_payment_required_error,
+    is_dead_key_error,
     get_key_suffix,
 )
 from ..core.backoff import ExponentialBackoff, BackoffConfig
@@ -171,11 +171,11 @@ class RotatingCredentialsMixin:
                         time.sleep(delay)
                         continue  # Retry with SAME key
 
-                    # Payment required (quota/credits exhausted) - key is dead
+                    # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
                     # for this process, rotate to next key
-                    if is_payment_required_error(e) and attempt < limit:
+                    if is_dead_key_error(e) and attempt < limit:
                         self.logger.warning(
-                            "402 Hit on key %s (Sync) [%s]. Marking key dead and rotating (%d/%d).",
+                            "Dead key (401/402/403) on key %s (Sync) [%s]. Marking key dead and rotating (%d/%d).",
                             get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
                         )
                         key_usage.mark_dead()
@@ -228,11 +228,11 @@ class RotatingCredentialsMixin:
                         await asyncio.sleep(delay)
                         continue  # Retry with SAME key
 
-                    # Payment required (quota/credits exhausted) - key is dead
+                    # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
                     # for this process, rotate to next key
-                    if is_payment_required_error(e) and attempt < limit:
+                    if is_dead_key_error(e) and attempt < limit:
                         self.logger.warning(
-                            "402 Hit on key %s (Async) [%s]. Marking key dead and rotating (%d/%d).",
+                            "Dead key (401/402/403) on key %s (Async) [%s]. Marking key dead and rotating (%d/%d).",
                             get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
                         )
                         key_usage.mark_dead()
@@ -300,9 +300,9 @@ class RotatingCredentialsMixin:
                         time.sleep(delay)
                         continue
 
-                    if is_payment_required_error(e) and attempt < limit:
+                    if is_dead_key_error(e) and attempt < limit:
                         self.logger.warning(
-                            "402 Hit on key %s (Sync Stream) [%s]. Marking key dead and rotating (%d/%d).",
+                            "Dead key (401/402/403) on key %s (Sync Stream) [%s]. Marking key dead and rotating (%d/%d).",
                             get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
                         )
                         key_usage.mark_dead()
@@ -366,9 +366,9 @@ class RotatingCredentialsMixin:
                         await asyncio.sleep(delay)
                         continue
 
-                    if is_payment_required_error(e) and attempt < limit:
+                    if is_dead_key_error(e) and attempt < limit:
                         self.logger.warning(
-                            "402 Hit on key %s (Async Stream) [%s]. Marking key dead and rotating (%d/%d).",
+                            "Dead key (401/402/403) on key %s (Async Stream) [%s]. Marking key dead and rotating (%d/%d).",
                             get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
                         )
                         key_usage.mark_dead()

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle/key_rotation/rotation_manager.py

@@ -1,4 +1,5 @@
 import time
+import asyncio
 import atexit
 import threading
 from threading import Lock, Event
@@ -17,6 +18,8 @@ from ..config.constants import (
     CLEANUP_INTERVAL_SECONDS,
     HISTORY_LOOKBACK_SECONDS,
     DEFAULT_COOLDOWN_SECONDS,
+    KEY_WAIT_MAX_SECONDS,
+    KEY_WAIT_POLL_SECONDS,
 )
 from ..core.utils import get_key_suffix, KeyEntry, normalize_key_entries
 from ..usage.usage_logger import AsyncUsageLogger
@@ -35,8 +38,10 @@ class RotatingKeyManager:
         cooldown_seconds: int = DEFAULT_COOLDOWN_SECONDS,
         limit_resolver: Optional[Callable[[str, Optional[str]], RateLimits]] = None,
         api_key_param: str = "api_key",
+        max_key_wait_seconds: float = KEY_WAIT_MAX_SECONDS,
     ):
         self.provider_name = provider_name
+        self.max_key_wait_seconds = max_key_wait_seconds
         self.logger = logger or default_logger
         self.strategy = strategy
         self.cooldown_seconds = cooldown_seconds
@@ -147,6 +152,64 @@ class RotatingKeyManager:
                     return key
             return None
     
+    def all_keys_dead(self) -> bool:
+        """True when every key has been marked dead (401/402/403) - nothing to wait for."""
+        return bool(self.keys) and all(k.dead for k in self.keys)
+
+    def unavailable_message(self, model_id: str) -> str:
+        dead = sum(1 for k in self.keys if k.dead)
+        if self.all_keys_dead():
+            return f"No usable keys for {model_id}: all {len(self.keys)} keys are dead (401/402/403)"
+        return (
+            f"No available keys for {model_id}: {dead}/{len(self.keys)} dead, the rest stayed "
+            f"rate-limited or cooling down for {self.max_key_wait_seconds:.0f}s"
+        )
+
+    def acquire_key(
+        self,
+        model_id: str,
+        default_limits: RateLimits,
+        estimated_tokens: int = 1000,
+        max_wait_seconds: Optional[float] = None,
+        poll_seconds: float = KEY_WAIT_POLL_SECONDS,
+    ) -> Optional[KeyUsage]:
+        """
+        get_key, but WAIT while every live key is inside a rate window or a
+        cooldown instead of handing back None at once.
+
+        Returns None only when every key is dead or max_wait elapses - so a
+        caller with one production key at its per-minute cap rides the window
+        out rather than aborting a long run.
+        """
+        max_wait = self.max_key_wait_seconds if max_wait_seconds is None else max_wait_seconds
+        deadline = time.monotonic() + max_wait
+        while True:
+            key = self.get_key(model_id, default_limits, estimated_tokens)
+            if key is not None:
+                return key
+            if self.all_keys_dead() or time.monotonic() >= deadline:
+                return None
+            time.sleep(poll_seconds)
+
+    async def acquire_key_async(
+        self,
+        model_id: str,
+        default_limits: RateLimits,
+        estimated_tokens: int = 1000,
+        max_wait_seconds: Optional[float] = None,
+        poll_seconds: float = KEY_WAIT_POLL_SECONDS,
+    ) -> Optional[KeyUsage]:
+        """acquire_key for async callers (non-blocking sleep)."""
+        max_wait = self.max_key_wait_seconds if max_wait_seconds is None else max_wait_seconds
+        deadline = time.monotonic() + max_wait
+        while True:
+            key = self.get_key(model_id, default_limits, estimated_tokens)
+            if key is not None:
+                return key
+            if self.all_keys_dead() or time.monotonic() >= deadline:
+                return None
+            await asyncio.sleep(poll_seconds)
+
     def get_specific_key(self, identifier: Union[int, str], model_id: str, estimated_tokens: int = 1000) -> Optional[KeyUsage]:
         """
         Get a specific key by index or identifier.

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keycycle
-Version: 0.3.2
+Version: 0.4.0
 Summary: A thread-safe key rotation and rate limiting manager for API keys.
 Author-email: William Jin <jinwilliam.jin@gmail.com>
 Classifier: Programming Language :: Python :: 3

{keycycle-0.3.2 → keycycle-0.4.0}/keycycle.egg-info/SOURCES.txt

@@ -42,6 +42,8 @@ tests/test_api_key_pinning.py
 tests/test_api_key_retrieval.py
 tests/test_cerebras_direct.py
 tests/test_custom_key_limits.py
+tests/test_dead_key_rotation.py
+tests/test_key_wait.py
 tests/test_multi_client_wrapper.py
 tests/test_payment_required_rotation.py
 tests/test_rate_limit_logic.py

{keycycle-0.3.2 → keycycle-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "keycycle"
-version = "0.3.2"
+version = "0.4.0"
 authors = [
   { name="William Jin", email="jinwilliam.jin@gmail.com" },
 ]

keycycle-0.4.0/tests/test_dead_key_rotation.py

@@ -0,0 +1,124 @@
+"""
+Unit tests for dead-key (HTTP 401/402/403) rotation.
+
+Extends test_payment_required_rotation.py: 401/403 mean the key itself is
+unusable (revoked, invalid, or its project denied — Gemini answers
+"Your project has been denied access" with a 403), exactly like a 402.
+Such a key must be marked dead and rotated off, never allowed to abort a
+whole run while dozens of healthy keys sit idle (emails_gen pilot,
+2026-09-07).
+"""
+import unittest
+from unittest.mock import MagicMock, patch
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+try:
+    from keycycle.core.utils import is_dead_key_error
+    from keycycle.config.dataclasses import KeyUsage, RateLimits
+    from keycycle.config.enums import RateLimitStrategy
+    from keycycle.key_rotation.rotation_manager import RotatingKeyManager
+    from keycycle.adapters.openai_adapter import RotatingOpenAIClient
+except ImportError:
+    from keycycle.keycycle.core.utils import is_dead_key_error
+    from keycycle.keycycle.config.dataclasses import KeyUsage, RateLimits
+    from keycycle.keycycle.config.enums import RateLimitStrategy
+    from keycycle.keycycle.key_rotation.rotation_manager import RotatingKeyManager
+    from keycycle.keycycle.adapters.openai_adapter import RotatingOpenAIClient
+
+
+class MockAPIError(Exception):
+    def __init__(self, message, status_code=None, body=None):
+        super().__init__(message)
+        self.status_code = status_code
+        self.body = body
+        self.response = body
+
+
+class TestDeadKeyDetection(unittest.TestCase):
+    def test_403_project_denied_is_dead(self):
+        e = MockAPIError(
+            "Error code: 403 - [{'error': {'code': 403, 'message': 'Your project has "
+            "been denied access. Please contact support.', 'status': 'PERMISSION_DENIED'}}]",
+            status_code=403,
+        )
+        self.assertTrue(is_dead_key_error(e))
+
+    def test_401_invalid_key_is_dead(self):
+        self.assertTrue(is_dead_key_error(MockAPIError("Unauthorized", status_code=401)))
+
+    def test_402_still_dead(self):
+        self.assertTrue(is_dead_key_error(MockAPIError("Payment Required", status_code=402)))
+
+    def test_429_never_dead_even_with_scary_message(self):
+        # A rate limit whose body mentions "forbidden" is still just a rate limit.
+        e = MockAPIError("429 rate limit: forbidden to exceed quota", status_code=429)
+        self.assertFalse(is_dead_key_error(e))
+
+    def test_string_fallback_without_status_code(self):
+        self.assertTrue(is_dead_key_error(Exception("Error code: 403 - PERMISSION_DENIED")))
+        self.assertFalse(is_dead_key_error(Exception("Something unrelated exploded")))
+
+
+class TestDeadKeyRotation(unittest.TestCase):
+    """A 403 on key A must mark A dead and let the request succeed on key B
+    (same fixture shape as TestOpenAIAdapter402Rotation)."""
+
+    def setUp(self):
+        self.api_keys = ["sk-key-one-AAAAAAAA", "sk-key-two-BBBBBBBB"]
+        mock_db = MagicMock()
+        mock_db.load_provider_history.return_value = []
+        self.manager = RotatingKeyManager(
+            api_keys=self.api_keys,
+            provider_name="test",
+            strategy=RateLimitStrategy.PER_MODEL,
+            db=mock_db,
+        )
+        limits = RateLimits(100, 6000, 10000, 1000000)
+        self.client = RotatingOpenAIClient(
+            manager=self.manager,
+            limit_resolver=lambda model_id, key_suffix: limits,
+            default_model="test-model",
+            max_retries=5,
+            provider="gemini",
+        )
+
+    def _factory(self, denied_suffixes):
+        def factory(api_key):
+            fake = MagicMock()
+            if api_key[-8:] in denied_suffixes:
+                fake.chat.completions.create.side_effect = MockAPIError(
+                    "Your project has been denied access. Please contact support.",
+                    status_code=403,
+                )
+            else:
+                response = MagicMock()
+                response.usage.total_tokens = 7
+                fake.chat.completions.create.return_value = response
+            return fake
+        return factory
+
+    def test_403_marks_key_dead_and_rotates(self):
+        with patch.object(
+            self.client, "_get_fresh_client", side_effect=self._factory({"AAAAAAAA"})
+        ), patch("keycycle.adapters.openai_adapter.time.sleep"):
+            result = self.client.chat.completions.create(model="test-model")
+
+        self.assertEqual(result.usage.total_tokens, 7)
+        dead = {k.api_key[-8:] for k in self.manager.keys if k.dead}
+        self.assertEqual(dead, {"AAAAAAAA"})
+
+    def test_all_403_raises_without_hanging(self):
+        with patch.object(
+            self.client, "_get_fresh_client",
+            side_effect=self._factory({"AAAAAAAA", "BBBBBBBB"}),
+        ), patch("keycycle.adapters.openai_adapter.time.sleep"):
+            with self.assertRaises(Exception):
+                self.client.chat.completions.create(model="test-model")
+        self.assertTrue(all(k.dead for k in self.manager.keys))
+
+
+if __name__ == "__main__":
+    unittest.main()

keycycle-0.4.0/tests/test_key_wait.py

@@ -0,0 +1,78 @@
+"""
+acquire_key: wait for a rate window / cooldown instead of failing at once.
+
+One production key at its per-minute cap (Cohere, emails_gen pilot
+2026-09-07) used to surface as "No available keys" after the first
+window filled, aborting a 2,000-document run 20 calls in.
+"""
+import time
+import unittest
+from unittest.mock import MagicMock
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+try:
+    from keycycle.config.dataclasses import RateLimits
+    from keycycle.config.enums import RateLimitStrategy
+    from keycycle.key_rotation.rotation_manager import RotatingKeyManager
+except ImportError:
+    from keycycle.keycycle.config.dataclasses import RateLimits
+    from keycycle.keycycle.config.enums import RateLimitStrategy
+    from keycycle.keycycle.key_rotation.rotation_manager import RotatingKeyManager
+
+
+def _manager(keys, cooldown_seconds=1):
+    db = MagicMock()
+    db.load_provider_history.return_value = []
+    return RotatingKeyManager(
+        api_keys=keys,
+        provider_name="test",
+        strategy=RateLimitStrategy.PER_MODEL,
+        db=db,
+        cooldown_seconds=cooldown_seconds,
+    )
+
+
+class TestAcquireKey(unittest.TestCase):
+    def setUp(self):
+        self.limits = RateLimits(100, 6000, 10000, 1000000)
+
+    def test_returns_immediately_when_a_key_is_free(self):
+        manager = _manager(["sk-only-key-AAAAAAAA"])
+        started = time.monotonic()
+        key = manager.acquire_key("m", self.limits, estimated_tokens=10)
+        self.assertIsNotNone(key)
+        self.assertLess(time.monotonic() - started, 0.5)
+
+    def test_waits_out_a_cooldown_then_returns_the_key(self):
+        manager = _manager(["sk-only-key-AAAAAAAA"], cooldown_seconds=1)
+        manager.keys[0].trigger_cooldown()
+        self.assertIsNone(manager.get_key("m", self.limits, 10))  # the old behavior
+
+        started = time.monotonic()
+        key = manager.acquire_key("m", self.limits, 10, max_wait_seconds=5, poll_seconds=0.1)
+        self.assertIsNotNone(key)
+        self.assertGreaterEqual(time.monotonic() - started, 0.5)
+
+    def test_all_dead_returns_none_without_waiting(self):
+        manager = _manager(["sk-a-AAAAAAAA", "sk-b-BBBBBBBB"])
+        for k in manager.keys:
+            k.mark_dead()
+        started = time.monotonic()
+        self.assertIsNone(manager.acquire_key("m", self.limits, 10, max_wait_seconds=5))
+        self.assertLess(time.monotonic() - started, 0.5)
+        self.assertIn("all 2 keys are dead", manager.unavailable_message("m"))
+
+    def test_gives_up_after_max_wait(self):
+        manager = _manager(["sk-only-key-AAAAAAAA"], cooldown_seconds=60)
+        manager.keys[0].trigger_cooldown()
+        started = time.monotonic()
+        self.assertIsNone(manager.acquire_key("m", self.limits, 10, max_wait_seconds=0.3, poll_seconds=0.1))
+        self.assertLess(time.monotonic() - started, 2)
+        self.assertIn("rate-limited or cooling down", manager.unavailable_message("m"))
+
+
+if __name__ == "__main__":
+    unittest.main()