keycycle 0.3.2__tar.gz → 0.5.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. {keycycle-0.3.2 → keycycle-0.5.0}/PKG-INFO +1 -1
  2. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/adapters/generic_adapter.py +51 -26
  3. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/adapters/openai_adapter.py +56 -27
  4. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/constants.py +7 -0
  5. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/__init__.py +2 -0
  6. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/utils.py +23 -0
  7. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/key_rotation/rotating_mixin.py +37 -19
  8. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/key_rotation/rotation_manager.py +107 -6
  9. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/legacy_multi_provider_wrapper.py +30 -4
  10. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/multi_client_wrapper.py +39 -3
  11. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/PKG-INFO +1 -1
  12. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/SOURCES.txt +4 -1
  13. {keycycle-0.3.2 → keycycle-0.5.0}/pyproject.toml +1 -1
  14. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_api_key_pinning.py +4 -2
  15. keycycle-0.5.0/tests/test_dead_key_rotation.py +124 -0
  16. keycycle-0.5.0/tests/test_key_wait.py +78 -0
  17. keycycle-0.5.0/tests/test_usage_tracking_toggle.py +503 -0
  18. {keycycle-0.3.2 → keycycle-0.5.0}/MANIFEST.in +0 -0
  19. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/__init__.py +0 -0
  20. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/adapters/__init__.py +0 -0
  21. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/__init__.py +0 -0
  22. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/dataclasses.py +0 -0
  23. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/enums.py +0 -0
  24. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/loader.py +0 -0
  25. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/log_config.py +0 -0
  26. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/cerebras.yaml +0 -0
  27. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/cohere.yaml +0 -0
  28. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/gemini.yaml +0 -0
  29. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/groq.yaml +0 -0
  30. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/mistral.yaml +0 -0
  31. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/moonshot.yaml +0 -0
  32. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/openrouter.yaml +0 -0
  33. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models/openrouter_models.yaml +0 -0
  34. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/config/models.py +0 -0
  35. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/backoff.py +0 -0
  36. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/core/exceptions.py +0 -0
  37. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/key_rotation/__init__.py +0 -0
  38. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/py.typed +0 -0
  39. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/usage/__init__.py +0 -0
  40. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/usage/db_logic.py +0 -0
  41. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/usage/usage_logger.py +0 -0
  42. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle/utils.py +0 -0
  43. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/dependency_links.txt +0 -0
  44. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/requires.txt +0 -0
  45. {keycycle-0.3.2 → keycycle-0.5.0}/keycycle.egg-info/top_level.txt +0 -0
  46. {keycycle-0.3.2 → keycycle-0.5.0}/setup.cfg +0 -0
  47. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_api_key_retrieval.py +0 -0
  48. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_cerebras_direct.py +0 -0
  49. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_custom_key_limits.py +0 -0
  50. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_multi_client_wrapper.py +0 -0
  51. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_payment_required_rotation.py +0 -0
  52. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_rate_limit_logic.py +0 -0
  53. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_specific_key_retrieval.py +0 -0
  54. {keycycle-0.3.2 → keycycle-0.5.0}/tests/test_streams_recording.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: keycycle
3
- Version: 0.3.2
3
+ Version: 0.5.0
4
4
  Summary: A thread-safe key rotation and rate limiting manager for API keys.
5
5
  Author-email: William Jin <jinwilliam.jin@gmail.com>
6
6
  Classifier: Programming Language :: Python :: 3
@@ -35,7 +35,7 @@ from ..config.constants import (
35
35
  from ..core.utils import (
36
36
  is_rate_limit_error,
37
37
  is_temporary_rate_limit_error,
38
- is_payment_required_error,
38
+ is_dead_key_error,
39
39
  get_key_suffix,
40
40
  )
41
41
  from ..core.backoff import ExponentialBackoff, BackoffConfig
@@ -275,13 +275,20 @@ class BaseGenericRotatingClient(Generic[T]):
275
275
 
276
276
  return self.config.client_class(**final_kwargs)
277
277
 
278
- def _record_usage(self, key_usage: KeyUsage, model_id: str, actual_tokens: int) -> None:
278
+ def _record_usage(
279
+ self,
280
+ key_usage: KeyUsage,
281
+ model_id: str,
282
+ actual_tokens: int,
283
+ track_usage: Optional[bool] = None,
284
+ ) -> None:
279
285
  """Record usage for a key."""
280
286
  self.manager.record_usage(
281
287
  key_obj=key_usage,
282
288
  model_id=model_id,
283
289
  actual_tokens=actual_tokens,
284
290
  estimated_tokens=self.config.estimated_tokens,
291
+ track_usage=track_usage,
285
292
  )
286
293
 
287
294
  def _extract_usage(self, response: Any) -> int:
@@ -301,13 +308,16 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
301
308
 
302
309
  def _execute(self, path: List[str], args: tuple, kwargs: dict) -> Any:
303
310
  """Execute a method call with key rotation."""
311
+ # Per-call tracking override - popped before the kwargs reach the real
312
+ # client, which may reject unknown keyword arguments.
313
+ track_usage = kwargs.pop("track_usage", None)
304
314
  model_id = self._get_model_id(kwargs)
305
315
  limits = self.limit_resolver(model_id, None)
306
316
 
307
317
  for attempt in range(self.config.max_retries + 1):
308
- key_usage = self.manager.get_key(model_id, limits, self.config.estimated_tokens)
318
+ key_usage = self.manager.acquire_key(model_id, limits, self.config.estimated_tokens)
309
319
  if not key_usage:
310
- raise RuntimeError(f"No available keys for {model_id}")
320
+ raise RuntimeError(self.manager.unavailable_message(model_id))
311
321
 
312
322
  # Create backoff for temporary rate limits
313
323
  temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -331,9 +341,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
331
341
  if hasattr(result, "__iter__") and not isinstance(result, (str, bytes, dict, list)):
332
342
  # Check if it looks like a generator/iterator
333
343
  if hasattr(result, "__next__") or inspect.isgenerator(result):
334
- return self._wrap_stream(result, key_usage, model_id)
344
+ return self._wrap_stream(result, key_usage, model_id, track_usage)
335
345
 
336
- self._record_usage(key_usage, model_id, self._extract_usage(result))
346
+ self._record_usage(
347
+ key_usage, model_id, self._extract_usage(result), track_usage
348
+ )
337
349
  return result
338
350
 
339
351
  except Exception as e:
@@ -348,11 +360,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
348
360
  time.sleep(delay)
349
361
  continue # Retry with SAME key
350
362
 
351
- # Payment required (quota/credits exhausted) - key is dead
363
+ # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
352
364
  # for this process, rotate to next key
353
- if is_payment_required_error(e) and attempt < self.config.max_retries:
365
+ if is_dead_key_error(e) and attempt < self.config.max_retries:
354
366
  logger.warning(
355
- "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
367
+ "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
356
368
  model_id, get_key_suffix(key_usage.api_key),
357
369
  attempt + 1, self.config.max_retries + 1
358
370
  )
@@ -373,7 +385,7 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
373
385
  time.sleep(KEY_ROTATION_DELAY_SECONDS)
374
386
  break # Break inner loop, continue outer loop with new key
375
387
 
376
- self._record_usage(key_usage, model_id, 0)
388
+ self._record_usage(key_usage, model_id, 0, track_usage)
377
389
  raise
378
390
  else:
379
391
  # Inner loop exhausted without success - continue to next key
@@ -389,7 +401,11 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
389
401
  raise RuntimeError(f"All retry attempts exhausted for {model_id}")
390
402
 
391
403
  def _wrap_stream(
392
- self, generator: Generator, key_usage: KeyUsage, model_id: str
404
+ self,
405
+ generator: Generator,
406
+ key_usage: KeyUsage,
407
+ model_id: str,
408
+ track_usage: Optional[bool] = None,
393
409
  ) -> Generator:
394
410
  """Wrap a streaming response to track usage and handle errors."""
395
411
  final_tokens = 0
@@ -401,9 +417,9 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
401
417
  final_tokens = chunk_tokens
402
418
  yield chunk
403
419
  except Exception as e:
404
- if is_payment_required_error(e):
420
+ if is_dead_key_error(e):
405
421
  logger.warning(
406
- "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
422
+ "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
407
423
  model_id, get_key_suffix(key_usage.api_key)
408
424
  )
409
425
  key_usage.mark_dead()
@@ -417,7 +433,7 @@ class SyncGenericRotatingClient(BaseGenericRotatingClient[T]):
417
433
  self.manager.force_rotate_index()
418
434
  raise
419
435
  finally:
420
- self._record_usage(key_usage, model_id, final_tokens)
436
+ self._record_usage(key_usage, model_id, final_tokens, track_usage)
421
437
 
422
438
 
423
439
  class SyncGenericProxyHelper:
@@ -442,13 +458,16 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
442
458
 
443
459
  async def _execute(self, path: List[str], args: tuple, kwargs: dict) -> Any:
444
460
  """Execute a method call with key rotation (async)."""
461
+ # Per-call tracking override - popped before the kwargs reach the real
462
+ # client, which may reject unknown keyword arguments.
463
+ track_usage = kwargs.pop("track_usage", None)
445
464
  model_id = self._get_model_id(kwargs)
446
465
  limits = self.limit_resolver(model_id, None)
447
466
 
448
467
  for attempt in range(self.config.max_retries + 1):
449
- key_usage = self.manager.get_key(model_id, limits, self.config.estimated_tokens)
468
+ key_usage = await self.manager.acquire_key_async(model_id, limits, self.config.estimated_tokens)
450
469
  if not key_usage:
451
- raise RuntimeError(f"No available keys for {model_id}")
470
+ raise RuntimeError(self.manager.unavailable_message(model_id))
452
471
 
453
472
  # Create backoff for temporary rate limits
454
473
  temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -470,9 +489,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
470
489
 
471
490
  # Handle async streaming responses
472
491
  if hasattr(result, "__aiter__"):
473
- return self._wrap_stream(result, key_usage, model_id)
492
+ return self._wrap_stream(result, key_usage, model_id, track_usage)
474
493
 
475
- self._record_usage(key_usage, model_id, self._extract_usage(result))
494
+ self._record_usage(
495
+ key_usage, model_id, self._extract_usage(result), track_usage
496
+ )
476
497
  return result
477
498
 
478
499
  except Exception as e:
@@ -487,11 +508,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
487
508
  await asyncio.sleep(delay)
488
509
  continue # Retry with SAME key
489
510
 
490
- # Payment required (quota/credits exhausted) - key is dead
511
+ # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
491
512
  # for this process, rotate to next key
492
- if is_payment_required_error(e) and attempt < self.config.max_retries:
513
+ if is_dead_key_error(e) and attempt < self.config.max_retries:
493
514
  logger.warning(
494
- "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
515
+ "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
495
516
  model_id, get_key_suffix(key_usage.api_key),
496
517
  attempt + 1, self.config.max_retries + 1
497
518
  )
@@ -512,7 +533,7 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
512
533
  await asyncio.sleep(KEY_ROTATION_DELAY_SECONDS)
513
534
  break # Break inner loop, continue outer loop with new key
514
535
 
515
- self._record_usage(key_usage, model_id, 0)
536
+ self._record_usage(key_usage, model_id, 0, track_usage)
516
537
  raise
517
538
  else:
518
539
  # Inner loop exhausted without success - continue to next key
@@ -528,7 +549,11 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
528
549
  raise RuntimeError(f"All retry attempts exhausted for {model_id}")
529
550
 
530
551
  async def _wrap_stream(
531
- self, generator: AsyncGenerator, key_usage: KeyUsage, model_id: str
552
+ self,
553
+ generator: AsyncGenerator,
554
+ key_usage: KeyUsage,
555
+ model_id: str,
556
+ track_usage: Optional[bool] = None,
532
557
  ) -> AsyncGenerator:
533
558
  """Wrap an async streaming response to track usage and handle errors."""
534
559
  final_tokens = 0
@@ -540,9 +565,9 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
540
565
  final_tokens = chunk_tokens
541
566
  yield chunk
542
567
  except Exception as e:
543
- if is_payment_required_error(e):
568
+ if is_dead_key_error(e):
544
569
  logger.warning(
545
- "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
570
+ "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
546
571
  model_id, get_key_suffix(key_usage.api_key)
547
572
  )
548
573
  key_usage.mark_dead()
@@ -556,7 +581,7 @@ class AsyncGenericRotatingClient(BaseGenericRotatingClient[T]):
556
581
  self.manager.force_rotate_index()
557
582
  raise
558
583
  finally:
559
- self._record_usage(key_usage, model_id, final_tokens)
584
+ self._record_usage(key_usage, model_id, final_tokens, track_usage)
560
585
 
561
586
 
562
587
  class AsyncGenericProxyHelper:
@@ -14,7 +14,7 @@ from ..config.constants import (
14
14
  from ..core.utils import (
15
15
  is_rate_limit_error,
16
16
  is_temporary_rate_limit_error,
17
- is_payment_required_error,
17
+ is_dead_key_error,
18
18
  get_key_suffix,
19
19
  )
20
20
  from ..core.backoff import ExponentialBackoff, BackoffConfig
@@ -97,12 +97,19 @@ class BaseRotatingClient:
97
97
  if self.base_url:
98
98
  self.client_kwargs['base_url'] = self.base_url
99
99
 
100
- def _record_usage(self, key_usage: KeyUsage, model_id: str, actual_tokens: int) -> None:
100
+ def _record_usage(
101
+ self,
102
+ key_usage: KeyUsage,
103
+ model_id: str,
104
+ actual_tokens: int,
105
+ track_usage: Optional[bool] = None,
106
+ ) -> None:
101
107
  self.manager.record_usage(
102
108
  key_obj=key_usage,
103
109
  model_id=model_id,
104
110
  actual_tokens=actual_tokens,
105
- estimated_tokens=self.estimated_tokens
111
+ estimated_tokens=self.estimated_tokens,
112
+ track_usage=track_usage
106
113
  )
107
114
 
108
115
  def _extract_usage(self, response: Any) -> int:
@@ -123,6 +130,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
123
130
  return SyncProxyHelper(self, [name])
124
131
 
125
132
  def _execute(self, path: List[str], args, kwargs):
133
+ # Per-call tracking override - popped before the kwargs reach the real
134
+ # OpenAI client, which rejects unknown keyword arguments.
135
+ track_usage = kwargs.pop('track_usage', None)
126
136
  model_id = kwargs.get('model', self.default_model)
127
137
  if 'model' not in kwargs:
128
138
  kwargs['model'] = model_id
@@ -132,9 +142,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
132
142
  kwargs['stream_options'] = {"include_usage": True}
133
143
 
134
144
  for attempt in range(self.max_retries + 1):
135
- key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
145
+ key_usage = self.manager.acquire_key(model_id, limits, self.estimated_tokens)
136
146
  if not key_usage:
137
- raise RuntimeError(f"No available keys for {model_id}")
147
+ raise RuntimeError(self.manager.unavailable_message(model_id))
138
148
 
139
149
  # Create backoff for temporary rate limits
140
150
  temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -154,9 +164,11 @@ class RotatingOpenAIClient(BaseRotatingClient):
154
164
  result = target(*args, **kwargs)
155
165
 
156
166
  if kwargs.get('stream', False):
157
- return self._wrap_stream(result, key_usage, model_id)
167
+ return self._wrap_stream(result, key_usage, model_id, track_usage)
158
168
 
159
- self._record_usage(key_usage, model_id, self._extract_usage(result))
169
+ self._record_usage(
170
+ key_usage, model_id, self._extract_usage(result), track_usage
171
+ )
160
172
  return result
161
173
 
162
174
  except Exception as e:
@@ -171,11 +183,11 @@ class RotatingOpenAIClient(BaseRotatingClient):
171
183
  time.sleep(delay)
172
184
  continue # Retry with SAME key
173
185
 
174
- # Payment required (quota/credits exhausted) - key is dead
186
+ # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
175
187
  # for this process, rotate to next key
176
- if is_payment_required_error(e) and attempt < self.max_retries:
188
+ if is_dead_key_error(e) and attempt < self.max_retries:
177
189
  logger.warning(
178
- "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
190
+ "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
179
191
  model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
180
192
  )
181
193
  key_usage.mark_dead()
@@ -194,7 +206,7 @@ class RotatingOpenAIClient(BaseRotatingClient):
194
206
  time.sleep(KEY_ROTATION_DELAY_SECONDS)
195
207
  break # Break inner loop, continue outer loop with new key
196
208
 
197
- self._record_usage(key_usage, model_id, 0)
209
+ self._record_usage(key_usage, model_id, 0, track_usage)
198
210
  raise
199
211
  else:
200
212
  # Inner loop exhausted without success - continue to next key
@@ -209,7 +221,13 @@ class RotatingOpenAIClient(BaseRotatingClient):
209
221
 
210
222
  raise RuntimeError(f"All retry attempts exhausted for {model_id}")
211
223
 
212
- def _wrap_stream(self, generator: Generator, key_usage: KeyUsage, model_id: str):
224
+ def _wrap_stream(
225
+ self,
226
+ generator: Generator,
227
+ key_usage: KeyUsage,
228
+ model_id: str,
229
+ track_usage: Optional[bool] = None,
230
+ ):
213
231
  final_tokens = 0
214
232
  try:
215
233
  for chunk in generator:
@@ -217,9 +235,9 @@ class RotatingOpenAIClient(BaseRotatingClient):
217
235
  final_tokens = chunk.usage.total_tokens
218
236
  yield chunk
219
237
  except Exception as e:
220
- if is_payment_required_error(e):
238
+ if is_dead_key_error(e):
221
239
  logger.warning(
222
- "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
240
+ "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
223
241
  model_id, get_key_suffix(key_usage.api_key)
224
242
  )
225
243
  key_usage.mark_dead()
@@ -233,7 +251,7 @@ class RotatingOpenAIClient(BaseRotatingClient):
233
251
  self.manager.force_rotate_index()
234
252
  raise
235
253
  finally:
236
- self._record_usage(key_usage, model_id, final_tokens)
254
+ self._record_usage(key_usage, model_id, final_tokens, track_usage)
237
255
 
238
256
  class SyncProxyHelper:
239
257
  def __init__(self, client: RotatingOpenAIClient, path: List[str]):
@@ -257,6 +275,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
257
275
  return AsyncProxyHelper(self, [name])
258
276
 
259
277
  async def _execute(self, path: List[str], args, kwargs: dict):
278
+ # Per-call tracking override - popped before the kwargs reach the real
279
+ # OpenAI client, which rejects unknown keyword arguments.
280
+ track_usage = kwargs.pop('track_usage', None)
260
281
  model_id = kwargs.get('model', self.default_model)
261
282
  if 'model' not in kwargs:
262
283
  kwargs['model'] = model_id
@@ -266,9 +287,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
266
287
  kwargs['stream_options'] = {"include_usage": True}
267
288
 
268
289
  for attempt in range(self.max_retries + 1):
269
- key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
290
+ key_usage = await self.manager.acquire_key_async(model_id, limits, self.estimated_tokens)
270
291
  if not key_usage:
271
- raise RuntimeError(f"No available keys for {model_id}")
292
+ raise RuntimeError(self.manager.unavailable_message(model_id))
272
293
 
273
294
  # Create backoff for temporary rate limits
274
295
  temp_backoff = ExponentialBackoff(BackoffConfig(
@@ -288,9 +309,11 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
288
309
  result = await target(*args, **kwargs)
289
310
 
290
311
  if kwargs.get('stream', False):
291
- return self._wrap_stream(result, key_usage, model_id)
312
+ return self._wrap_stream(result, key_usage, model_id, track_usage)
292
313
 
293
- self._record_usage(key_usage, model_id, self._extract_usage(result))
314
+ self._record_usage(
315
+ key_usage, model_id, self._extract_usage(result), track_usage
316
+ )
294
317
  return result
295
318
 
296
319
  except Exception as e:
@@ -305,11 +328,11 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
305
328
  await asyncio.sleep(delay)
306
329
  continue # Retry with SAME key
307
330
 
308
- # Payment required (quota/credits exhausted) - key is dead
331
+ # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
309
332
  # for this process, rotate to next key
310
- if is_payment_required_error(e) and attempt < self.max_retries:
333
+ if is_dead_key_error(e) and attempt < self.max_retries:
311
334
  logger.warning(
312
- "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
335
+ "Dead key (401/402/403) hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
313
336
  model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
314
337
  )
315
338
  key_usage.mark_dead()
@@ -328,7 +351,7 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
328
351
  await asyncio.sleep(KEY_ROTATION_DELAY_SECONDS)
329
352
  break # Break inner loop, continue outer loop with new key
330
353
 
331
- self._record_usage(key_usage, model_id, 0)
354
+ self._record_usage(key_usage, model_id, 0, track_usage)
332
355
  raise
333
356
  else:
334
357
  # Inner loop exhausted without success - continue to next key
@@ -343,7 +366,13 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
343
366
 
344
367
  raise RuntimeError(f"All retry attempts exhausted for {model_id}")
345
368
 
346
- async def _wrap_stream(self, generator: AsyncGenerator, key_usage: KeyUsage, model_id: str):
369
+ async def _wrap_stream(
370
+ self,
371
+ generator: AsyncGenerator,
372
+ key_usage: KeyUsage,
373
+ model_id: str,
374
+ track_usage: Optional[bool] = None,
375
+ ):
347
376
  final_tokens = 0
348
377
  try:
349
378
  async for chunk in generator:
@@ -351,9 +380,9 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
351
380
  final_tokens = chunk.usage.total_tokens
352
381
  yield chunk
353
382
  except Exception as e:
354
- if is_payment_required_error(e):
383
+ if is_dead_key_error(e):
355
384
  logger.warning(
356
- "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
385
+ "Dead key (401/402/403) hit during streaming for %s on key ...%s. Marking key dead.",
357
386
  model_id, get_key_suffix(key_usage.api_key)
358
387
  )
359
388
  key_usage.mark_dead()
@@ -367,7 +396,7 @@ class RotatingAsyncOpenAIClient(BaseRotatingClient):
367
396
  self.manager.force_rotate_index()
368
397
  raise
369
398
  finally:
370
- self._record_usage(key_usage, model_id, final_tokens)
399
+ self._record_usage(key_usage, model_id, final_tokens, track_usage)
371
400
 
372
401
  class AsyncProxyHelper:
373
402
  def __init__(self, client: RotatingAsyncOpenAIClient, path: List[str]):
@@ -31,5 +31,12 @@ TEMP_RATE_LIMIT_MULTIPLIER = 2.0
31
31
  # Key rotation delay (seconds to wait after rotating to a new key)
32
32
  KEY_ROTATION_DELAY_SECONDS = 0.5
33
33
 
34
+ # How long acquire_key() polls for a live key to leave its rate window /
35
+ # cooldown before giving up. A single production key at its per-minute cap
36
+ # is the common case (Cohere, 2026-09-07): the right answer is to wait for
37
+ # the window, not to abort the caller's whole run.
38
+ KEY_WAIT_MAX_SECONDS = 120.0
39
+ KEY_WAIT_POLL_SECONDS = 1.0
40
+
34
41
  # Usage logging batch size
35
42
  USAGE_LOG_BATCH_SIZE = 50
@@ -14,6 +14,7 @@ from .utils import (
14
14
  get_key_suffix,
15
15
  is_rate_limit_error,
16
16
  is_auth_error,
17
+ is_dead_key_error,
17
18
  validate_api_key,
18
19
  )
19
20
  from .backoff import ExponentialBackoff, BackoffConfig
@@ -33,6 +34,7 @@ __all__ = [
33
34
  "get_key_suffix",
34
35
  "is_rate_limit_error",
35
36
  "is_auth_error",
37
+ "is_dead_key_error",
36
38
  "validate_api_key",
37
39
  # Backoff
38
40
  "ExponentialBackoff",
@@ -241,6 +241,29 @@ def is_auth_error(e: Exception) -> bool:
241
241
  return False
242
242
 
243
243
 
244
+ # A key that answers 401/402/403 is unusable for good (revoked, unpaid,
245
+ # or its project denied) — no cooldown will fix it, so rotate off it.
246
+ DEAD_KEY_STATUS_CODES = AUTH_STATUS_CODES | PAYMENT_REQUIRED_STATUS_CODES
247
+
248
+
249
+ def is_dead_key_error(e: Exception) -> bool:
250
+ """
251
+ Detect errors that mean THIS KEY is dead, not merely busy: HTTP 401/403
252
+ (invalid, revoked, or project-denied key — e.g. Gemini's
253
+ "Your project has been denied access") and 402 (credits exhausted).
254
+
255
+ A status_code, when present, is authoritative: 429s are never dead keys
256
+ whatever their message says. Without one, fall back to the string
257
+ heuristics of is_payment_required_error / is_auth_error.
258
+ """
259
+ status = getattr(e, "status_code", None)
260
+ if status in DEAD_KEY_STATUS_CODES:
261
+ return True
262
+ if status == 429:
263
+ return False
264
+ return is_payment_required_error(e) or is_auth_error(e)
265
+
266
+
244
267
  def validate_api_key(api_key: str) -> bool:
245
268
  """
246
269
  Basic validation of API key format.
@@ -15,7 +15,7 @@ from ..config.constants import (
15
15
  from ..core.utils import (
16
16
  is_rate_limit_error,
17
17
  is_temporary_rate_limit_error,
18
- is_payment_required_error,
18
+ is_dead_key_error,
19
19
  get_key_suffix,
20
20
  )
21
21
  from ..core.backoff import ExponentialBackoff, BackoffConfig
@@ -115,7 +115,12 @@ class RotatingCredentialsMixin:
115
115
  def _get_retry_limit(self) -> int:
116
116
  return min(self._max_retries, len(self.wrapper.manager.keys) - 1)
117
117
 
118
- def _record_usage(self, key_obj: KeyUsage, response: Optional["ModelResponse"]):
118
+ def _record_usage(
119
+ self,
120
+ key_obj: KeyUsage,
121
+ response: Optional["ModelResponse"],
122
+ track_usage: Optional[bool] = None,
123
+ ):
119
124
  """
120
125
  Extracts usage from the response and reports it to the manager.
121
126
  Falls back to estimated_tokens if response is None or usage is missing.
@@ -136,7 +141,8 @@ class RotatingCredentialsMixin:
136
141
  key_obj=key_obj,
137
142
  model_id=self.model_id,
138
143
  actual_tokens=actual_tokens,
139
- estimated_tokens=self._estimated_tokens
144
+ estimated_tokens=self._estimated_tokens,
145
+ track_usage=track_usage
140
146
  )
141
147
 
142
148
  def _create_temp_backoff(self) -> ExponentialBackoff:
@@ -148,6 +154,9 @@ class RotatingCredentialsMixin:
148
154
  ))
149
155
 
150
156
  def invoke(self, *args, **kwargs) -> "ModelResponse":
157
+ # Per-call tracking override - popped before the kwargs reach the
158
+ # underlying model, which does not know this keyword.
159
+ track_usage = kwargs.pop("track_usage", None)
151
160
  limit = self._get_retry_limit()
152
161
 
153
162
  for attempt in range(limit + 1):
@@ -157,7 +166,7 @@ class RotatingCredentialsMixin:
157
166
  for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
158
167
  try:
159
168
  response = super().invoke(*args, **kwargs)
160
- self._record_usage(key_usage, response)
169
+ self._record_usage(key_usage, response, track_usage)
161
170
  return response
162
171
  except Exception as e:
163
172
  # Check for temporary rate limit first - retry with SAME key
@@ -171,11 +180,11 @@ class RotatingCredentialsMixin:
171
180
  time.sleep(delay)
172
181
  continue # Retry with SAME key
173
182
 
174
- # Payment required (quota/credits exhausted) - key is dead
183
+ # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
175
184
  # for this process, rotate to next key
176
- if is_payment_required_error(e) and attempt < limit:
185
+ if is_dead_key_error(e) and attempt < limit:
177
186
  self.logger.warning(
178
- "402 Hit on key %s (Sync) [%s]. Marking key dead and rotating (%d/%d).",
187
+ "Dead key (401/402/403) on key %s (Sync) [%s]. Marking key dead and rotating (%d/%d).",
179
188
  get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
180
189
  )
181
190
  key_usage.mark_dead()
@@ -205,6 +214,9 @@ class RotatingCredentialsMixin:
205
214
  raise
206
215
 
207
216
  async def ainvoke(self, *args, **kwargs) -> "ModelResponse":
217
+ # Per-call tracking override - popped before the kwargs reach the
218
+ # underlying model, which does not know this keyword.
219
+ track_usage = kwargs.pop("track_usage", None)
208
220
  limit = self._get_retry_limit()
209
221
 
210
222
  for attempt in range(limit + 1):
@@ -214,7 +226,7 @@ class RotatingCredentialsMixin:
214
226
  for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
215
227
  try:
216
228
  response = await super().ainvoke(*args, **kwargs)
217
- self._record_usage(key_usage, response)
229
+ self._record_usage(key_usage, response, track_usage)
218
230
  return response
219
231
  except Exception as e:
220
232
  # Check for temporary rate limit first - retry with SAME key
@@ -228,11 +240,11 @@ class RotatingCredentialsMixin:
228
240
  await asyncio.sleep(delay)
229
241
  continue # Retry with SAME key
230
242
 
231
- # Payment required (quota/credits exhausted) - key is dead
243
+ # Dead key (401/403 revoked/denied, 402 credits exhausted) - key is dead
232
244
  # for this process, rotate to next key
233
- if is_payment_required_error(e) and attempt < limit:
245
+ if is_dead_key_error(e) and attempt < limit:
234
246
  self.logger.warning(
235
- "402 Hit on key %s (Async) [%s]. Marking key dead and rotating (%d/%d).",
247
+ "Dead key (401/402/403) on key %s (Async) [%s]. Marking key dead and rotating (%d/%d).",
236
248
  get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
237
249
  )
238
250
  key_usage.mark_dead()
@@ -262,6 +274,9 @@ class RotatingCredentialsMixin:
262
274
  raise
263
275
 
264
276
  def invoke_stream(self, *args, **kwargs) -> Iterator["ModelResponse"]:
277
+ # Per-call tracking override - popped before the kwargs reach the
278
+ # underlying model, which does not know this keyword.
279
+ track_usage = kwargs.pop("track_usage", None)
265
280
  limit = self._get_retry_limit()
266
281
 
267
282
  for attempt in range(limit + 1):
@@ -280,10 +295,10 @@ class RotatingCredentialsMixin:
280
295
  yield chunk
281
296
  if final_usage:
282
297
  dummy_response = _UsageResponse(final_usage)
283
- self._record_usage(key_usage, dummy_response)
298
+ self._record_usage(key_usage, dummy_response, track_usage)
284
299
  else:
285
300
  # If the stream didn't return usage data, fallback to estimation
286
- self._record_usage(key_usage, None)
301
+ self._record_usage(key_usage, None, track_usage)
287
302
 
288
303
  return
289
304
  except Exception as e:
@@ -300,9 +315,9 @@ class RotatingCredentialsMixin:
300
315
  time.sleep(delay)
301
316
  continue
302
317
 
303
- if is_payment_required_error(e) and attempt < limit:
318
+ if is_dead_key_error(e) and attempt < limit:
304
319
  self.logger.warning(
305
- "402 Hit on key %s (Sync Stream) [%s]. Marking key dead and rotating (%d/%d).",
320
+ "Dead key (401/402/403) on key %s (Sync Stream) [%s]. Marking key dead and rotating (%d/%d).",
306
321
  get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
307
322
  )
308
323
  key_usage.mark_dead()
@@ -326,6 +341,9 @@ class RotatingCredentialsMixin:
326
341
  raise
327
342
 
328
343
  async def ainvoke_stream(self, *args, **kwargs) -> AsyncIterator["ModelResponse"]:
344
+ # Per-call tracking override - popped before the kwargs reach the
345
+ # underlying model, which does not know this keyword.
346
+ track_usage = kwargs.pop("track_usage", None)
329
347
  limit = self._get_retry_limit()
330
348
 
331
349
  for attempt in range(limit + 1):
@@ -347,9 +365,9 @@ class RotatingCredentialsMixin:
347
365
  # Stream completed successfully. Record usage.
348
366
  if final_usage:
349
367
  dummy_response = _UsageResponse(final_usage)
350
- self._record_usage(key_usage, dummy_response)
368
+ self._record_usage(key_usage, dummy_response, track_usage)
351
369
  else:
352
- self._record_usage(key_usage, None)
370
+ self._record_usage(key_usage, None, track_usage)
353
371
 
354
372
  return
355
373
  except Exception as e:
@@ -366,9 +384,9 @@ class RotatingCredentialsMixin:
366
384
  await asyncio.sleep(delay)
367
385
  continue
368
386
 
369
- if is_payment_required_error(e) and attempt < limit:
387
+ if is_dead_key_error(e) and attempt < limit:
370
388
  self.logger.warning(
371
- "402 Hit on key %s (Async Stream) [%s]. Marking key dead and rotating (%d/%d).",
389
+ "Dead key (401/402/403) on key %s (Async Stream) [%s]. Marking key dead and rotating (%d/%d).",
372
390
  get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
373
391
  )
374
392
  key_usage.mark_dead()