Mirrowel commited on
Commit
bd84d38
·
1 Parent(s): 4dfb828

feat(rotation): ✨ add sequential rotation mode with provider-specific quota parsing

Browse files

Introduces a new credential rotation mode system that allows providers to choose between "balanced" (distribute load evenly) and "sequential" (use until exhausted) strategies. Sequential mode is particularly beneficial for providers with cache-preserving features like Antigravity's thinking signature caches.

Key changes:
- Added ROTATION_MODE_{PROVIDER} environment variable support with comprehensive documentation in .env.example
- Implemented provider-specific quota error parsing for Antigravity and Gemini CLI providers, extracting retry_after from Google RPC error format (handles compound durations like "143h4m52.73s")
- Extended ProviderInterface with rotation mode configuration and parse_quota_error() method
- Updated UsageManager to support sequential credential selection that preserves sticky credential usage until quota exhaustion
- Enhanced error_handler.py classify_error() to attempt provider-specific parsing before falling back to generic classification
- Added rotation mode management UI in settings_tool.py with visual indicators for configured vs default modes
- Preserved long-term cooldowns during daily reset to prevent premature quota retry
- Updated all classify_error() call sites to pass provider parameter for context-aware parsing

Provider defaults:
- Antigravity: sequential (preserves thinking caches, handles weekly quota reset)
- Gemini CLI: balanced (short cooldowns in seconds/minutes)
- All others: balanced (standard per-minute rate limits)

The sequential mode ensures the same credential is reused until it hits a cooldown (429 error), at which point the system switches to the next available credential. This maximizes cache hit rates for providers that maintain request context across API calls.

.env.example CHANGED
@@ -159,6 +159,32 @@ MAX_CONCURRENT_REQUESTS_PER_KEY_GEMINI=1
159
  MAX_CONCURRENT_REQUESTS_PER_KEY_ANTHROPIC=1
160
  MAX_CONCURRENT_REQUESTS_PER_KEY_IFLOW=1
161
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
162
  # ------------------------------------------------------------------------------
163
  # | [ADVANCED] Proxy Configuration |
164
  # ------------------------------------------------------------------------------
 
159
  MAX_CONCURRENT_REQUESTS_PER_KEY_ANTHROPIC=1
160
  MAX_CONCURRENT_REQUESTS_PER_KEY_IFLOW=1
161
 
162
+ # --- Credential Rotation Mode ---
163
+ # Controls how credentials are rotated when multiple are available for a provider.
164
+ # This affects how the proxy selects the next credential to use for requests.
165
+ #
166
+ # Available modes:
167
+ # balanced - (Default) Rotate credentials evenly across requests to distribute load.
168
+ # Best for API keys with per-minute rate limits.
169
+ # sequential - Use one credential until it's exhausted (429 error), then switch to next.
170
+ # Best for credentials with daily/weekly quotas (e.g., free tier accounts).
171
+ # When a credential hits quota, it's put on cooldown based on the reset time
172
+ # parsed from the provider's error response.
173
+ #
174
+ # Format: ROTATION_MODE_<PROVIDER_NAME>=<mode>
175
+ #
176
+ # Provider Defaults:
177
+ # - antigravity: sequential (free tier accounts with daily quotas)
178
+ # - All others: balanced
179
+ #
180
+ # Example:
181
+ # ROTATION_MODE_GEMINI=sequential # Use Gemini keys until quota exhausted
182
+ # ROTATION_MODE_OPENAI=balanced # Distribute load across OpenAI keys (default)
183
+ # ROTATION_MODE_ANTIGRAVITY=balanced # Override Antigravity's sequential default
184
+ #
185
+ # ROTATION_MODE_GEMINI=balanced
186
+ # ROTATION_MODE_ANTIGRAVITY=sequential
187
+
188
  # ------------------------------------------------------------------------------
189
  # | [ADVANCED] Proxy Configuration |
190
  # ------------------------------------------------------------------------------
src/proxy_app/settings_tool.py CHANGED
@@ -17,37 +17,38 @@ console = Console()
17
 
18
  def clear_screen():
19
  """
20
- Cross-platform terminal clear that works robustly on both
21
  classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac).
22
-
23
  Uses native OS commands instead of ANSI escape sequences:
24
  - Windows (conhost & Windows Terminal): cls
25
  - Unix-like systems (Linux, Mac): clear
26
  """
27
- os.system('cls' if os.name == 'nt' else 'clear')
28
 
29
 
30
  class AdvancedSettings:
31
  """Manages pending changes to .env"""
32
-
33
  def __init__(self):
34
  self.env_file = Path.cwd() / ".env"
35
  self.pending_changes = {} # key -> value (None means delete)
36
  self.load_current_settings()
37
-
38
  def load_current_settings(self):
39
  """Load current .env values into env vars"""
40
  from dotenv import load_dotenv
 
41
  load_dotenv(override=True)
42
-
43
  def set(self, key: str, value: str):
44
  """Stage a change"""
45
  self.pending_changes[key] = value
46
-
47
  def remove(self, key: str):
48
  """Stage a removal"""
49
  self.pending_changes[key] = None
50
-
51
  def save(self):
52
  """Write pending changes to .env"""
53
  for key, value in self.pending_changes.items():
@@ -57,14 +58,14 @@ class AdvancedSettings:
57
  else:
58
  # Set key
59
  set_key(str(self.env_file), key, value)
60
-
61
  self.pending_changes.clear()
62
  self.load_current_settings()
63
-
64
  def discard(self):
65
  """Discard pending changes"""
66
  self.pending_changes.clear()
67
-
68
  def has_pending(self) -> bool:
69
  """Check if there are pending changes"""
70
  return bool(self.pending_changes)
@@ -72,14 +73,14 @@ class AdvancedSettings:
72
 
73
  class CustomProviderManager:
74
  """Manages custom provider API bases"""
75
-
76
  def __init__(self, settings: AdvancedSettings):
77
  self.settings = settings
78
-
79
  def get_current_providers(self) -> Dict[str, str]:
80
  """Get currently configured custom providers"""
81
  from proxy_app.provider_urls import PROVIDER_URL_MAP
82
-
83
  providers = {}
84
  for key, value in os.environ.items():
85
  if key.endswith("_API_BASE"):
@@ -88,16 +89,16 @@ class CustomProviderManager:
88
  if provider not in PROVIDER_URL_MAP:
89
  providers[provider] = value
90
  return providers
91
-
92
  def add_provider(self, name: str, api_base: str):
93
  """Add PROVIDER_API_BASE"""
94
  key = f"{name.upper()}_API_BASE"
95
  self.settings.set(key, api_base)
96
-
97
  def edit_provider(self, name: str, api_base: str):
98
  """Edit PROVIDER_API_BASE"""
99
  self.add_provider(name, api_base)
100
-
101
  def remove_provider(self, name: str):
102
  """Remove PROVIDER_API_BASE"""
103
  key = f"{name.upper()}_API_BASE"
@@ -106,10 +107,10 @@ class CustomProviderManager:
106
 
107
  class ModelDefinitionManager:
108
  """Manages PROVIDER_MODELS"""
109
-
110
  def __init__(self, settings: AdvancedSettings):
111
  self.settings = settings
112
-
113
  def get_current_provider_models(self, provider: str) -> Optional[Dict]:
114
  """Get currently configured models for a provider"""
115
  key = f"{provider.upper()}_MODELS"
@@ -120,7 +121,7 @@ class ModelDefinitionManager:
120
  except (json.JSONDecodeError, ValueError):
121
  return None
122
  return None
123
-
124
  def get_all_providers_with_models(self) -> Dict[str, int]:
125
  """Get all providers with model definitions"""
126
  providers = {}
@@ -136,13 +137,13 @@ class ModelDefinitionManager:
136
  except (json.JSONDecodeError, ValueError):
137
  pass
138
  return providers
139
-
140
  def set_models(self, provider: str, models: Dict[str, Dict[str, Any]]):
141
  """Set PROVIDER_MODELS"""
142
  key = f"{provider.upper()}_MODELS"
143
  value = json.dumps(models)
144
  self.settings.set(key, value)
145
-
146
  def remove_models(self, provider: str):
147
  """Remove PROVIDER_MODELS"""
148
  key = f"{provider.upper()}_MODELS"
@@ -151,10 +152,10 @@ class ModelDefinitionManager:
151
 
152
  class ConcurrencyManager:
153
  """Manages MAX_CONCURRENT_REQUESTS_PER_KEY_PROVIDER"""
154
-
155
  def __init__(self, settings: AdvancedSettings):
156
  self.settings = settings
157
-
158
  def get_current_limits(self) -> Dict[str, int]:
159
  """Get currently configured concurrency limits"""
160
  limits = {}
@@ -166,18 +167,73 @@ class ConcurrencyManager:
166
  except (json.JSONDecodeError, ValueError):
167
  pass
168
  return limits
169
-
170
  def set_limit(self, provider: str, limit: int):
171
  """Set concurrency limit"""
172
  key = f"MAX_CONCURRENT_REQUESTS_PER_KEY_{provider.upper()}"
173
  self.settings.set(key, str(limit))
174
-
175
  def remove_limit(self, provider: str):
176
  """Remove concurrency limit (reset to default)"""
177
  key = f"MAX_CONCURRENT_REQUESTS_PER_KEY_{provider.upper()}"
178
  self.settings.remove(key)
179
 
180
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
181
  # =============================================================================
182
  # PROVIDER-SPECIFIC SETTINGS DEFINITIONS
183
  # =============================================================================
@@ -294,24 +350,26 @@ PROVIDER_SETTINGS_MAP = {
294
 
295
  class ProviderSettingsManager:
296
  """Manages provider-specific configuration settings"""
297
-
298
  def __init__(self, settings: AdvancedSettings):
299
  self.settings = settings
300
-
301
  def get_available_providers(self) -> List[str]:
302
  """Get list of providers with specific settings available"""
303
  return list(PROVIDER_SETTINGS_MAP.keys())
304
-
305
- def get_provider_settings_definitions(self, provider: str) -> Dict[str, Dict[str, Any]]:
 
 
306
  """Get settings definitions for a provider"""
307
  return PROVIDER_SETTINGS_MAP.get(provider, {})
308
-
309
  def get_current_value(self, key: str, definition: Dict[str, Any]) -> Any:
310
  """Get current value of a setting from environment"""
311
  env_value = os.getenv(key)
312
  if env_value is None:
313
  return definition.get("default")
314
-
315
  setting_type = definition.get("type", "str")
316
  try:
317
  if setting_type == "bool":
@@ -322,7 +380,7 @@ class ProviderSettingsManager:
322
  return env_value
323
  except (ValueError, AttributeError):
324
  return definition.get("default")
325
-
326
  def get_all_current_values(self, provider: str) -> Dict[str, Any]:
327
  """Get all current values for a provider"""
328
  definitions = self.get_provider_settings_definitions(provider)
@@ -330,7 +388,7 @@ class ProviderSettingsManager:
330
  for key, definition in definitions.items():
331
  values[key] = self.get_current_value(key, definition)
332
  return values
333
-
334
  def set_value(self, key: str, value: Any, definition: Dict[str, Any]):
335
  """Set a setting value, converting to string for .env storage"""
336
  setting_type = definition.get("type", "str")
@@ -339,11 +397,11 @@ class ProviderSettingsManager:
339
  else:
340
  str_value = str(value)
341
  self.settings.set(key, str_value)
342
-
343
  def reset_to_default(self, key: str):
344
  """Remove a setting to reset it to default"""
345
  self.settings.remove(key)
346
-
347
  def get_modified_settings(self, provider: str) -> Dict[str, Any]:
348
  """Get settings that differ from defaults"""
349
  definitions = self.get_provider_settings_definitions(provider)
@@ -358,80 +416,96 @@ class ProviderSettingsManager:
358
 
359
  class SettingsTool:
360
  """Main settings tool TUI"""
361
-
362
  def __init__(self):
363
  self.console = Console()
364
  self.settings = AdvancedSettings()
365
  self.provider_mgr = CustomProviderManager(self.settings)
366
  self.model_mgr = ModelDefinitionManager(self.settings)
367
  self.concurrency_mgr = ConcurrencyManager(self.settings)
 
368
  self.provider_settings_mgr = ProviderSettingsManager(self.settings)
369
  self.running = True
370
-
371
  def get_available_providers(self) -> List[str]:
372
  """Get list of providers that have credentials configured"""
373
  env_file = Path.cwd() / ".env"
374
  providers = set()
375
-
376
  # Scan for providers with API keys from local .env
377
  if env_file.exists():
378
  try:
379
- with open(env_file, 'r', encoding='utf-8') as f:
380
  for line in f:
381
  line = line.strip()
382
- if "_API_KEY" in line and "PROXY_API_KEY" not in line and "=" in line:
 
 
 
 
383
  provider = line.split("_API_KEY")[0].strip().lower()
384
  providers.add(provider)
385
  except (IOError, OSError):
386
  pass
387
-
388
  # Also check for OAuth providers from files
389
  oauth_dir = Path("oauth_credentials")
390
  if oauth_dir.exists():
391
  for file in oauth_dir.glob("*_oauth_*.json"):
392
  provider = file.name.split("_oauth_")[0]
393
  providers.add(provider)
394
-
395
  return sorted(list(providers))
396
 
397
  def run(self):
398
  """Main loop"""
399
  while self.running:
400
  self.show_main_menu()
401
-
402
  def show_main_menu(self):
403
  """Display settings categories"""
404
  clear_screen()
405
-
406
- self.console.print(Panel.fit(
407
- "[bold cyan]🔧 Advanced Settings Configuration[/bold cyan]",
408
- border_style="cyan"
409
- ))
410
-
 
 
411
  self.console.print()
412
  self.console.print("[bold]⚙️ Configuration Categories[/bold]")
413
  self.console.print()
414
  self.console.print(" 1. 🌐 Custom Provider API Bases")
415
  self.console.print(" 2. 📦 Provider Model Definitions")
416
  self.console.print(" 3. ⚡ Concurrency Limits")
417
- self.console.print(" 4. 🔬 Provider-Specific Settings")
418
- self.console.print(" 5. 💾 Save & Exit")
419
- self.console.print(" 6. 🚫 Exit Without Saving")
420
-
 
421
  self.console.print()
422
  self.console.print("━" * 70)
423
-
424
  if self.settings.has_pending():
425
- self.console.print("[yellow]ℹ️ Changes are pending until you select \"Save & Exit\"[/yellow]")
 
 
426
  else:
427
  self.console.print("[dim]ℹ️ No pending changes[/dim]")
428
-
429
  self.console.print()
430
- self.console.print("[dim]⚠️ Model filters not supported - edit .env for IGNORE_MODELS_* / WHITELIST_MODELS_*[/dim]")
 
 
431
  self.console.print()
432
-
433
- choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5", "6"], show_choices=False)
434
-
 
 
 
 
435
  if choice == "1":
436
  self.manage_custom_providers()
437
  elif choice == "2":
@@ -439,34 +513,38 @@ class SettingsTool:
439
  elif choice == "3":
440
  self.manage_concurrency_limits()
441
  elif choice == "4":
442
- self.manage_provider_settings()
443
  elif choice == "5":
444
- self.save_and_exit()
445
  elif choice == "6":
 
 
446
  self.exit_without_saving()
447
-
448
  def manage_custom_providers(self):
449
  """Manage custom provider API bases"""
450
  while True:
451
  clear_screen()
452
-
453
  providers = self.provider_mgr.get_current_providers()
454
-
455
- self.console.print(Panel.fit(
456
- "[bold cyan]🌐 Custom Provider API Bases[/bold cyan]",
457
- border_style="cyan"
458
- ))
459
-
 
 
460
  self.console.print()
461
  self.console.print("[bold]📋 Configured Custom Providers[/bold]")
462
  self.console.print("━" * 70)
463
-
464
  if providers:
465
  for name, base in providers.items():
466
  self.console.print(f" • {name:15} {base}")
467
  else:
468
  self.console.print(" [dim]No custom providers configured[/dim]")
469
-
470
  self.console.print()
471
  self.console.print("━" * 70)
472
  self.console.print()
@@ -476,94 +554,116 @@ class SettingsTool:
476
  self.console.print(" 2. ✏️ Edit Existing Provider")
477
  self.console.print(" 3. 🗑️ Remove Provider")
478
  self.console.print(" 4. ↩️ Back to Settings Menu")
479
-
480
  self.console.print()
481
  self.console.print("━" * 70)
482
  self.console.print()
483
-
484
- choice = Prompt.ask("Select option", choices=["1", "2", "3", "4"], show_choices=False)
485
-
 
 
486
  if choice == "1":
487
  name = Prompt.ask("Provider name (e.g., 'opencode')").strip().lower()
488
  if name:
489
  api_base = Prompt.ask("API Base URL").strip()
490
  if api_base:
491
  self.provider_mgr.add_provider(name, api_base)
492
- self.console.print(f"\n[green]✅ Custom provider '{name}' configured![/green]")
493
- self.console.print(f" To use: set {name.upper()}_API_KEY in credentials")
 
 
 
 
494
  input("\nPress Enter to continue...")
495
-
496
  elif choice == "2":
497
  if not providers:
498
  self.console.print("\n[yellow]No providers to edit[/yellow]")
499
  input("\nPress Enter to continue...")
500
  continue
501
-
502
  # Show numbered list
503
  self.console.print("\n[bold]Select provider to edit:[/bold]")
504
  providers_list = list(providers.keys())
505
  for idx, prov in enumerate(providers_list, 1):
506
  self.console.print(f" {idx}. {prov}")
507
-
508
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers_list) + 1)])
 
 
 
509
  name = providers_list[choice_idx - 1]
510
  current_base = providers.get(name, "")
511
-
512
  self.console.print(f"\nCurrent API Base: {current_base}")
513
- new_base = Prompt.ask("New API Base [press Enter to keep current]", default=current_base).strip()
514
-
 
 
515
  if new_base and new_base != current_base:
516
  self.provider_mgr.edit_provider(name, new_base)
517
- self.console.print(f"\n[green]✅ Custom provider '{name}' updated![/green]")
 
 
518
  else:
519
  self.console.print("\n[yellow]No changes made[/yellow]")
520
  input("\nPress Enter to continue...")
521
-
522
  elif choice == "3":
523
  if not providers:
524
  self.console.print("\n[yellow]No providers to remove[/yellow]")
525
  input("\nPress Enter to continue...")
526
  continue
527
-
528
  # Show numbered list
529
  self.console.print("\n[bold]Select provider to remove:[/bold]")
530
  providers_list = list(providers.keys())
531
  for idx, prov in enumerate(providers_list, 1):
532
  self.console.print(f" {idx}. {prov}")
533
-
534
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers_list) + 1)])
 
 
 
535
  name = providers_list[choice_idx - 1]
536
-
537
  if Confirm.ask(f"Remove '{name}'?"):
538
  self.provider_mgr.remove_provider(name)
539
- self.console.print(f"\n[green]✅ Provider '{name}' removed![/green]")
 
 
540
  input("\nPress Enter to continue...")
541
-
542
  elif choice == "4":
543
  break
544
-
545
  def manage_model_definitions(self):
546
  """Manage provider model definitions"""
547
  while True:
548
  clear_screen()
549
-
550
  all_providers = self.model_mgr.get_all_providers_with_models()
551
-
552
- self.console.print(Panel.fit(
553
- "[bold cyan]📦 Provider Model Definitions[/bold cyan]",
554
- border_style="cyan"
555
- ))
556
-
 
 
557
  self.console.print()
558
  self.console.print("[bold]📋 Configured Provider Models[/bold]")
559
  self.console.print("━" * 70)
560
-
561
  if all_providers:
562
  for provider, count in all_providers.items():
563
- self.console.print(f" • {provider:15} {count} model{'s' if count > 1 else ''}")
 
 
564
  else:
565
  self.console.print(" [dim]No model definitions configured[/dim]")
566
-
567
  self.console.print()
568
  self.console.print("━" * 70)
569
  self.console.print()
@@ -574,13 +674,15 @@ class SettingsTool:
574
  self.console.print(" 3. 👁️ View Provider Models")
575
  self.console.print(" 4. 🗑️ Remove Provider Models")
576
  self.console.print(" 5. ↩️ Back to Settings Menu")
577
-
578
  self.console.print()
579
  self.console.print("━" * 70)
580
  self.console.print()
581
-
582
- choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5"], show_choices=False)
583
-
 
 
584
  if choice == "1":
585
  self.add_model_definitions()
586
  elif choice == "2":
@@ -600,57 +702,71 @@ class SettingsTool:
600
  self.console.print("\n[yellow]No providers to remove[/yellow]")
601
  input("\nPress Enter to continue...")
602
  continue
603
-
604
  # Show numbered list
605
- self.console.print("\n[bold]Select provider to remove models from:[/bold]")
 
 
606
  providers_list = list(all_providers.keys())
607
  for idx, prov in enumerate(providers_list, 1):
608
  self.console.print(f" {idx}. {prov}")
609
-
610
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers_list) + 1)])
 
 
 
611
  provider = providers_list[choice_idx - 1]
612
-
613
  if Confirm.ask(f"Remove all model definitions for '{provider}'?"):
614
  self.model_mgr.remove_models(provider)
615
- self.console.print(f"\n[green]✅ Model definitions removed for '{provider}'![/green]")
 
 
616
  input("\nPress Enter to continue...")
617
  elif choice == "5":
618
  break
619
-
620
  def add_model_definitions(self):
621
  """Add model definitions for a provider"""
622
  # Get available providers from credentials
623
  available_providers = self.get_available_providers()
624
-
625
  if not available_providers:
626
- self.console.print("\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]")
 
 
627
  input("\nPress Enter to continue...")
628
  return
629
-
630
  # Show provider selection menu
631
  self.console.print("\n[bold]Select provider:[/bold]")
632
  for idx, prov in enumerate(available_providers, 1):
633
  self.console.print(f" {idx}. {prov}")
634
- self.console.print(f" {len(available_providers) + 1}. Enter custom provider name")
635
-
636
- choice = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(available_providers) + 2)])
637
-
 
 
 
 
 
638
  if choice == len(available_providers) + 1:
639
  provider = Prompt.ask("Provider name").strip().lower()
640
  else:
641
  provider = available_providers[choice - 1]
642
-
643
  if not provider:
644
  return
645
-
646
  self.console.print("\nHow would you like to define models?")
647
  self.console.print(" 1. Simple list (names only)")
648
  self.console.print(" 2. Advanced (names with IDs and options)")
649
-
650
  mode = Prompt.ask("Select mode", choices=["1", "2"], show_choices=False)
651
-
652
  models = {}
653
-
654
  if mode == "1":
655
  # Simple mode
656
  while True:
@@ -667,13 +783,19 @@ class SettingsTool:
667
  break
668
  if name:
669
  model_def = {}
670
- model_id = Prompt.ask(f"Model ID [press Enter to use '{name}']", default=name).strip()
 
 
671
  if model_id and model_id != name:
672
  model_def["id"] = model_id
673
-
674
  # Optional: model options
675
- if Confirm.ask("Add model options (e.g., temperature limits)?", default=False):
676
- self.console.print("\nEnter options as key=value pairs (one per line, 'done' to finish):")
 
 
 
 
677
  options = {}
678
  while True:
679
  opt = Prompt.ask("Option").strip()
@@ -690,121 +812,143 @@ class SettingsTool:
690
  options[key.strip()] = value
691
  if options:
692
  model_def["options"] = options
693
-
694
  models[name] = model_def
695
-
696
  if models:
697
  self.model_mgr.set_models(provider, models)
698
- self.console.print(f"\n[green]✅ Model definitions saved for '{provider}'![/green]")
 
 
699
  else:
700
  self.console.print("\n[yellow]No models added[/yellow]")
701
-
702
  input("\nPress Enter to continue...")
703
-
704
  def edit_model_definitions(self, providers: List[str]):
705
  """Edit existing model definitions"""
706
  # Show numbered list
707
  self.console.print("\n[bold]Select provider to edit:[/bold]")
708
  for idx, prov in enumerate(providers, 1):
709
  self.console.print(f" {idx}. {prov}")
710
-
711
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers) + 1)])
 
 
712
  provider = providers[choice_idx - 1]
713
-
714
  current_models = self.model_mgr.get_current_provider_models(provider)
715
  if not current_models:
716
  self.console.print(f"\n[yellow]No models found for '{provider}'[/yellow]")
717
  input("\nPress Enter to continue...")
718
  return
719
-
720
  # Convert to dict if list
721
  if isinstance(current_models, list):
722
  current_models = {m: {} for m in current_models}
723
-
724
  while True:
725
  clear_screen()
726
  self.console.print(f"[bold]Editing models for: {provider}[/bold]\n")
727
  self.console.print("Current models:")
728
  for i, (name, definition) in enumerate(current_models.items(), 1):
729
- model_id = definition.get("id", name) if isinstance(definition, dict) else name
 
 
730
  self.console.print(f" {i}. {name} (ID: {model_id})")
731
-
732
  self.console.print("\nOptions:")
733
  self.console.print(" 1. Add new model")
734
  self.console.print(" 2. Edit existing model")
735
  self.console.print(" 3. Remove model")
736
  self.console.print(" 4. Done")
737
-
738
- choice = Prompt.ask("\nSelect option", choices=["1", "2", "3", "4"], show_choices=False)
739
-
 
 
740
  if choice == "1":
741
  name = Prompt.ask("New model name").strip()
742
  if name and name not in current_models:
743
  model_id = Prompt.ask("Model ID", default=name).strip()
744
  current_models[name] = {"id": model_id} if model_id != name else {}
745
-
746
  elif choice == "2":
747
  # Show numbered list
748
  models_list = list(current_models.keys())
749
  self.console.print("\n[bold]Select model to edit:[/bold]")
750
  for idx, model_name in enumerate(models_list, 1):
751
  self.console.print(f" {idx}. {model_name}")
752
-
753
- model_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(models_list) + 1)])
 
 
 
754
  name = models_list[model_idx - 1]
755
-
756
  current_def = current_models[name]
757
- current_id = current_def.get("id", name) if isinstance(current_def, dict) else name
758
-
 
 
 
 
759
  new_id = Prompt.ask("Model ID", default=current_id).strip()
760
  current_models[name] = {"id": new_id} if new_id != name else {}
761
-
762
  elif choice == "3":
763
  # Show numbered list
764
  models_list = list(current_models.keys())
765
  self.console.print("\n[bold]Select model to remove:[/bold]")
766
  for idx, model_name in enumerate(models_list, 1):
767
  self.console.print(f" {idx}. {model_name}")
768
-
769
- model_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(models_list) + 1)])
 
 
 
770
  name = models_list[model_idx - 1]
771
-
772
  if Confirm.ask(f"Remove '{name}'?"):
773
  del current_models[name]
774
-
775
  elif choice == "4":
776
  break
777
-
778
  if current_models:
779
  self.model_mgr.set_models(provider, current_models)
780
  self.console.print(f"\n[green]✅ Models updated for '{provider}'![/green]")
781
  else:
782
- self.console.print("\n[yellow]No models left - removing definition[/yellow]")
 
 
783
  self.model_mgr.remove_models(provider)
784
-
785
  input("\nPress Enter to continue...")
786
-
787
  def view_model_definitions(self, providers: List[str]):
788
  """View model definitions for a provider"""
789
  # Show numbered list
790
  self.console.print("\n[bold]Select provider to view:[/bold]")
791
  for idx, prov in enumerate(providers, 1):
792
  self.console.print(f" {idx}. {prov}")
793
-
794
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers) + 1)])
 
 
795
  provider = providers[choice_idx - 1]
796
-
797
  models = self.model_mgr.get_current_provider_models(provider)
798
  if not models:
799
  self.console.print(f"\n[yellow]No models found for '{provider}'[/yellow]")
800
  input("\nPress Enter to continue...")
801
  return
802
-
803
  clear_screen()
804
  self.console.print(f"[bold]Provider: {provider}[/bold]\n")
805
  self.console.print("[bold]📦 Configured Models:[/bold]")
806
  self.console.print("━" * 50)
807
-
808
  # Handle both dict and list formats
809
  if isinstance(models, dict):
810
  for name, definition in models.items():
@@ -822,74 +966,88 @@ class SettingsTool:
822
  for name in models:
823
  self.console.print(f" Name: {name}")
824
  self.console.print()
825
-
826
  input("Press Enter to return...")
827
-
828
  def manage_provider_settings(self):
829
  """Manage provider-specific settings (Antigravity, Gemini CLI)"""
830
  while True:
831
  clear_screen()
832
-
833
  available_providers = self.provider_settings_mgr.get_available_providers()
834
-
835
- self.console.print(Panel.fit(
836
- "[bold cyan]🔬 Provider-Specific Settings[/bold cyan]",
837
- border_style="cyan"
838
- ))
839
-
 
 
840
  self.console.print()
841
- self.console.print("[bold]📋 Available Providers with Custom Settings[/bold]")
 
 
842
  self.console.print("━" * 70)
843
-
844
  for provider in available_providers:
845
  modified = self.provider_settings_mgr.get_modified_settings(provider)
846
- status = f"[yellow]{len(modified)} modified[/yellow]" if modified else "[dim]defaults[/dim]"
 
 
 
 
847
  display_name = provider.replace("_", " ").title()
848
  self.console.print(f" • {display_name:20} {status}")
849
-
850
  self.console.print()
851
  self.console.print("━" * 70)
852
  self.console.print()
853
  self.console.print("[bold]⚙️ Select Provider to Configure[/bold]")
854
  self.console.print()
855
-
856
  for idx, provider in enumerate(available_providers, 1):
857
  display_name = provider.replace("_", " ").title()
858
  self.console.print(f" {idx}. {display_name}")
859
- self.console.print(f" {len(available_providers) + 1}. ↩️ Back to Settings Menu")
860
-
 
 
861
  self.console.print()
862
  self.console.print("━" * 70)
863
  self.console.print()
864
-
865
  choices = [str(i) for i in range(1, len(available_providers) + 2)]
866
  choice = Prompt.ask("Select option", choices=choices, show_choices=False)
867
  choice_idx = int(choice)
868
-
869
  if choice_idx == len(available_providers) + 1:
870
  break
871
-
872
  provider = available_providers[choice_idx - 1]
873
  self._manage_single_provider_settings(provider)
874
-
875
  def _manage_single_provider_settings(self, provider: str):
876
  """Manage settings for a single provider"""
877
  while True:
878
  clear_screen()
879
-
880
  display_name = provider.replace("_", " ").title()
881
- definitions = self.provider_settings_mgr.get_provider_settings_definitions(provider)
 
 
882
  current_values = self.provider_settings_mgr.get_all_current_values(provider)
883
-
884
- self.console.print(Panel.fit(
885
- f"[bold cyan]🔬 {display_name} Settings[/bold cyan]",
886
- border_style="cyan"
887
- ))
888
-
 
 
889
  self.console.print()
890
  self.console.print("[bold]📋 Current Settings[/bold]")
891
  self.console.print("━" * 70)
892
-
893
  # Display all settings with current values
894
  settings_list = list(definitions.keys())
895
  for idx, key in enumerate(settings_list, 1):
@@ -898,25 +1056,35 @@ class SettingsTool:
898
  default = definition.get("default")
899
  setting_type = definition.get("type", "str")
900
  description = definition.get("description", "")
901
-
902
  # Format value display
903
  if setting_type == "bool":
904
- value_display = "[green]✓ Enabled[/green]" if current else "[red]✗ Disabled[/red]"
 
 
 
 
905
  elif setting_type == "int":
906
  value_display = f"[cyan]{current}[/cyan]"
907
  else:
908
- value_display = f"[cyan]{current or '(not set)'}[/cyan]" if current else "[dim](not set)[/dim]"
909
-
 
 
 
 
910
  # Check if modified from default
911
  modified = current != default
912
  mod_marker = "[yellow]*[/yellow]" if modified else " "
913
-
914
  # Short key name for display (strip provider prefix)
915
  short_key = key.replace(f"{provider.upper()}_", "")
916
-
917
- self.console.print(f" {mod_marker}{idx:2}. {short_key:35} {value_display}")
 
 
918
  self.console.print(f" [dim]{description}[/dim]")
919
-
920
  self.console.print()
921
  self.console.print("━" * 70)
922
  self.console.print("[dim]* = modified from default[/dim]")
@@ -927,13 +1095,17 @@ class SettingsTool:
927
  self.console.print(" R. 🔄 Reset Setting to Default")
928
  self.console.print(" A. 🔄 Reset All to Defaults")
929
  self.console.print(" B. ↩️ Back to Provider Selection")
930
-
931
  self.console.print()
932
  self.console.print("━" * 70)
933
  self.console.print()
934
-
935
- choice = Prompt.ask("Select action", choices=["e", "r", "a", "b", "E", "R", "A", "B"], show_choices=False).lower()
936
-
 
 
 
 
937
  if choice == "b":
938
  break
939
  elif choice == "e":
@@ -942,26 +1114,31 @@ class SettingsTool:
942
  self._reset_provider_setting(provider, settings_list, definitions)
943
  elif choice == "a":
944
  self._reset_all_provider_settings(provider, settings_list)
945
-
946
- def _edit_provider_setting(self, provider: str, settings_list: List[str], definitions: Dict[str, Dict[str, Any]]):
 
 
 
 
 
947
  """Edit a single provider setting"""
948
  self.console.print("\n[bold]Select setting number to edit:[/bold]")
949
-
950
  choices = [str(i) for i in range(1, len(settings_list) + 1)]
951
  choice = IntPrompt.ask("Setting number", choices=choices)
952
  key = settings_list[choice - 1]
953
  definition = definitions[key]
954
-
955
  current = self.provider_settings_mgr.get_current_value(key, definition)
956
  default = definition.get("default")
957
  setting_type = definition.get("type", "str")
958
  short_key = key.replace(f"{provider.upper()}_", "")
959
-
960
  self.console.print(f"\n[bold]Editing: {short_key}[/bold]")
961
  self.console.print(f"Current value: [cyan]{current}[/cyan]")
962
  self.console.print(f"Default value: [dim]{default}[/dim]")
963
  self.console.print(f"Type: {setting_type}")
964
-
965
  if setting_type == "bool":
966
  new_value = Confirm.ask("\nEnable this setting?", default=current)
967
  self.provider_settings_mgr.set_value(key, new_value, definition)
@@ -972,71 +1149,252 @@ class SettingsTool:
972
  self.provider_settings_mgr.set_value(key, new_value, definition)
973
  self.console.print(f"\n[green]✅ {short_key} set to {new_value}![/green]")
974
  else:
975
- new_value = Prompt.ask("\nNew value", default=str(current) if current else "").strip()
 
 
976
  if new_value:
977
  self.provider_settings_mgr.set_value(key, new_value, definition)
978
  self.console.print(f"\n[green]✅ {short_key} updated![/green]")
979
  else:
980
  self.console.print("\n[yellow]No changes made[/yellow]")
981
-
982
  input("\nPress Enter to continue...")
983
-
984
- def _reset_provider_setting(self, provider: str, settings_list: List[str], definitions: Dict[str, Dict[str, Any]]):
 
 
 
 
 
985
  """Reset a single provider setting to default"""
986
  self.console.print("\n[bold]Select setting number to reset:[/bold]")
987
-
988
  choices = [str(i) for i in range(1, len(settings_list) + 1)]
989
  choice = IntPrompt.ask("Setting number", choices=choices)
990
  key = settings_list[choice - 1]
991
  definition = definitions[key]
992
-
993
  default = definition.get("default")
994
  short_key = key.replace(f"{provider.upper()}_", "")
995
-
996
  if Confirm.ask(f"\nReset {short_key} to default ({default})?"):
997
  self.provider_settings_mgr.reset_to_default(key)
998
  self.console.print(f"\n[green]✅ {short_key} reset to default![/green]")
999
  else:
1000
  self.console.print("\n[yellow]No changes made[/yellow]")
1001
-
1002
  input("\nPress Enter to continue...")
1003
-
1004
  def _reset_all_provider_settings(self, provider: str, settings_list: List[str]):
1005
  """Reset all provider settings to defaults"""
1006
  display_name = provider.replace("_", " ").title()
1007
-
1008
- if Confirm.ask(f"\n[bold red]Reset ALL {display_name} settings to defaults?[/bold red]"):
 
 
1009
  for key in settings_list:
1010
  self.provider_settings_mgr.reset_to_default(key)
1011
- self.console.print(f"\n[green]✅ All {display_name} settings reset to defaults![/green]")
 
 
1012
  else:
1013
  self.console.print("\n[yellow]No changes made[/yellow]")
1014
-
1015
  input("\nPress Enter to continue...")
1016
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1017
  def manage_concurrency_limits(self):
1018
  """Manage concurrency limits"""
1019
  while True:
1020
  clear_screen()
1021
-
1022
  limits = self.concurrency_mgr.get_current_limits()
1023
-
1024
- self.console.print(Panel.fit(
1025
- "[bold cyan]⚡ Concurrency Limits Configuration[/bold cyan]",
1026
- border_style="cyan"
1027
- ))
1028
-
 
 
1029
  self.console.print()
1030
  self.console.print("[bold]📋 Current Concurrency Settings[/bold]")
1031
  self.console.print("━" * 70)
1032
-
1033
  if limits:
1034
  for provider, limit in limits.items():
1035
  self.console.print(f" • {provider:15} {limit} requests/key")
1036
  self.console.print(f" • Default: 1 request/key (all others)")
1037
  else:
1038
  self.console.print(" • Default: 1 request/key (all providers)")
1039
-
1040
  self.console.print()
1041
  self.console.print("━" * 70)
1042
  self.console.print()
@@ -1046,96 +1404,128 @@ class SettingsTool:
1046
  self.console.print(" 2. ✏️ Edit Existing Limit")
1047
  self.console.print(" 3. 🗑️ Remove Limit (reset to default)")
1048
  self.console.print(" 4. ↩️ Back to Settings Menu")
1049
-
1050
  self.console.print()
1051
  self.console.print("━" * 70)
1052
  self.console.print()
1053
-
1054
- choice = Prompt.ask("Select option", choices=["1", "2", "3", "4"], show_choices=False)
1055
-
 
 
1056
  if choice == "1":
1057
  # Get available providers
1058
  available_providers = self.get_available_providers()
1059
-
1060
  if not available_providers:
1061
- self.console.print("\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]")
 
 
1062
  input("\nPress Enter to continue...")
1063
  continue
1064
-
1065
  # Show provider selection menu
1066
  self.console.print("\n[bold]Select provider:[/bold]")
1067
  for idx, prov in enumerate(available_providers, 1):
1068
  self.console.print(f" {idx}. {prov}")
1069
- self.console.print(f" {len(available_providers) + 1}. Enter custom provider name")
1070
-
1071
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(available_providers) + 2)])
1072
-
 
 
 
 
 
1073
  if choice_idx == len(available_providers) + 1:
1074
  provider = Prompt.ask("Provider name").strip().lower()
1075
  else:
1076
  provider = available_providers[choice_idx - 1]
1077
-
1078
  if provider:
1079
- limit = IntPrompt.ask("Max concurrent requests per key (1-100)", default=1)
 
 
1080
  if 1 <= limit <= 100:
1081
  self.concurrency_mgr.set_limit(provider, limit)
1082
- self.console.print(f"\n[green]✅ Concurrency limit set for '{provider}': {limit} requests/key[/green]")
 
 
1083
  else:
1084
- self.console.print("\n[red]❌ Limit must be between 1-100[/red]")
 
 
1085
  input("\nPress Enter to continue...")
1086
-
1087
  elif choice == "2":
1088
  if not limits:
1089
  self.console.print("\n[yellow]No limits to edit[/yellow]")
1090
  input("\nPress Enter to continue...")
1091
  continue
1092
-
1093
  # Show numbered list
1094
  self.console.print("\n[bold]Select provider to edit:[/bold]")
1095
  limits_list = list(limits.keys())
1096
  for idx, prov in enumerate(limits_list, 1):
1097
  self.console.print(f" {idx}. {prov}")
1098
-
1099
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(limits_list) + 1)])
 
 
 
1100
  provider = limits_list[choice_idx - 1]
1101
  current_limit = limits.get(provider, 1)
1102
-
1103
  self.console.print(f"\nCurrent limit: {current_limit} requests/key")
1104
- new_limit = IntPrompt.ask("New limit (1-100) [press Enter to keep current]", default=current_limit)
1105
-
 
 
 
1106
  if 1 <= new_limit <= 100:
1107
  if new_limit != current_limit:
1108
  self.concurrency_mgr.set_limit(provider, new_limit)
1109
- self.console.print(f"\n[green]✅ Concurrency limit updated for '{provider}': {new_limit} requests/key[/green]")
 
 
1110
  else:
1111
  self.console.print("\n[yellow]No changes made[/yellow]")
1112
  else:
1113
  self.console.print("\n[red]Limit must be between 1-100[/red]")
1114
  input("\nPress Enter to continue...")
1115
-
1116
  elif choice == "3":
1117
  if not limits:
1118
  self.console.print("\n[yellow]No limits to remove[/yellow]")
1119
  input("\nPress Enter to continue...")
1120
  continue
1121
-
1122
  # Show numbered list
1123
- self.console.print("\n[bold]Select provider to remove limit from:[/bold]")
 
 
1124
  limits_list = list(limits.keys())
1125
  for idx, prov in enumerate(limits_list, 1):
1126
  self.console.print(f" {idx}. {prov}")
1127
-
1128
- choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(limits_list) + 1)])
 
 
 
1129
  provider = limits_list[choice_idx - 1]
1130
-
1131
- if Confirm.ask(f"Remove concurrency limit for '{provider}' (reset to default 1)?"):
 
 
1132
  self.concurrency_mgr.remove_limit(provider)
1133
- self.console.print(f"\n[green]✅ Limit removed for '{provider}' - using default (1 request/key)[/green]")
 
 
1134
  input("\nPress Enter to continue...")
1135
-
1136
  elif choice == "4":
1137
  break
1138
-
1139
  def save_and_exit(self):
1140
  """Save pending changes and exit"""
1141
  if self.settings.has_pending():
@@ -1150,9 +1540,9 @@ class SettingsTool:
1150
  else:
1151
  self.console.print("\n[dim]No changes to save[/dim]")
1152
  input("\nPress Enter to return to launcher...")
1153
-
1154
  self.running = False
1155
-
1156
  def exit_without_saving(self):
1157
  """Exit without saving"""
1158
  if self.settings.has_pending():
 
17
 
18
  def clear_screen():
19
  """
20
+ Cross-platform terminal clear that works robustly on both
21
  classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac).
22
+
23
  Uses native OS commands instead of ANSI escape sequences:
24
  - Windows (conhost & Windows Terminal): cls
25
  - Unix-like systems (Linux, Mac): clear
26
  """
27
+ os.system("cls" if os.name == "nt" else "clear")
28
 
29
 
30
  class AdvancedSettings:
31
  """Manages pending changes to .env"""
32
+
33
  def __init__(self):
34
  self.env_file = Path.cwd() / ".env"
35
  self.pending_changes = {} # key -> value (None means delete)
36
  self.load_current_settings()
37
+
38
  def load_current_settings(self):
39
  """Load current .env values into env vars"""
40
  from dotenv import load_dotenv
41
+
42
  load_dotenv(override=True)
43
+
44
  def set(self, key: str, value: str):
45
  """Stage a change"""
46
  self.pending_changes[key] = value
47
+
48
  def remove(self, key: str):
49
  """Stage a removal"""
50
  self.pending_changes[key] = None
51
+
52
  def save(self):
53
  """Write pending changes to .env"""
54
  for key, value in self.pending_changes.items():
 
58
  else:
59
  # Set key
60
  set_key(str(self.env_file), key, value)
61
+
62
  self.pending_changes.clear()
63
  self.load_current_settings()
64
+
65
  def discard(self):
66
  """Discard pending changes"""
67
  self.pending_changes.clear()
68
+
69
  def has_pending(self) -> bool:
70
  """Check if there are pending changes"""
71
  return bool(self.pending_changes)
 
73
 
74
  class CustomProviderManager:
75
  """Manages custom provider API bases"""
76
+
77
  def __init__(self, settings: AdvancedSettings):
78
  self.settings = settings
79
+
80
  def get_current_providers(self) -> Dict[str, str]:
81
  """Get currently configured custom providers"""
82
  from proxy_app.provider_urls import PROVIDER_URL_MAP
83
+
84
  providers = {}
85
  for key, value in os.environ.items():
86
  if key.endswith("_API_BASE"):
 
89
  if provider not in PROVIDER_URL_MAP:
90
  providers[provider] = value
91
  return providers
92
+
93
  def add_provider(self, name: str, api_base: str):
94
  """Add PROVIDER_API_BASE"""
95
  key = f"{name.upper()}_API_BASE"
96
  self.settings.set(key, api_base)
97
+
98
  def edit_provider(self, name: str, api_base: str):
99
  """Edit PROVIDER_API_BASE"""
100
  self.add_provider(name, api_base)
101
+
102
  def remove_provider(self, name: str):
103
  """Remove PROVIDER_API_BASE"""
104
  key = f"{name.upper()}_API_BASE"
 
107
 
108
  class ModelDefinitionManager:
109
  """Manages PROVIDER_MODELS"""
110
+
111
  def __init__(self, settings: AdvancedSettings):
112
  self.settings = settings
113
+
114
  def get_current_provider_models(self, provider: str) -> Optional[Dict]:
115
  """Get currently configured models for a provider"""
116
  key = f"{provider.upper()}_MODELS"
 
121
  except (json.JSONDecodeError, ValueError):
122
  return None
123
  return None
124
+
125
  def get_all_providers_with_models(self) -> Dict[str, int]:
126
  """Get all providers with model definitions"""
127
  providers = {}
 
137
  except (json.JSONDecodeError, ValueError):
138
  pass
139
  return providers
140
+
141
  def set_models(self, provider: str, models: Dict[str, Dict[str, Any]]):
142
  """Set PROVIDER_MODELS"""
143
  key = f"{provider.upper()}_MODELS"
144
  value = json.dumps(models)
145
  self.settings.set(key, value)
146
+
147
  def remove_models(self, provider: str):
148
  """Remove PROVIDER_MODELS"""
149
  key = f"{provider.upper()}_MODELS"
 
152
 
153
  class ConcurrencyManager:
154
  """Manages MAX_CONCURRENT_REQUESTS_PER_KEY_PROVIDER"""
155
+
156
  def __init__(self, settings: AdvancedSettings):
157
  self.settings = settings
158
+
159
  def get_current_limits(self) -> Dict[str, int]:
160
  """Get currently configured concurrency limits"""
161
  limits = {}
 
167
  except (json.JSONDecodeError, ValueError):
168
  pass
169
  return limits
170
+
171
  def set_limit(self, provider: str, limit: int):
172
  """Set concurrency limit"""
173
  key = f"MAX_CONCURRENT_REQUESTS_PER_KEY_{provider.upper()}"
174
  self.settings.set(key, str(limit))
175
+
176
  def remove_limit(self, provider: str):
177
  """Remove concurrency limit (reset to default)"""
178
  key = f"MAX_CONCURRENT_REQUESTS_PER_KEY_{provider.upper()}"
179
  self.settings.remove(key)
180
 
181
 
182
+ class RotationModeManager:
183
+ """Manages ROTATION_MODE_PROVIDER settings for sequential/balanced credential rotation"""
184
+
185
+ VALID_MODES = ["balanced", "sequential"]
186
+
187
+ def __init__(self, settings: AdvancedSettings):
188
+ self.settings = settings
189
+
190
+ def get_current_modes(self) -> Dict[str, str]:
191
+ """Get currently configured rotation modes"""
192
+ modes = {}
193
+ for key, value in os.environ.items():
194
+ if key.startswith("ROTATION_MODE_"):
195
+ provider = key.replace("ROTATION_MODE_", "").lower()
196
+ if value.lower() in self.VALID_MODES:
197
+ modes[provider] = value.lower()
198
+ return modes
199
+
200
+ def get_default_mode(self, provider: str) -> str:
201
+ """Get the default rotation mode for a provider"""
202
+ # Import here to avoid circular imports
203
+ try:
204
+ from rotator_library.providers.provider_interface import (
205
+ LLMProviderInterface,
206
+ )
207
+
208
+ return LLMProviderInterface.get_rotation_mode(provider)
209
+ except ImportError:
210
+ # Fallback defaults if import fails
211
+ if provider.lower() == "antigravity":
212
+ return "sequential"
213
+ return "balanced"
214
+
215
+ def get_effective_mode(self, provider: str) -> str:
216
+ """Get the effective rotation mode (configured or default)"""
217
+ configured = self.get_current_modes().get(provider.lower())
218
+ if configured:
219
+ return configured
220
+ return self.get_default_mode(provider)
221
+
222
+ def set_mode(self, provider: str, mode: str):
223
+ """Set rotation mode for a provider"""
224
+ if mode.lower() not in self.VALID_MODES:
225
+ raise ValueError(
226
+ f"Invalid rotation mode: {mode}. Must be one of {self.VALID_MODES}"
227
+ )
228
+ key = f"ROTATION_MODE_{provider.upper()}"
229
+ self.settings.set(key, mode.lower())
230
+
231
+ def remove_mode(self, provider: str):
232
+ """Remove rotation mode (reset to provider default)"""
233
+ key = f"ROTATION_MODE_{provider.upper()}"
234
+ self.settings.remove(key)
235
+
236
+
237
  # =============================================================================
238
  # PROVIDER-SPECIFIC SETTINGS DEFINITIONS
239
  # =============================================================================
 
350
 
351
  class ProviderSettingsManager:
352
  """Manages provider-specific configuration settings"""
353
+
354
  def __init__(self, settings: AdvancedSettings):
355
  self.settings = settings
356
+
357
  def get_available_providers(self) -> List[str]:
358
  """Get list of providers with specific settings available"""
359
  return list(PROVIDER_SETTINGS_MAP.keys())
360
+
361
+ def get_provider_settings_definitions(
362
+ self, provider: str
363
+ ) -> Dict[str, Dict[str, Any]]:
364
  """Get settings definitions for a provider"""
365
  return PROVIDER_SETTINGS_MAP.get(provider, {})
366
+
367
  def get_current_value(self, key: str, definition: Dict[str, Any]) -> Any:
368
  """Get current value of a setting from environment"""
369
  env_value = os.getenv(key)
370
  if env_value is None:
371
  return definition.get("default")
372
+
373
  setting_type = definition.get("type", "str")
374
  try:
375
  if setting_type == "bool":
 
380
  return env_value
381
  except (ValueError, AttributeError):
382
  return definition.get("default")
383
+
384
  def get_all_current_values(self, provider: str) -> Dict[str, Any]:
385
  """Get all current values for a provider"""
386
  definitions = self.get_provider_settings_definitions(provider)
 
388
  for key, definition in definitions.items():
389
  values[key] = self.get_current_value(key, definition)
390
  return values
391
+
392
  def set_value(self, key: str, value: Any, definition: Dict[str, Any]):
393
  """Set a setting value, converting to string for .env storage"""
394
  setting_type = definition.get("type", "str")
 
397
  else:
398
  str_value = str(value)
399
  self.settings.set(key, str_value)
400
+
401
  def reset_to_default(self, key: str):
402
  """Remove a setting to reset it to default"""
403
  self.settings.remove(key)
404
+
405
  def get_modified_settings(self, provider: str) -> Dict[str, Any]:
406
  """Get settings that differ from defaults"""
407
  definitions = self.get_provider_settings_definitions(provider)
 
416
 
417
  class SettingsTool:
418
  """Main settings tool TUI"""
419
+
420
  def __init__(self):
421
  self.console = Console()
422
  self.settings = AdvancedSettings()
423
  self.provider_mgr = CustomProviderManager(self.settings)
424
  self.model_mgr = ModelDefinitionManager(self.settings)
425
  self.concurrency_mgr = ConcurrencyManager(self.settings)
426
+ self.rotation_mgr = RotationModeManager(self.settings)
427
  self.provider_settings_mgr = ProviderSettingsManager(self.settings)
428
  self.running = True
429
+
430
  def get_available_providers(self) -> List[str]:
431
  """Get list of providers that have credentials configured"""
432
  env_file = Path.cwd() / ".env"
433
  providers = set()
434
+
435
  # Scan for providers with API keys from local .env
436
  if env_file.exists():
437
  try:
438
+ with open(env_file, "r", encoding="utf-8") as f:
439
  for line in f:
440
  line = line.strip()
441
+ if (
442
+ "_API_KEY" in line
443
+ and "PROXY_API_KEY" not in line
444
+ and "=" in line
445
+ ):
446
  provider = line.split("_API_KEY")[0].strip().lower()
447
  providers.add(provider)
448
  except (IOError, OSError):
449
  pass
450
+
451
  # Also check for OAuth providers from files
452
  oauth_dir = Path("oauth_credentials")
453
  if oauth_dir.exists():
454
  for file in oauth_dir.glob("*_oauth_*.json"):
455
  provider = file.name.split("_oauth_")[0]
456
  providers.add(provider)
457
+
458
  return sorted(list(providers))
459
 
460
  def run(self):
461
  """Main loop"""
462
  while self.running:
463
  self.show_main_menu()
464
+
465
  def show_main_menu(self):
466
  """Display settings categories"""
467
  clear_screen()
468
+
469
+ self.console.print(
470
+ Panel.fit(
471
+ "[bold cyan]🔧 Advanced Settings Configuration[/bold cyan]",
472
+ border_style="cyan",
473
+ )
474
+ )
475
+
476
  self.console.print()
477
  self.console.print("[bold]⚙️ Configuration Categories[/bold]")
478
  self.console.print()
479
  self.console.print(" 1. 🌐 Custom Provider API Bases")
480
  self.console.print(" 2. 📦 Provider Model Definitions")
481
  self.console.print(" 3. ⚡ Concurrency Limits")
482
+ self.console.print(" 4. 🔄 Rotation Modes")
483
+ self.console.print(" 5. 🔬 Provider-Specific Settings")
484
+ self.console.print(" 6. 💾 Save & Exit")
485
+ self.console.print(" 7. 🚫 Exit Without Saving")
486
+
487
  self.console.print()
488
  self.console.print("━" * 70)
489
+
490
  if self.settings.has_pending():
491
+ self.console.print(
492
+ '[yellow]ℹ️ Changes are pending until you select "Save & Exit"[/yellow]'
493
+ )
494
  else:
495
  self.console.print("[dim]ℹ️ No pending changes[/dim]")
496
+
497
  self.console.print()
498
+ self.console.print(
499
+ "[dim]⚠️ Model filters not supported - edit .env for IGNORE_MODELS_* / WHITELIST_MODELS_*[/dim]"
500
+ )
501
  self.console.print()
502
+
503
+ choice = Prompt.ask(
504
+ "Select option",
505
+ choices=["1", "2", "3", "4", "5", "6", "7"],
506
+ show_choices=False,
507
+ )
508
+
509
  if choice == "1":
510
  self.manage_custom_providers()
511
  elif choice == "2":
 
513
  elif choice == "3":
514
  self.manage_concurrency_limits()
515
  elif choice == "4":
516
+ self.manage_rotation_modes()
517
  elif choice == "5":
518
+ self.manage_provider_settings()
519
  elif choice == "6":
520
+ self.save_and_exit()
521
+ elif choice == "7":
522
  self.exit_without_saving()
523
+
524
  def manage_custom_providers(self):
525
  """Manage custom provider API bases"""
526
  while True:
527
  clear_screen()
528
+
529
  providers = self.provider_mgr.get_current_providers()
530
+
531
+ self.console.print(
532
+ Panel.fit(
533
+ "[bold cyan]🌐 Custom Provider API Bases[/bold cyan]",
534
+ border_style="cyan",
535
+ )
536
+ )
537
+
538
  self.console.print()
539
  self.console.print("[bold]📋 Configured Custom Providers[/bold]")
540
  self.console.print("━" * 70)
541
+
542
  if providers:
543
  for name, base in providers.items():
544
  self.console.print(f" • {name:15} {base}")
545
  else:
546
  self.console.print(" [dim]No custom providers configured[/dim]")
547
+
548
  self.console.print()
549
  self.console.print("━" * 70)
550
  self.console.print()
 
554
  self.console.print(" 2. ✏️ Edit Existing Provider")
555
  self.console.print(" 3. 🗑️ Remove Provider")
556
  self.console.print(" 4. ↩️ Back to Settings Menu")
557
+
558
  self.console.print()
559
  self.console.print("━" * 70)
560
  self.console.print()
561
+
562
+ choice = Prompt.ask(
563
+ "Select option", choices=["1", "2", "3", "4"], show_choices=False
564
+ )
565
+
566
  if choice == "1":
567
  name = Prompt.ask("Provider name (e.g., 'opencode')").strip().lower()
568
  if name:
569
  api_base = Prompt.ask("API Base URL").strip()
570
  if api_base:
571
  self.provider_mgr.add_provider(name, api_base)
572
+ self.console.print(
573
+ f"\n[green]✅ Custom provider '{name}' configured![/green]"
574
+ )
575
+ self.console.print(
576
+ f" To use: set {name.upper()}_API_KEY in credentials"
577
+ )
578
  input("\nPress Enter to continue...")
579
+
580
  elif choice == "2":
581
  if not providers:
582
  self.console.print("\n[yellow]No providers to edit[/yellow]")
583
  input("\nPress Enter to continue...")
584
  continue
585
+
586
  # Show numbered list
587
  self.console.print("\n[bold]Select provider to edit:[/bold]")
588
  providers_list = list(providers.keys())
589
  for idx, prov in enumerate(providers_list, 1):
590
  self.console.print(f" {idx}. {prov}")
591
+
592
+ choice_idx = IntPrompt.ask(
593
+ "Select option",
594
+ choices=[str(i) for i in range(1, len(providers_list) + 1)],
595
+ )
596
  name = providers_list[choice_idx - 1]
597
  current_base = providers.get(name, "")
598
+
599
  self.console.print(f"\nCurrent API Base: {current_base}")
600
+ new_base = Prompt.ask(
601
+ "New API Base [press Enter to keep current]", default=current_base
602
+ ).strip()
603
+
604
  if new_base and new_base != current_base:
605
  self.provider_mgr.edit_provider(name, new_base)
606
+ self.console.print(
607
+ f"\n[green]✅ Custom provider '{name}' updated![/green]"
608
+ )
609
  else:
610
  self.console.print("\n[yellow]No changes made[/yellow]")
611
  input("\nPress Enter to continue...")
612
+
613
  elif choice == "3":
614
  if not providers:
615
  self.console.print("\n[yellow]No providers to remove[/yellow]")
616
  input("\nPress Enter to continue...")
617
  continue
618
+
619
  # Show numbered list
620
  self.console.print("\n[bold]Select provider to remove:[/bold]")
621
  providers_list = list(providers.keys())
622
  for idx, prov in enumerate(providers_list, 1):
623
  self.console.print(f" {idx}. {prov}")
624
+
625
+ choice_idx = IntPrompt.ask(
626
+ "Select option",
627
+ choices=[str(i) for i in range(1, len(providers_list) + 1)],
628
+ )
629
  name = providers_list[choice_idx - 1]
630
+
631
  if Confirm.ask(f"Remove '{name}'?"):
632
  self.provider_mgr.remove_provider(name)
633
+ self.console.print(
634
+ f"\n[green]✅ Provider '{name}' removed![/green]"
635
+ )
636
  input("\nPress Enter to continue...")
637
+
638
  elif choice == "4":
639
  break
640
+
641
  def manage_model_definitions(self):
642
  """Manage provider model definitions"""
643
  while True:
644
  clear_screen()
645
+
646
  all_providers = self.model_mgr.get_all_providers_with_models()
647
+
648
+ self.console.print(
649
+ Panel.fit(
650
+ "[bold cyan]📦 Provider Model Definitions[/bold cyan]",
651
+ border_style="cyan",
652
+ )
653
+ )
654
+
655
  self.console.print()
656
  self.console.print("[bold]📋 Configured Provider Models[/bold]")
657
  self.console.print("━" * 70)
658
+
659
  if all_providers:
660
  for provider, count in all_providers.items():
661
+ self.console.print(
662
+ f" • {provider:15} {count} model{'s' if count > 1 else ''}"
663
+ )
664
  else:
665
  self.console.print(" [dim]No model definitions configured[/dim]")
666
+
667
  self.console.print()
668
  self.console.print("━" * 70)
669
  self.console.print()
 
674
  self.console.print(" 3. 👁️ View Provider Models")
675
  self.console.print(" 4. 🗑️ Remove Provider Models")
676
  self.console.print(" 5. ↩️ Back to Settings Menu")
677
+
678
  self.console.print()
679
  self.console.print("━" * 70)
680
  self.console.print()
681
+
682
+ choice = Prompt.ask(
683
+ "Select option", choices=["1", "2", "3", "4", "5"], show_choices=False
684
+ )
685
+
686
  if choice == "1":
687
  self.add_model_definitions()
688
  elif choice == "2":
 
702
  self.console.print("\n[yellow]No providers to remove[/yellow]")
703
  input("\nPress Enter to continue...")
704
  continue
705
+
706
  # Show numbered list
707
+ self.console.print(
708
+ "\n[bold]Select provider to remove models from:[/bold]"
709
+ )
710
  providers_list = list(all_providers.keys())
711
  for idx, prov in enumerate(providers_list, 1):
712
  self.console.print(f" {idx}. {prov}")
713
+
714
+ choice_idx = IntPrompt.ask(
715
+ "Select option",
716
+ choices=[str(i) for i in range(1, len(providers_list) + 1)],
717
+ )
718
  provider = providers_list[choice_idx - 1]
719
+
720
  if Confirm.ask(f"Remove all model definitions for '{provider}'?"):
721
  self.model_mgr.remove_models(provider)
722
+ self.console.print(
723
+ f"\n[green]✅ Model definitions removed for '{provider}'![/green]"
724
+ )
725
  input("\nPress Enter to continue...")
726
  elif choice == "5":
727
  break
728
+
729
  def add_model_definitions(self):
730
  """Add model definitions for a provider"""
731
  # Get available providers from credentials
732
  available_providers = self.get_available_providers()
733
+
734
  if not available_providers:
735
+ self.console.print(
736
+ "\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]"
737
+ )
738
  input("\nPress Enter to continue...")
739
  return
740
+
741
  # Show provider selection menu
742
  self.console.print("\n[bold]Select provider:[/bold]")
743
  for idx, prov in enumerate(available_providers, 1):
744
  self.console.print(f" {idx}. {prov}")
745
+ self.console.print(
746
+ f" {len(available_providers) + 1}. Enter custom provider name"
747
+ )
748
+
749
+ choice = IntPrompt.ask(
750
+ "Select option",
751
+ choices=[str(i) for i in range(1, len(available_providers) + 2)],
752
+ )
753
+
754
  if choice == len(available_providers) + 1:
755
  provider = Prompt.ask("Provider name").strip().lower()
756
  else:
757
  provider = available_providers[choice - 1]
758
+
759
  if not provider:
760
  return
761
+
762
  self.console.print("\nHow would you like to define models?")
763
  self.console.print(" 1. Simple list (names only)")
764
  self.console.print(" 2. Advanced (names with IDs and options)")
765
+
766
  mode = Prompt.ask("Select mode", choices=["1", "2"], show_choices=False)
767
+
768
  models = {}
769
+
770
  if mode == "1":
771
  # Simple mode
772
  while True:
 
783
  break
784
  if name:
785
  model_def = {}
786
+ model_id = Prompt.ask(
787
+ f"Model ID [press Enter to use '{name}']", default=name
788
+ ).strip()
789
  if model_id and model_id != name:
790
  model_def["id"] = model_id
791
+
792
  # Optional: model options
793
+ if Confirm.ask(
794
+ "Add model options (e.g., temperature limits)?", default=False
795
+ ):
796
+ self.console.print(
797
+ "\nEnter options as key=value pairs (one per line, 'done' to finish):"
798
+ )
799
  options = {}
800
  while True:
801
  opt = Prompt.ask("Option").strip()
 
812
  options[key.strip()] = value
813
  if options:
814
  model_def["options"] = options
815
+
816
  models[name] = model_def
817
+
818
  if models:
819
  self.model_mgr.set_models(provider, models)
820
+ self.console.print(
821
+ f"\n[green]✅ Model definitions saved for '{provider}'![/green]"
822
+ )
823
  else:
824
  self.console.print("\n[yellow]No models added[/yellow]")
825
+
826
  input("\nPress Enter to continue...")
827
+
828
  def edit_model_definitions(self, providers: List[str]):
829
  """Edit existing model definitions"""
830
  # Show numbered list
831
  self.console.print("\n[bold]Select provider to edit:[/bold]")
832
  for idx, prov in enumerate(providers, 1):
833
  self.console.print(f" {idx}. {prov}")
834
+
835
+ choice_idx = IntPrompt.ask(
836
+ "Select option", choices=[str(i) for i in range(1, len(providers) + 1)]
837
+ )
838
  provider = providers[choice_idx - 1]
839
+
840
  current_models = self.model_mgr.get_current_provider_models(provider)
841
  if not current_models:
842
  self.console.print(f"\n[yellow]No models found for '{provider}'[/yellow]")
843
  input("\nPress Enter to continue...")
844
  return
845
+
846
  # Convert to dict if list
847
  if isinstance(current_models, list):
848
  current_models = {m: {} for m in current_models}
849
+
850
  while True:
851
  clear_screen()
852
  self.console.print(f"[bold]Editing models for: {provider}[/bold]\n")
853
  self.console.print("Current models:")
854
  for i, (name, definition) in enumerate(current_models.items(), 1):
855
+ model_id = (
856
+ definition.get("id", name) if isinstance(definition, dict) else name
857
+ )
858
  self.console.print(f" {i}. {name} (ID: {model_id})")
859
+
860
  self.console.print("\nOptions:")
861
  self.console.print(" 1. Add new model")
862
  self.console.print(" 2. Edit existing model")
863
  self.console.print(" 3. Remove model")
864
  self.console.print(" 4. Done")
865
+
866
+ choice = Prompt.ask(
867
+ "\nSelect option", choices=["1", "2", "3", "4"], show_choices=False
868
+ )
869
+
870
  if choice == "1":
871
  name = Prompt.ask("New model name").strip()
872
  if name and name not in current_models:
873
  model_id = Prompt.ask("Model ID", default=name).strip()
874
  current_models[name] = {"id": model_id} if model_id != name else {}
875
+
876
  elif choice == "2":
877
  # Show numbered list
878
  models_list = list(current_models.keys())
879
  self.console.print("\n[bold]Select model to edit:[/bold]")
880
  for idx, model_name in enumerate(models_list, 1):
881
  self.console.print(f" {idx}. {model_name}")
882
+
883
+ model_idx = IntPrompt.ask(
884
+ "Select option",
885
+ choices=[str(i) for i in range(1, len(models_list) + 1)],
886
+ )
887
  name = models_list[model_idx - 1]
888
+
889
  current_def = current_models[name]
890
+ current_id = (
891
+ current_def.get("id", name)
892
+ if isinstance(current_def, dict)
893
+ else name
894
+ )
895
+
896
  new_id = Prompt.ask("Model ID", default=current_id).strip()
897
  current_models[name] = {"id": new_id} if new_id != name else {}
898
+
899
  elif choice == "3":
900
  # Show numbered list
901
  models_list = list(current_models.keys())
902
  self.console.print("\n[bold]Select model to remove:[/bold]")
903
  for idx, model_name in enumerate(models_list, 1):
904
  self.console.print(f" {idx}. {model_name}")
905
+
906
+ model_idx = IntPrompt.ask(
907
+ "Select option",
908
+ choices=[str(i) for i in range(1, len(models_list) + 1)],
909
+ )
910
  name = models_list[model_idx - 1]
911
+
912
  if Confirm.ask(f"Remove '{name}'?"):
913
  del current_models[name]
914
+
915
  elif choice == "4":
916
  break
917
+
918
  if current_models:
919
  self.model_mgr.set_models(provider, current_models)
920
  self.console.print(f"\n[green]✅ Models updated for '{provider}'![/green]")
921
  else:
922
+ self.console.print(
923
+ "\n[yellow]No models left - removing definition[/yellow]"
924
+ )
925
  self.model_mgr.remove_models(provider)
926
+
927
  input("\nPress Enter to continue...")
928
+
929
  def view_model_definitions(self, providers: List[str]):
930
  """View model definitions for a provider"""
931
  # Show numbered list
932
  self.console.print("\n[bold]Select provider to view:[/bold]")
933
  for idx, prov in enumerate(providers, 1):
934
  self.console.print(f" {idx}. {prov}")
935
+
936
+ choice_idx = IntPrompt.ask(
937
+ "Select option", choices=[str(i) for i in range(1, len(providers) + 1)]
938
+ )
939
  provider = providers[choice_idx - 1]
940
+
941
  models = self.model_mgr.get_current_provider_models(provider)
942
  if not models:
943
  self.console.print(f"\n[yellow]No models found for '{provider}'[/yellow]")
944
  input("\nPress Enter to continue...")
945
  return
946
+
947
  clear_screen()
948
  self.console.print(f"[bold]Provider: {provider}[/bold]\n")
949
  self.console.print("[bold]📦 Configured Models:[/bold]")
950
  self.console.print("━" * 50)
951
+
952
  # Handle both dict and list formats
953
  if isinstance(models, dict):
954
  for name, definition in models.items():
 
966
  for name in models:
967
  self.console.print(f" Name: {name}")
968
  self.console.print()
969
+
970
  input("Press Enter to return...")
971
+
972
  def manage_provider_settings(self):
973
  """Manage provider-specific settings (Antigravity, Gemini CLI)"""
974
  while True:
975
  clear_screen()
976
+
977
  available_providers = self.provider_settings_mgr.get_available_providers()
978
+
979
+ self.console.print(
980
+ Panel.fit(
981
+ "[bold cyan]🔬 Provider-Specific Settings[/bold cyan]",
982
+ border_style="cyan",
983
+ )
984
+ )
985
+
986
  self.console.print()
987
+ self.console.print(
988
+ "[bold]📋 Available Providers with Custom Settings[/bold]"
989
+ )
990
  self.console.print("━" * 70)
991
+
992
  for provider in available_providers:
993
  modified = self.provider_settings_mgr.get_modified_settings(provider)
994
+ status = (
995
+ f"[yellow]{len(modified)} modified[/yellow]"
996
+ if modified
997
+ else "[dim]defaults[/dim]"
998
+ )
999
  display_name = provider.replace("_", " ").title()
1000
  self.console.print(f" • {display_name:20} {status}")
1001
+
1002
  self.console.print()
1003
  self.console.print("━" * 70)
1004
  self.console.print()
1005
  self.console.print("[bold]⚙️ Select Provider to Configure[/bold]")
1006
  self.console.print()
1007
+
1008
  for idx, provider in enumerate(available_providers, 1):
1009
  display_name = provider.replace("_", " ").title()
1010
  self.console.print(f" {idx}. {display_name}")
1011
+ self.console.print(
1012
+ f" {len(available_providers) + 1}. ↩️ Back to Settings Menu"
1013
+ )
1014
+
1015
  self.console.print()
1016
  self.console.print("━" * 70)
1017
  self.console.print()
1018
+
1019
  choices = [str(i) for i in range(1, len(available_providers) + 2)]
1020
  choice = Prompt.ask("Select option", choices=choices, show_choices=False)
1021
  choice_idx = int(choice)
1022
+
1023
  if choice_idx == len(available_providers) + 1:
1024
  break
1025
+
1026
  provider = available_providers[choice_idx - 1]
1027
  self._manage_single_provider_settings(provider)
1028
+
1029
  def _manage_single_provider_settings(self, provider: str):
1030
  """Manage settings for a single provider"""
1031
  while True:
1032
  clear_screen()
1033
+
1034
  display_name = provider.replace("_", " ").title()
1035
+ definitions = self.provider_settings_mgr.get_provider_settings_definitions(
1036
+ provider
1037
+ )
1038
  current_values = self.provider_settings_mgr.get_all_current_values(provider)
1039
+
1040
+ self.console.print(
1041
+ Panel.fit(
1042
+ f"[bold cyan]🔬 {display_name} Settings[/bold cyan]",
1043
+ border_style="cyan",
1044
+ )
1045
+ )
1046
+
1047
  self.console.print()
1048
  self.console.print("[bold]📋 Current Settings[/bold]")
1049
  self.console.print("━" * 70)
1050
+
1051
  # Display all settings with current values
1052
  settings_list = list(definitions.keys())
1053
  for idx, key in enumerate(settings_list, 1):
 
1056
  default = definition.get("default")
1057
  setting_type = definition.get("type", "str")
1058
  description = definition.get("description", "")
1059
+
1060
  # Format value display
1061
  if setting_type == "bool":
1062
+ value_display = (
1063
+ "[green]✓ Enabled[/green]"
1064
+ if current
1065
+ else "[red]✗ Disabled[/red]"
1066
+ )
1067
  elif setting_type == "int":
1068
  value_display = f"[cyan]{current}[/cyan]"
1069
  else:
1070
+ value_display = (
1071
+ f"[cyan]{current or '(not set)'}[/cyan]"
1072
+ if current
1073
+ else "[dim](not set)[/dim]"
1074
+ )
1075
+
1076
  # Check if modified from default
1077
  modified = current != default
1078
  mod_marker = "[yellow]*[/yellow]" if modified else " "
1079
+
1080
  # Short key name for display (strip provider prefix)
1081
  short_key = key.replace(f"{provider.upper()}_", "")
1082
+
1083
+ self.console.print(
1084
+ f" {mod_marker}{idx:2}. {short_key:35} {value_display}"
1085
+ )
1086
  self.console.print(f" [dim]{description}[/dim]")
1087
+
1088
  self.console.print()
1089
  self.console.print("━" * 70)
1090
  self.console.print("[dim]* = modified from default[/dim]")
 
1095
  self.console.print(" R. 🔄 Reset Setting to Default")
1096
  self.console.print(" A. 🔄 Reset All to Defaults")
1097
  self.console.print(" B. ↩️ Back to Provider Selection")
1098
+
1099
  self.console.print()
1100
  self.console.print("━" * 70)
1101
  self.console.print()
1102
+
1103
+ choice = Prompt.ask(
1104
+ "Select action",
1105
+ choices=["e", "r", "a", "b", "E", "R", "A", "B"],
1106
+ show_choices=False,
1107
+ ).lower()
1108
+
1109
  if choice == "b":
1110
  break
1111
  elif choice == "e":
 
1114
  self._reset_provider_setting(provider, settings_list, definitions)
1115
  elif choice == "a":
1116
  self._reset_all_provider_settings(provider, settings_list)
1117
+
1118
+ def _edit_provider_setting(
1119
+ self,
1120
+ provider: str,
1121
+ settings_list: List[str],
1122
+ definitions: Dict[str, Dict[str, Any]],
1123
+ ):
1124
  """Edit a single provider setting"""
1125
  self.console.print("\n[bold]Select setting number to edit:[/bold]")
1126
+
1127
  choices = [str(i) for i in range(1, len(settings_list) + 1)]
1128
  choice = IntPrompt.ask("Setting number", choices=choices)
1129
  key = settings_list[choice - 1]
1130
  definition = definitions[key]
1131
+
1132
  current = self.provider_settings_mgr.get_current_value(key, definition)
1133
  default = definition.get("default")
1134
  setting_type = definition.get("type", "str")
1135
  short_key = key.replace(f"{provider.upper()}_", "")
1136
+
1137
  self.console.print(f"\n[bold]Editing: {short_key}[/bold]")
1138
  self.console.print(f"Current value: [cyan]{current}[/cyan]")
1139
  self.console.print(f"Default value: [dim]{default}[/dim]")
1140
  self.console.print(f"Type: {setting_type}")
1141
+
1142
  if setting_type == "bool":
1143
  new_value = Confirm.ask("\nEnable this setting?", default=current)
1144
  self.provider_settings_mgr.set_value(key, new_value, definition)
 
1149
  self.provider_settings_mgr.set_value(key, new_value, definition)
1150
  self.console.print(f"\n[green]✅ {short_key} set to {new_value}![/green]")
1151
  else:
1152
+ new_value = Prompt.ask(
1153
+ "\nNew value", default=str(current) if current else ""
1154
+ ).strip()
1155
  if new_value:
1156
  self.provider_settings_mgr.set_value(key, new_value, definition)
1157
  self.console.print(f"\n[green]✅ {short_key} updated![/green]")
1158
  else:
1159
  self.console.print("\n[yellow]No changes made[/yellow]")
1160
+
1161
  input("\nPress Enter to continue...")
1162
+
1163
+ def _reset_provider_setting(
1164
+ self,
1165
+ provider: str,
1166
+ settings_list: List[str],
1167
+ definitions: Dict[str, Dict[str, Any]],
1168
+ ):
1169
  """Reset a single provider setting to default"""
1170
  self.console.print("\n[bold]Select setting number to reset:[/bold]")
1171
+
1172
  choices = [str(i) for i in range(1, len(settings_list) + 1)]
1173
  choice = IntPrompt.ask("Setting number", choices=choices)
1174
  key = settings_list[choice - 1]
1175
  definition = definitions[key]
1176
+
1177
  default = definition.get("default")
1178
  short_key = key.replace(f"{provider.upper()}_", "")
1179
+
1180
  if Confirm.ask(f"\nReset {short_key} to default ({default})?"):
1181
  self.provider_settings_mgr.reset_to_default(key)
1182
  self.console.print(f"\n[green]✅ {short_key} reset to default![/green]")
1183
  else:
1184
  self.console.print("\n[yellow]No changes made[/yellow]")
1185
+
1186
  input("\nPress Enter to continue...")
1187
+
1188
  def _reset_all_provider_settings(self, provider: str, settings_list: List[str]):
1189
  """Reset all provider settings to defaults"""
1190
  display_name = provider.replace("_", " ").title()
1191
+
1192
+ if Confirm.ask(
1193
+ f"\n[bold red]Reset ALL {display_name} settings to defaults?[/bold red]"
1194
+ ):
1195
  for key in settings_list:
1196
  self.provider_settings_mgr.reset_to_default(key)
1197
+ self.console.print(
1198
+ f"\n[green]✅ All {display_name} settings reset to defaults![/green]"
1199
+ )
1200
  else:
1201
  self.console.print("\n[yellow]No changes made[/yellow]")
1202
+
1203
  input("\nPress Enter to continue...")
1204
+
1205
+ def manage_rotation_modes(self):
1206
+ """Manage credential rotation modes (sequential vs balanced)"""
1207
+ while True:
1208
+ clear_screen()
1209
+
1210
+ modes = self.rotation_mgr.get_current_modes()
1211
+ available_providers = self.get_available_providers()
1212
+
1213
+ self.console.print(
1214
+ Panel.fit(
1215
+ "[bold cyan]🔄 Credential Rotation Mode Configuration[/bold cyan]",
1216
+ border_style="cyan",
1217
+ )
1218
+ )
1219
+
1220
+ self.console.print()
1221
+ self.console.print("[bold]📋 Rotation Modes Explained[/bold]")
1222
+ self.console.print("━" * 70)
1223
+ self.console.print(
1224
+ " [cyan]balanced[/cyan] - Rotate credentials evenly across requests (default)"
1225
+ )
1226
+ self.console.print(
1227
+ " [cyan]sequential[/cyan] - Use one credential until exhausted (429), then switch"
1228
+ )
1229
+ self.console.print()
1230
+ self.console.print("[bold]📋 Current Rotation Mode Settings[/bold]")
1231
+ self.console.print("━" * 70)
1232
+
1233
+ if modes:
1234
+ for provider, mode in modes.items():
1235
+ default_mode = self.rotation_mgr.get_default_mode(provider)
1236
+ is_custom = mode != default_mode
1237
+ marker = "[yellow]*[/yellow]" if is_custom else " "
1238
+ mode_display = (
1239
+ f"[green]{mode}[/green]"
1240
+ if mode == "sequential"
1241
+ else f"[blue]{mode}[/blue]"
1242
+ )
1243
+ self.console.print(f" {marker}• {provider:20} {mode_display}")
1244
+
1245
+ # Show providers with default modes
1246
+ providers_with_defaults = [p for p in available_providers if p not in modes]
1247
+ if providers_with_defaults:
1248
+ self.console.print()
1249
+ self.console.print("[dim]Providers using default modes:[/dim]")
1250
+ for provider in providers_with_defaults:
1251
+ default_mode = self.rotation_mgr.get_default_mode(provider)
1252
+ mode_display = (
1253
+ f"[green]{default_mode}[/green]"
1254
+ if default_mode == "sequential"
1255
+ else f"[blue]{default_mode}[/blue]"
1256
+ )
1257
+ self.console.print(
1258
+ f" • {provider:20} {mode_display} [dim](default)[/dim]"
1259
+ )
1260
+
1261
+ self.console.print()
1262
+ self.console.print("━" * 70)
1263
+ self.console.print(
1264
+ "[dim]* = custom setting (differs from provider default)[/dim]"
1265
+ )
1266
+ self.console.print()
1267
+ self.console.print("[bold]⚙️ Actions[/bold]")
1268
+ self.console.print()
1269
+ self.console.print(" 1. ➕ Set Rotation Mode for Provider")
1270
+ self.console.print(" 2. 🗑️ Reset to Provider Default")
1271
+ self.console.print(" 3. ↩️ Back to Settings Menu")
1272
+
1273
+ self.console.print()
1274
+ self.console.print("━" * 70)
1275
+ self.console.print()
1276
+
1277
+ choice = Prompt.ask(
1278
+ "Select option", choices=["1", "2", "3"], show_choices=False
1279
+ )
1280
+
1281
+ if choice == "1":
1282
+ if not available_providers:
1283
+ self.console.print(
1284
+ "\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]"
1285
+ )
1286
+ input("\nPress Enter to continue...")
1287
+ continue
1288
+
1289
+ # Show provider selection menu
1290
+ self.console.print("\n[bold]Select provider:[/bold]")
1291
+ for idx, prov in enumerate(available_providers, 1):
1292
+ current_mode = self.rotation_mgr.get_effective_mode(prov)
1293
+ mode_display = (
1294
+ f"[green]{current_mode}[/green]"
1295
+ if current_mode == "sequential"
1296
+ else f"[blue]{current_mode}[/blue]"
1297
+ )
1298
+ self.console.print(f" {idx}. {prov} ({mode_display})")
1299
+ self.console.print(
1300
+ f" {len(available_providers) + 1}. Enter custom provider name"
1301
+ )
1302
+
1303
+ choice_idx = IntPrompt.ask(
1304
+ "Select option",
1305
+ choices=[str(i) for i in range(1, len(available_providers) + 2)],
1306
+ )
1307
+
1308
+ if choice_idx == len(available_providers) + 1:
1309
+ provider = Prompt.ask("Provider name").strip().lower()
1310
+ else:
1311
+ provider = available_providers[choice_idx - 1]
1312
+
1313
+ if provider:
1314
+ current_mode = self.rotation_mgr.get_effective_mode(provider)
1315
+ self.console.print(
1316
+ f"\nCurrent mode for {provider}: [cyan]{current_mode}[/cyan]"
1317
+ )
1318
+ self.console.print("\nSelect new rotation mode:")
1319
+ self.console.print(
1320
+ " 1. [blue]balanced[/blue] - Rotate credentials evenly"
1321
+ )
1322
+ self.console.print(
1323
+ " 2. [green]sequential[/green] - Use until exhausted"
1324
+ )
1325
+
1326
+ mode_choice = Prompt.ask(
1327
+ "Select mode", choices=["1", "2"], show_choices=False
1328
+ )
1329
+ new_mode = "balanced" if mode_choice == "1" else "sequential"
1330
+
1331
+ self.rotation_mgr.set_mode(provider, new_mode)
1332
+ self.console.print(
1333
+ f"\n[green]✅ Rotation mode for '{provider}' set to {new_mode}![/green]"
1334
+ )
1335
+ input("\nPress Enter to continue...")
1336
+
1337
+ elif choice == "2":
1338
+ if not modes:
1339
+ self.console.print(
1340
+ "\n[yellow]No custom rotation modes to reset[/yellow]"
1341
+ )
1342
+ input("\nPress Enter to continue...")
1343
+ continue
1344
+
1345
+ # Show numbered list
1346
+ self.console.print(
1347
+ "\n[bold]Select provider to reset to default:[/bold]"
1348
+ )
1349
+ modes_list = list(modes.keys())
1350
+ for idx, prov in enumerate(modes_list, 1):
1351
+ default_mode = self.rotation_mgr.get_default_mode(prov)
1352
+ self.console.print(
1353
+ f" {idx}. {prov} (will reset to: {default_mode})"
1354
+ )
1355
+
1356
+ choice_idx = IntPrompt.ask(
1357
+ "Select option",
1358
+ choices=[str(i) for i in range(1, len(modes_list) + 1)],
1359
+ )
1360
+ provider = modes_list[choice_idx - 1]
1361
+ default_mode = self.rotation_mgr.get_default_mode(provider)
1362
+
1363
+ if Confirm.ask(f"Reset '{provider}' to default mode ({default_mode})?"):
1364
+ self.rotation_mgr.remove_mode(provider)
1365
+ self.console.print(
1366
+ f"\n[green]✅ Rotation mode for '{provider}' reset to default ({default_mode})![/green]"
1367
+ )
1368
+ input("\nPress Enter to continue...")
1369
+
1370
+ elif choice == "3":
1371
+ break
1372
+
1373
  def manage_concurrency_limits(self):
1374
  """Manage concurrency limits"""
1375
  while True:
1376
  clear_screen()
1377
+
1378
  limits = self.concurrency_mgr.get_current_limits()
1379
+
1380
+ self.console.print(
1381
+ Panel.fit(
1382
+ "[bold cyan]⚡ Concurrency Limits Configuration[/bold cyan]",
1383
+ border_style="cyan",
1384
+ )
1385
+ )
1386
+
1387
  self.console.print()
1388
  self.console.print("[bold]📋 Current Concurrency Settings[/bold]")
1389
  self.console.print("━" * 70)
1390
+
1391
  if limits:
1392
  for provider, limit in limits.items():
1393
  self.console.print(f" • {provider:15} {limit} requests/key")
1394
  self.console.print(f" • Default: 1 request/key (all others)")
1395
  else:
1396
  self.console.print(" • Default: 1 request/key (all providers)")
1397
+
1398
  self.console.print()
1399
  self.console.print("━" * 70)
1400
  self.console.print()
 
1404
  self.console.print(" 2. ✏️ Edit Existing Limit")
1405
  self.console.print(" 3. 🗑️ Remove Limit (reset to default)")
1406
  self.console.print(" 4. ↩️ Back to Settings Menu")
1407
+
1408
  self.console.print()
1409
  self.console.print("━" * 70)
1410
  self.console.print()
1411
+
1412
+ choice = Prompt.ask(
1413
+ "Select option", choices=["1", "2", "3", "4"], show_choices=False
1414
+ )
1415
+
1416
  if choice == "1":
1417
  # Get available providers
1418
  available_providers = self.get_available_providers()
1419
+
1420
  if not available_providers:
1421
+ self.console.print(
1422
+ "\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]"
1423
+ )
1424
  input("\nPress Enter to continue...")
1425
  continue
1426
+
1427
  # Show provider selection menu
1428
  self.console.print("\n[bold]Select provider:[/bold]")
1429
  for idx, prov in enumerate(available_providers, 1):
1430
  self.console.print(f" {idx}. {prov}")
1431
+ self.console.print(
1432
+ f" {len(available_providers) + 1}. Enter custom provider name"
1433
+ )
1434
+
1435
+ choice_idx = IntPrompt.ask(
1436
+ "Select option",
1437
+ choices=[str(i) for i in range(1, len(available_providers) + 2)],
1438
+ )
1439
+
1440
  if choice_idx == len(available_providers) + 1:
1441
  provider = Prompt.ask("Provider name").strip().lower()
1442
  else:
1443
  provider = available_providers[choice_idx - 1]
1444
+
1445
  if provider:
1446
+ limit = IntPrompt.ask(
1447
+ "Max concurrent requests per key (1-100)", default=1
1448
+ )
1449
  if 1 <= limit <= 100:
1450
  self.concurrency_mgr.set_limit(provider, limit)
1451
+ self.console.print(
1452
+ f"\n[green]✅ Concurrency limit set for '{provider}': {limit} requests/key[/green]"
1453
+ )
1454
  else:
1455
+ self.console.print(
1456
+ "\n[red]❌ Limit must be between 1-100[/red]"
1457
+ )
1458
  input("\nPress Enter to continue...")
1459
+
1460
  elif choice == "2":
1461
  if not limits:
1462
  self.console.print("\n[yellow]No limits to edit[/yellow]")
1463
  input("\nPress Enter to continue...")
1464
  continue
1465
+
1466
  # Show numbered list
1467
  self.console.print("\n[bold]Select provider to edit:[/bold]")
1468
  limits_list = list(limits.keys())
1469
  for idx, prov in enumerate(limits_list, 1):
1470
  self.console.print(f" {idx}. {prov}")
1471
+
1472
+ choice_idx = IntPrompt.ask(
1473
+ "Select option",
1474
+ choices=[str(i) for i in range(1, len(limits_list) + 1)],
1475
+ )
1476
  provider = limits_list[choice_idx - 1]
1477
  current_limit = limits.get(provider, 1)
1478
+
1479
  self.console.print(f"\nCurrent limit: {current_limit} requests/key")
1480
+ new_limit = IntPrompt.ask(
1481
+ "New limit (1-100) [press Enter to keep current]",
1482
+ default=current_limit,
1483
+ )
1484
+
1485
  if 1 <= new_limit <= 100:
1486
  if new_limit != current_limit:
1487
  self.concurrency_mgr.set_limit(provider, new_limit)
1488
+ self.console.print(
1489
+ f"\n[green]✅ Concurrency limit updated for '{provider}': {new_limit} requests/key[/green]"
1490
+ )
1491
  else:
1492
  self.console.print("\n[yellow]No changes made[/yellow]")
1493
  else:
1494
  self.console.print("\n[red]Limit must be between 1-100[/red]")
1495
  input("\nPress Enter to continue...")
1496
+
1497
  elif choice == "3":
1498
  if not limits:
1499
  self.console.print("\n[yellow]No limits to remove[/yellow]")
1500
  input("\nPress Enter to continue...")
1501
  continue
1502
+
1503
  # Show numbered list
1504
+ self.console.print(
1505
+ "\n[bold]Select provider to remove limit from:[/bold]"
1506
+ )
1507
  limits_list = list(limits.keys())
1508
  for idx, prov in enumerate(limits_list, 1):
1509
  self.console.print(f" {idx}. {prov}")
1510
+
1511
+ choice_idx = IntPrompt.ask(
1512
+ "Select option",
1513
+ choices=[str(i) for i in range(1, len(limits_list) + 1)],
1514
+ )
1515
  provider = limits_list[choice_idx - 1]
1516
+
1517
+ if Confirm.ask(
1518
+ f"Remove concurrency limit for '{provider}' (reset to default 1)?"
1519
+ ):
1520
  self.concurrency_mgr.remove_limit(provider)
1521
+ self.console.print(
1522
+ f"\n[green]✅ Limit removed for '{provider}' - using default (1 request/key)[/green]"
1523
+ )
1524
  input("\nPress Enter to continue...")
1525
+
1526
  elif choice == "4":
1527
  break
1528
+
1529
  def save_and_exit(self):
1530
  """Save pending changes and exit"""
1531
  if self.settings.has_pending():
 
1540
  else:
1541
  self.console.print("\n[dim]No changes to save[/dim]")
1542
  input("\nPress Enter to return to launcher...")
1543
+
1544
  self.running = False
1545
+
1546
  def exit_without_saving(self):
1547
  """Exit without saving"""
1548
  if self.settings.has_pending():
src/rotator_library/client.py CHANGED
@@ -139,8 +139,28 @@ class RotatingClient:
139
  self.max_retries = max_retries
140
  self.global_timeout = global_timeout
141
  self.abort_on_callback_error = abort_on_callback_error
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
142
  self.usage_manager = UsageManager(
143
- file_path=usage_file_path, rotation_tolerance=rotation_tolerance
 
 
144
  )
145
  self._model_list_cache = {}
146
  self._provider_plugins = PROVIDER_PLUGINS
@@ -1070,7 +1090,7 @@ class RotatingClient:
1070
  if request
1071
  else {},
1072
  )
1073
- classified_error = classify_error(e)
1074
 
1075
  # Extract a clean error message for the user-facing log
1076
  error_message = str(e).split("\n")[0]
@@ -1114,7 +1134,7 @@ class RotatingClient:
1114
  if request
1115
  else {},
1116
  )
1117
- classified_error = classify_error(e)
1118
  error_message = str(e).split("\n")[0]
1119
 
1120
  # Provider-level error: don't increment consecutive failures
@@ -1170,7 +1190,7 @@ class RotatingClient:
1170
  else {},
1171
  )
1172
 
1173
- classified_error = classify_error(e)
1174
  error_message = str(e).split("\n")[0]
1175
 
1176
  lib_logger.warning(
@@ -1239,7 +1259,7 @@ class RotatingClient:
1239
  )
1240
  raise last_exception
1241
 
1242
- classified_error = classify_error(e)
1243
  error_message = str(e).split("\n")[0]
1244
 
1245
  lib_logger.warning(
@@ -1566,7 +1586,9 @@ class RotatingClient:
1566
  last_exception = e
1567
  # If the exception is our custom wrapper, unwrap the original error
1568
  original_exc = getattr(e, "data", e)
1569
- classified_error = classify_error(original_exc)
 
 
1570
  error_message = str(original_exc).split("\n")[0]
1571
 
1572
  log_failure(
@@ -1623,7 +1645,7 @@ class RotatingClient:
1623
  if request
1624
  else {},
1625
  )
1626
- classified_error = classify_error(e)
1627
  error_message = str(e).split("\n")[0]
1628
 
1629
  # Provider-level error: don't increment consecutive failures
@@ -1673,7 +1695,7 @@ class RotatingClient:
1673
  if request
1674
  else {},
1675
  )
1676
- classified_error = classify_error(e)
1677
  error_message = str(e).split("\n")[0]
1678
 
1679
  # Record in accumulator
@@ -1812,7 +1834,9 @@ class RotatingClient:
1812
  cleaned_str = None
1813
  # The actual exception might be wrapped in our StreamedAPIError.
1814
  original_exc = getattr(e, "data", e)
1815
- classified_error = classify_error(original_exc)
 
 
1816
 
1817
  # Check if this error should trigger rotation
1818
  if not should_rotate_on_error(classified_error):
@@ -1939,7 +1963,7 @@ class RotatingClient:
1939
  if request
1940
  else {},
1941
  )
1942
- classified_error = classify_error(e)
1943
  error_message_text = str(e).split("\n")[0]
1944
 
1945
  # Record error in accumulator (server errors are transient, not abnormal)
@@ -1990,7 +2014,7 @@ class RotatingClient:
1990
  if request
1991
  else {},
1992
  )
1993
- classified_error = classify_error(e)
1994
  error_message_text = str(e).split("\n")[0]
1995
 
1996
  # Record error in accumulator
@@ -2232,7 +2256,7 @@ class RotatingClient:
2232
  self._model_list_cache[provider] = final_models
2233
  return final_models
2234
  except Exception as e:
2235
- classified_error = classify_error(e)
2236
  cred_display = mask_credential(credential)
2237
  lib_logger.debug(
2238
  f"Failed to get models for provider {provider} with credential {cred_display}: {classified_error.error_type}. Trying next credential."
 
139
  self.max_retries = max_retries
140
  self.global_timeout = global_timeout
141
  self.abort_on_callback_error = abort_on_callback_error
142
+
143
+ # Build provider rotation modes map
144
+ # Each provider can specify its preferred rotation mode ("balanced" or "sequential")
145
+ provider_rotation_modes = {}
146
+ for provider in self.all_credentials.keys():
147
+ provider_class = self._provider_plugins.get(provider)
148
+ if provider_class and hasattr(provider_class, "get_rotation_mode"):
149
+ # Use class method to get rotation mode (checks env var + class default)
150
+ mode = provider_class.get_rotation_mode(provider)
151
+ else:
152
+ # Fallback: check environment variable directly
153
+ env_key = f"ROTATION_MODE_{provider.upper()}"
154
+ mode = os.getenv(env_key, "balanced")
155
+
156
+ provider_rotation_modes[provider] = mode
157
+ if mode != "balanced":
158
+ lib_logger.info(f"Provider '{provider}' using rotation mode: {mode}")
159
+
160
  self.usage_manager = UsageManager(
161
+ file_path=usage_file_path,
162
+ rotation_tolerance=rotation_tolerance,
163
+ provider_rotation_modes=provider_rotation_modes,
164
  )
165
  self._model_list_cache = {}
166
  self._provider_plugins = PROVIDER_PLUGINS
 
1090
  if request
1091
  else {},
1092
  )
1093
+ classified_error = classify_error(e, provider=provider)
1094
 
1095
  # Extract a clean error message for the user-facing log
1096
  error_message = str(e).split("\n")[0]
 
1134
  if request
1135
  else {},
1136
  )
1137
+ classified_error = classify_error(e, provider=provider)
1138
  error_message = str(e).split("\n")[0]
1139
 
1140
  # Provider-level error: don't increment consecutive failures
 
1190
  else {},
1191
  )
1192
 
1193
+ classified_error = classify_error(e, provider=provider)
1194
  error_message = str(e).split("\n")[0]
1195
 
1196
  lib_logger.warning(
 
1259
  )
1260
  raise last_exception
1261
 
1262
+ classified_error = classify_error(e, provider=provider)
1263
  error_message = str(e).split("\n")[0]
1264
 
1265
  lib_logger.warning(
 
1586
  last_exception = e
1587
  # If the exception is our custom wrapper, unwrap the original error
1588
  original_exc = getattr(e, "data", e)
1589
+ classified_error = classify_error(
1590
+ original_exc, provider=provider
1591
+ )
1592
  error_message = str(original_exc).split("\n")[0]
1593
 
1594
  log_failure(
 
1645
  if request
1646
  else {},
1647
  )
1648
+ classified_error = classify_error(e, provider=provider)
1649
  error_message = str(e).split("\n")[0]
1650
 
1651
  # Provider-level error: don't increment consecutive failures
 
1695
  if request
1696
  else {},
1697
  )
1698
+ classified_error = classify_error(e, provider=provider)
1699
  error_message = str(e).split("\n")[0]
1700
 
1701
  # Record in accumulator
 
1834
  cleaned_str = None
1835
  # The actual exception might be wrapped in our StreamedAPIError.
1836
  original_exc = getattr(e, "data", e)
1837
+ classified_error = classify_error(
1838
+ original_exc, provider=provider
1839
+ )
1840
 
1841
  # Check if this error should trigger rotation
1842
  if not should_rotate_on_error(classified_error):
 
1963
  if request
1964
  else {},
1965
  )
1966
+ classified_error = classify_error(e, provider=provider)
1967
  error_message_text = str(e).split("\n")[0]
1968
 
1969
  # Record error in accumulator (server errors are transient, not abnormal)
 
2014
  if request
2015
  else {},
2016
  )
2017
+ classified_error = classify_error(e, provider=provider)
2018
  error_message_text = str(e).split("\n")[0]
2019
 
2020
  # Record error in accumulator
 
2256
  self._model_list_cache[provider] = final_models
2257
  return final_models
2258
  except Exception as e:
2259
+ classified_error = classify_error(e, provider=provider)
2260
  cred_display = mask_credential(credential)
2261
  lib_logger.debug(
2262
  f"Failed to get models for provider {provider} with credential {cred_display}: {classified_error.error_type}. Trying next credential."
src/rotator_library/error_handler.py CHANGED
@@ -1,6 +1,7 @@
1
  import re
2
  import json
3
  import os
 
4
  from typing import Optional, Dict, Any
5
  import httpx
6
 
@@ -17,6 +18,8 @@ from litellm.exceptions import (
17
  ContextWindowExceededError,
18
  )
19
 
 
 
20
 
21
  def _parse_duration_string(duration_str: str) -> Optional[int]:
22
  """
@@ -513,11 +516,15 @@ def get_retry_after(error: Exception) -> Optional[int]:
513
  return None
514
 
515
 
516
- def classify_error(e: Exception) -> ClassifiedError:
517
  """
518
  Classifies an exception into a structured ClassifiedError object.
519
  Now handles both litellm and httpx exceptions.
520
 
 
 
 
 
521
  Error types and their typical handling:
522
  - rate_limit (429): Rotate key, may retry with backoff
523
  - server_error (5xx): Retry with backoff, then rotate
@@ -528,7 +535,60 @@ def classify_error(e: Exception) -> ClassifiedError:
528
  - context_window_exceeded: Don't retry - request too large
529
  - api_connection: Retry with backoff, then rotate
530
  - unknown: Rotate key (safer to try another)
 
 
 
 
 
 
 
531
  """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
532
  status_code = getattr(e, "status_code", None)
533
 
534
  if isinstance(e, httpx.HTTPStatusError): # [NEW] Handle httpx errors first
 
1
  import re
2
  import json
3
  import os
4
+ import logging
5
  from typing import Optional, Dict, Any
6
  import httpx
7
 
 
18
  ContextWindowExceededError,
19
  )
20
 
21
+ lib_logger = logging.getLogger("rotator_library")
22
+
23
 
24
  def _parse_duration_string(duration_str: str) -> Optional[int]:
25
  """
 
516
  return None
517
 
518
 
519
+ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedError:
520
  """
521
  Classifies an exception into a structured ClassifiedError object.
522
  Now handles both litellm and httpx exceptions.
523
 
524
+ If provider is specified and has a parse_quota_error() method,
525
+ attempts provider-specific error parsing first before falling back
526
+ to generic classification.
527
+
528
  Error types and their typical handling:
529
  - rate_limit (429): Rotate key, may retry with backoff
530
  - server_error (5xx): Retry with backoff, then rotate
 
535
  - context_window_exceeded: Don't retry - request too large
536
  - api_connection: Retry with backoff, then rotate
537
  - unknown: Rotate key (safer to try another)
538
+
539
+ Args:
540
+ e: The exception to classify
541
+ provider: Optional provider name for provider-specific error parsing
542
+
543
+ Returns:
544
+ ClassifiedError with error_type, status_code, retry_after, etc.
545
  """
546
+ # Try provider-specific parsing first for 429/rate limit errors
547
+ if provider:
548
+ try:
549
+ from .providers import PROVIDER_PLUGINS
550
+
551
+ provider_class = PROVIDER_PLUGINS.get(provider)
552
+
553
+ if provider_class and hasattr(provider_class, "parse_quota_error"):
554
+ # Get error body if available
555
+ error_body = None
556
+ if hasattr(e, "response") and hasattr(e.response, "text"):
557
+ try:
558
+ error_body = e.response.text
559
+ except Exception:
560
+ pass
561
+ elif hasattr(e, "body"):
562
+ error_body = str(e.body)
563
+
564
+ quota_info = provider_class.parse_quota_error(e, error_body)
565
+
566
+ if quota_info and quota_info.get("retry_after"):
567
+ retry_after = quota_info["retry_after"]
568
+ reason = quota_info.get("reason", "QUOTA_EXHAUSTED")
569
+ reset_ts = quota_info.get("reset_timestamp")
570
+
571
+ # Log the parsed result with human-readable duration
572
+ hours = retry_after / 3600
573
+ lib_logger.info(
574
+ f"Provider '{provider}' parsed quota error: "
575
+ f"retry_after={retry_after}s ({hours:.1f}h), reason={reason}"
576
+ + (f", resets at {reset_ts}" if reset_ts else "")
577
+ )
578
+
579
+ return ClassifiedError(
580
+ error_type="quota_exceeded",
581
+ original_exception=e,
582
+ status_code=429,
583
+ retry_after=retry_after,
584
+ )
585
+ except Exception as parse_error:
586
+ lib_logger.debug(
587
+ f"Provider-specific error parsing failed for '{provider}': {parse_error}"
588
+ )
589
+ # Fall through to generic classification
590
+
591
+ # Generic classification logic
592
  status_code = getattr(e, "status_code", None)
593
 
594
  if isinstance(e, httpx.HTTPStatusError): # [NEW] Handle httpx errors first
src/rotator_library/providers/antigravity_provider.py CHANGED
@@ -494,6 +494,147 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface):
494
 
495
  skip_cost_calculation = True
496
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
497
  def __init__(self):
498
  super().__init__()
499
  self.model_definitions = ModelDefinitions()
 
494
 
495
  skip_cost_calculation = True
496
 
497
+ # Sequential mode by default - preserves thinking signature caches between requests
498
+ default_rotation_mode: str = "sequential"
499
+
500
+ @staticmethod
501
+ def parse_quota_error(
502
+ error: Exception, error_body: Optional[str] = None
503
+ ) -> Optional[Dict[str, Any]]:
504
+ """
505
+ Parse Antigravity/Google RPC quota errors.
506
+
507
+ Handles the Google Cloud API error format with ErrorInfo and RetryInfo details.
508
+
509
+ Example error format:
510
+ {
511
+ "error": {
512
+ "code": 429,
513
+ "details": [
514
+ {
515
+ "@type": "type.googleapis.com/google.rpc.ErrorInfo",
516
+ "reason": "QUOTA_EXHAUSTED",
517
+ "metadata": {
518
+ "quotaResetDelay": "143h4m52.730699158s",
519
+ "quotaResetTimeStamp": "2025-12-11T22:53:16Z"
520
+ }
521
+ },
522
+ {
523
+ "@type": "type.googleapis.com/google.rpc.RetryInfo",
524
+ "retryDelay": "515092.730699158s"
525
+ }
526
+ ]
527
+ }
528
+ }
529
+
530
+ Args:
531
+ error: The caught exception
532
+ error_body: Optional raw response body string
533
+
534
+ Returns:
535
+ None if not a parseable quota error, otherwise:
536
+ {
537
+ "retry_after": int,
538
+ "reason": str,
539
+ "reset_timestamp": str | None,
540
+ }
541
+ """
542
+ import re as regex_module
543
+
544
+ def parse_duration(duration_str: str) -> Optional[int]:
545
+ """Parse duration strings like '143h4m52.73s' or '515092.73s' to seconds."""
546
+ if not duration_str:
547
+ return None
548
+
549
+ # Handle pure seconds format: "515092.730699158s"
550
+ pure_seconds_match = regex_module.match(r"^([\d.]+)s$", duration_str)
551
+ if pure_seconds_match:
552
+ return int(float(pure_seconds_match.group(1)))
553
+
554
+ # Handle compound format: "143h4m52.730699158s"
555
+ total_seconds = 0
556
+ patterns = [
557
+ (r"(\d+)h", 3600), # hours
558
+ (r"(\d+)m", 60), # minutes
559
+ (r"([\d.]+)s", 1), # seconds
560
+ ]
561
+ for pattern, multiplier in patterns:
562
+ match = regex_module.search(pattern, duration_str)
563
+ if match:
564
+ total_seconds += float(match.group(1)) * multiplier
565
+
566
+ return int(total_seconds) if total_seconds > 0 else None
567
+
568
+ # Get error body from exception if not provided
569
+ body = error_body
570
+ if not body:
571
+ # Try to extract from various exception attributes
572
+ if hasattr(error, "response") and hasattr(error.response, "text"):
573
+ body = error.response.text
574
+ elif hasattr(error, "body"):
575
+ body = str(error.body)
576
+ elif hasattr(error, "message"):
577
+ body = str(error.message)
578
+ else:
579
+ body = str(error)
580
+
581
+ # Try to find JSON in the body
582
+ try:
583
+ # Handle cases where JSON is embedded in a larger string
584
+ json_match = regex_module.search(r"\{[\s\S]*\}", body)
585
+ if not json_match:
586
+ return None
587
+
588
+ data = json.loads(json_match.group(0))
589
+ except (json.JSONDecodeError, AttributeError, TypeError):
590
+ return None
591
+
592
+ # Navigate to error.details
593
+ error_obj = data.get("error", data)
594
+ details = error_obj.get("details", [])
595
+
596
+ if not details:
597
+ return None
598
+
599
+ result = {
600
+ "retry_after": None,
601
+ "reason": None,
602
+ "reset_timestamp": None,
603
+ }
604
+
605
+ for detail in details:
606
+ detail_type = detail.get("@type", "")
607
+
608
+ # Parse RetryInfo - most authoritative source for retry delay
609
+ if "RetryInfo" in detail_type:
610
+ retry_delay = detail.get("retryDelay")
611
+ if retry_delay:
612
+ parsed = parse_duration(retry_delay)
613
+ if parsed:
614
+ result["retry_after"] = parsed
615
+
616
+ # Parse ErrorInfo - contains reason and quota reset metadata
617
+ elif "ErrorInfo" in detail_type:
618
+ result["reason"] = detail.get("reason")
619
+ metadata = detail.get("metadata", {})
620
+
621
+ # Get quotaResetDelay as fallback if RetryInfo not present
622
+ if not result["retry_after"]:
623
+ quota_delay = metadata.get("quotaResetDelay")
624
+ if quota_delay:
625
+ parsed = parse_duration(quota_delay)
626
+ if parsed:
627
+ result["retry_after"] = parsed
628
+
629
+ # Capture reset timestamp for logging
630
+ result["reset_timestamp"] = metadata.get("quotaResetTimeStamp")
631
+
632
+ # Return None if we couldn't extract retry_after
633
+ if not result["retry_after"]:
634
+ return None
635
+
636
+ return result
637
+
638
  def __init__(self):
639
  super().__init__()
640
  self.model_definitions = ModelDefinitions()
src/rotator_library/providers/gemini_cli_provider.py CHANGED
@@ -186,6 +186,31 @@ def _env_int(key: str, default: int) -> int:
186
  class GeminiCliProvider(GeminiAuthBase, ProviderInterface):
187
  skip_cost_calculation = True
188
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
189
  def __init__(self):
190
  super().__init__()
191
  self.model_definitions = ModelDefinitions()
 
186
  class GeminiCliProvider(GeminiAuthBase, ProviderInterface):
187
  skip_cost_calculation = True
188
 
189
+ # Balanced by default - Gemini CLI has short cooldowns (seconds, not hours)
190
+ default_rotation_mode: str = "balanced"
191
+
192
+ @staticmethod
193
+ def parse_quota_error(
194
+ error: Exception, error_body: Optional[str] = None
195
+ ) -> Optional[Dict[str, Any]]:
196
+ """
197
+ Parse Gemini CLI quota errors.
198
+
199
+ Uses the same Google RPC format as Antigravity but typically has
200
+ much shorter cooldown durations (seconds to minutes, not hours).
201
+
202
+ Args:
203
+ error: The caught exception
204
+ error_body: Optional raw response body string
205
+
206
+ Returns:
207
+ Same format as AntigravityProvider.parse_quota_error()
208
+ """
209
+ # Reuse the same parsing logic as Antigravity since both use Google RPC format
210
+ from .antigravity_provider import AntigravityProvider
211
+
212
+ return AntigravityProvider.parse_quota_error(error, error_body)
213
+
214
  def __init__(self):
215
  super().__init__()
216
  self.model_definitions = ModelDefinitions()
src/rotator_library/providers/provider_interface.py CHANGED
@@ -1,5 +1,6 @@
1
  from abc import ABC, abstractmethod
2
  from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 
3
  import httpx
4
  import litellm
5
 
@@ -12,6 +13,11 @@ class ProviderInterface(ABC):
12
 
13
  skip_cost_calculation: bool = False
14
 
 
 
 
 
 
15
  @abstractmethod
16
  async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]:
17
  """
@@ -153,3 +159,69 @@ class ProviderInterface(ABC):
153
  Tier name string (e.g., "free-tier", "paid-tier") or None if unknown
154
  """
155
  return None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  from abc import ABC, abstractmethod
2
  from typing import List, Dict, Any, Optional, AsyncGenerator, Union
3
+ import os
4
  import httpx
5
  import litellm
6
 
 
13
 
14
  skip_cost_calculation: bool = False
15
 
16
+ # Default rotation mode for this provider ("balanced" or "sequential")
17
+ # - "balanced": Rotate credentials to distribute load evenly
18
+ # - "sequential": Use one credential until exhausted, then switch to next
19
+ default_rotation_mode: str = "balanced"
20
+
21
  @abstractmethod
22
  async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]:
23
  """
 
159
  Tier name string (e.g., "free-tier", "paid-tier") or None if unknown
160
  """
161
  return None
162
+
163
+ # =========================================================================
164
+ # Sequential Rotation Support
165
+ # =========================================================================
166
+
167
+ @classmethod
168
+ def get_rotation_mode(cls, provider_name: str) -> str:
169
+ """
170
+ Get the rotation mode for this provider.
171
+
172
+ Checks ROTATION_MODE_{PROVIDER} environment variable first,
173
+ then falls back to the class's default_rotation_mode.
174
+
175
+ Args:
176
+ provider_name: The provider name (e.g., "antigravity", "gemini_cli")
177
+
178
+ Returns:
179
+ "balanced" or "sequential"
180
+ """
181
+ env_key = f"ROTATION_MODE_{provider_name.upper()}"
182
+ return os.getenv(env_key, cls.default_rotation_mode)
183
+
184
+ @staticmethod
185
+ def parse_quota_error(
186
+ error: Exception, error_body: Optional[str] = None
187
+ ) -> Optional[Dict[str, Any]]:
188
+ """
189
+ Parse a quota/rate-limit error and extract structured information.
190
+
191
+ Providers should override this method to handle their specific error formats.
192
+ This allows the error_handler to use provider-specific parsing when available,
193
+ falling back to generic parsing otherwise.
194
+
195
+ Args:
196
+ error: The caught exception
197
+ error_body: Optional raw response body string
198
+
199
+ Returns:
200
+ None if not a parseable quota error, otherwise:
201
+ {
202
+ "retry_after": int, # seconds until quota resets
203
+ "reason": str, # e.g., "QUOTA_EXHAUSTED", "RATE_LIMITED"
204
+ "reset_timestamp": str | None, # ISO timestamp if available
205
+ }
206
+ """
207
+ return None # Default: no provider-specific parsing
208
+
209
+ # TODO: Implement provider-specific quota reset schedules
210
+ # Different providers have different quota reset periods:
211
+ # - Most providers: Daily reset at a specific time
212
+ # - Antigravity free tier: Weekly reset
213
+ # - Antigravity paid tier: 5-hour rolling window
214
+ #
215
+ # Future implementation should add:
216
+ # @classmethod
217
+ # def get_quota_reset_behavior(cls) -> Dict[str, Any]:
218
+ # """
219
+ # Get provider-specific quota reset behavior.
220
+ # Returns:
221
+ # {
222
+ # "type": "daily" | "weekly" | "rolling",
223
+ # "reset_time_utc": "03:00", # For daily/weekly
224
+ # "rolling_hours": 5, # For rolling
225
+ # }
226
+ # """
227
+ # return {"type": "daily", "reset_time_utc": "03:00"}
src/rotator_library/usage_manager.py CHANGED
@@ -5,7 +5,7 @@ import logging
5
  import asyncio
6
  import random
7
  from datetime import date, datetime, timezone, time as dt_time
8
- from typing import Any, Dict, List, Optional, Set
9
  import aiofiles
10
  import litellm
11
 
@@ -42,6 +42,10 @@ class UsageManager:
42
 
43
  This ensures lower-usage credentials are preferred while tolerance controls how much
44
  randomness is introduced into the selection process.
 
 
 
 
45
  """
46
 
47
  def __init__(
@@ -49,6 +53,7 @@ class UsageManager:
49
  file_path: str = "key_usage.json",
50
  daily_reset_time_utc: Optional[str] = "03:00",
51
  rotation_tolerance: float = 0.0,
 
52
  ):
53
  """
54
  Initialize the UsageManager.
@@ -60,9 +65,13 @@ class UsageManager:
60
  - 0.0: Deterministic, least-used credential always selected
61
  - tolerance = 2.0 - 4.0 (default, recommended): Balanced randomness, can pick credentials within 2 uses of max
62
  - 5.0+: High randomness, more unpredictable selection patterns
 
 
 
63
  """
64
  self.file_path = file_path
65
  self.rotation_tolerance = rotation_tolerance
 
66
  self.key_states: Dict[str, Dict[str, Any]] = {}
67
 
68
  self._data_lock = asyncio.Lock()
@@ -81,6 +90,72 @@ class UsageManager:
81
  else:
82
  self.daily_reset_time_utc = None
83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  async def _lazy_init(self):
85
  """Initializes the usage data by loading it from the file asynchronously."""
86
  async with self._init_lock:
@@ -144,14 +219,63 @@ class UsageManager:
144
  )
145
  needs_saving = True
146
 
147
- # Reset cooldowns
148
- data["model_cooldowns"] = {}
149
- data["key_cooldown_until"] = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
150
 
151
  # Reset consecutive failures
152
  if "failures" in data:
153
  data["failures"] = {}
154
 
 
 
 
 
 
 
 
 
 
 
 
 
 
155
  # Archive global stats from the previous day's 'daily'
156
  daily_data = data.get("daily", {})
157
  if daily_data:
@@ -336,15 +460,30 @@ class UsageManager:
336
  elif key_state["models_in_use"].get(model, 0) < max_concurrent:
337
  tier2_keys.append((key, usage_count))
338
 
339
- # Apply weighted random selection or deterministic sorting
340
- selection_method = (
341
- "weighted-random"
342
- if self.rotation_tolerance > 0
343
- else "least-used"
344
- )
345
 
346
- if self.rotation_tolerance > 0:
347
- # Weighted random selection within each tier
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
348
  if tier1_keys:
349
  selected_key = self._select_weighted_random(
350
  tier1_keys, self.rotation_tolerance
@@ -361,6 +500,7 @@ class UsageManager:
361
  ]
362
  else:
363
  # Deterministic: sort by usage within each tier
 
364
  tier1_keys.sort(key=lambda x: x[1])
365
  tier2_keys.sort(key=lambda x: x[1])
366
 
@@ -452,13 +592,30 @@ class UsageManager:
452
  elif key_state["models_in_use"].get(model, 0) < max_concurrent:
453
  tier2_keys.append((key, usage_count))
454
 
455
- # Apply weighted random selection or deterministic sorting
456
- selection_method = (
457
- "weighted-random" if self.rotation_tolerance > 0 else "least-used"
458
- )
459
 
460
- if self.rotation_tolerance > 0:
461
- # Weighted random selection within each tier
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
462
  if tier1_keys:
463
  selected_key = self._select_weighted_random(
464
  tier1_keys, self.rotation_tolerance
@@ -475,6 +632,7 @@ class UsageManager:
475
  ]
476
  else:
477
  # Deterministic: sort by usage within each tier
 
478
  tier1_keys.sort(key=lambda x: x[1])
479
  tier2_keys.sort(key=lambda x: x[1])
480
 
@@ -726,10 +884,24 @@ class UsageManager:
726
  if classified_error.error_type in ["rate_limit", "quota_exceeded"]:
727
  # Rate limit / Quota errors: use retry_after if available, otherwise default to 60s
728
  cooldown_seconds = classified_error.retry_after or 60
729
- lib_logger.info(
730
- f"Rate limit error on key {mask_credential(key)} for model {model}. "
731
- f"Using {'provided' if classified_error.retry_after else 'default'} retry_after: {cooldown_seconds}s"
732
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
733
  elif classified_error.error_type == "authentication":
734
  # Apply a 5-minute key-level lockout for auth errors
735
  key_data["key_cooldown_until"] = time.time() + 300
 
5
  import asyncio
6
  import random
7
  from datetime import date, datetime, timezone, time as dt_time
8
+ from typing import Any, Dict, List, Optional, Set, Tuple
9
  import aiofiles
10
  import litellm
11
 
 
42
 
43
  This ensures lower-usage credentials are preferred while tolerance controls how much
44
  randomness is introduced into the selection process.
45
+
46
+ Additionally, providers can specify a rotation mode:
47
+ - "balanced" (default): Rotate credentials to distribute load evenly
48
+ - "sequential": Use one credential until exhausted (preserves caching)
49
  """
50
 
51
  def __init__(
 
53
  file_path: str = "key_usage.json",
54
  daily_reset_time_utc: Optional[str] = "03:00",
55
  rotation_tolerance: float = 0.0,
56
+ provider_rotation_modes: Optional[Dict[str, str]] = None,
57
  ):
58
  """
59
  Initialize the UsageManager.
 
65
  - 0.0: Deterministic, least-used credential always selected
66
  - tolerance = 2.0 - 4.0 (default, recommended): Balanced randomness, can pick credentials within 2 uses of max
67
  - 5.0+: High randomness, more unpredictable selection patterns
68
+ provider_rotation_modes: Dict mapping provider names to rotation modes.
69
+ - "balanced": Rotate credentials to distribute load evenly (default)
70
+ - "sequential": Use one credential until exhausted (preserves caching)
71
  """
72
  self.file_path = file_path
73
  self.rotation_tolerance = rotation_tolerance
74
+ self.provider_rotation_modes = provider_rotation_modes or {}
75
  self.key_states: Dict[str, Dict[str, Any]] = {}
76
 
77
  self._data_lock = asyncio.Lock()
 
90
  else:
91
  self.daily_reset_time_utc = None
92
 
93
+ def _get_rotation_mode(self, provider: str) -> str:
94
+ """
95
+ Get the rotation mode for a provider.
96
+
97
+ Args:
98
+ provider: Provider name (e.g., "antigravity", "gemini_cli")
99
+
100
+ Returns:
101
+ "balanced" or "sequential"
102
+ """
103
+ return self.provider_rotation_modes.get(provider, "balanced")
104
+
105
+ def _select_sequential(
106
+ self,
107
+ candidates: List[Tuple[str, int]],
108
+ credential_priorities: Optional[Dict[str, int]] = None,
109
+ ) -> str:
110
+ """
111
+ Select credential in strict sequential order for cache-preserving rotation.
112
+
113
+ This method ensures the same credential is reused until it hits a cooldown,
114
+ which preserves provider-side caching (e.g., thinking signature caches).
115
+
116
+ Selection logic:
117
+ 1. Sort by priority (lowest number = highest priority)
118
+ 2. Within same priority, sort by last_used_ts (most recent first = sticky)
119
+ 3. Return the first candidate
120
+
121
+ Args:
122
+ candidates: List of (credential_id, usage_count) tuples
123
+ credential_priorities: Optional dict mapping credentials to priority levels
124
+
125
+ Returns:
126
+ Selected credential ID
127
+ """
128
+ if not candidates:
129
+ raise ValueError("Cannot select from empty candidate list")
130
+
131
+ if len(candidates) == 1:
132
+ return candidates[0][0]
133
+
134
+ def sort_key(item: Tuple[str, int]) -> Tuple[int, float]:
135
+ cred, _ = item
136
+ # Priority: lower is better (1 = highest priority)
137
+ priority = (
138
+ credential_priorities.get(cred, 999) if credential_priorities else 999
139
+ )
140
+ # Last used: higher (more recent) is better for stickiness
141
+ last_used = (
142
+ self._usage_data.get(cred, {}).get("last_used_ts", 0)
143
+ if self._usage_data
144
+ else 0
145
+ )
146
+ # Negative last_used so most recent sorts first
147
+ return (priority, -last_used)
148
+
149
+ sorted_candidates = sorted(candidates, key=sort_key)
150
+ selected = sorted_candidates[0][0]
151
+
152
+ lib_logger.debug(
153
+ f"Sequential selection: chose {mask_credential(selected)} "
154
+ f"(priority={credential_priorities.get(selected, 999) if credential_priorities else 'N/A'})"
155
+ )
156
+
157
+ return selected
158
+
159
  async def _lazy_init(self):
160
  """Initializes the usage data by loading it from the file asynchronously."""
161
  async with self._init_lock:
 
219
  )
220
  needs_saving = True
221
 
222
+ # Reset cooldowns - BUT preserve unexpired long-term cooldowns
223
+ # This is important for quota errors with long cooldowns (e.g., 143 hours)
224
+ now_ts = time.time()
225
+ if "model_cooldowns" in data:
226
+ active_cooldowns = {
227
+ model: end_time
228
+ for model, end_time in data["model_cooldowns"].items()
229
+ if end_time > now_ts
230
+ }
231
+ if active_cooldowns:
232
+ # Calculate how long the longest cooldown has remaining
233
+ max_remaining = max(
234
+ end_time - now_ts
235
+ for end_time in active_cooldowns.values()
236
+ )
237
+ hours_remaining = max_remaining / 3600
238
+ lib_logger.info(
239
+ f"Preserving {len(active_cooldowns)} active cooldown(s) "
240
+ f"for key {mask_credential(key)} during daily reset "
241
+ f"(longest: {hours_remaining:.1f}h remaining)"
242
+ )
243
+ data["model_cooldowns"] = active_cooldowns
244
+ else:
245
+ data["model_cooldowns"] = {}
246
+
247
+ # Clear key-level cooldown only if expired
248
+ if data.get("key_cooldown_until"):
249
+ if data["key_cooldown_until"] <= now_ts:
250
+ data["key_cooldown_until"] = None
251
+ else:
252
+ hours_remaining = (
253
+ data["key_cooldown_until"] - now_ts
254
+ ) / 3600
255
+ lib_logger.info(
256
+ f"Preserving key-level cooldown for {mask_credential(key)} "
257
+ f"during daily reset ({hours_remaining:.1f}h remaining)"
258
+ )
259
+ else:
260
+ data["key_cooldown_until"] = None
261
 
262
  # Reset consecutive failures
263
  if "failures" in data:
264
  data["failures"] = {}
265
 
266
+ # TODO: Implement provider-specific reset schedules
267
+ # Different providers have different quota reset periods:
268
+ # - Most providers: Daily reset at daily_reset_time_utc
269
+ # - Antigravity free tier: Weekly reset
270
+ # - Antigravity paid tier: 5-hour rolling window
271
+ #
272
+ # Future implementation should:
273
+ # 1. Group credentials by provider (extracted from key path or metadata)
274
+ # 2. Check each provider's get_quota_reset_behavior()
275
+ # 3. Apply provider-specific reset logic instead of universal daily reset
276
+ #
277
+ # For now, we preserve unexpired cooldowns which handles long cooldowns correctly.
278
+
279
  # Archive global stats from the previous day's 'daily'
280
  daily_data = data.get("daily", {})
281
  if daily_data:
 
460
  elif key_state["models_in_use"].get(model, 0) < max_concurrent:
461
  tier2_keys.append((key, usage_count))
462
 
463
+ # Determine selection method based on provider's rotation mode
464
+ provider = model.split("/")[0] if "/" in model else ""
465
+ rotation_mode = self._get_rotation_mode(provider)
 
 
 
466
 
467
+ if rotation_mode == "sequential":
468
+ # Sequential mode: stick with same credential until exhausted
469
+ selection_method = "sequential"
470
+ if tier1_keys:
471
+ selected_key = self._select_sequential(
472
+ tier1_keys, credential_priorities
473
+ )
474
+ tier1_keys = [
475
+ (k, u) for k, u in tier1_keys if k == selected_key
476
+ ]
477
+ if tier2_keys:
478
+ selected_key = self._select_sequential(
479
+ tier2_keys, credential_priorities
480
+ )
481
+ tier2_keys = [
482
+ (k, u) for k, u in tier2_keys if k == selected_key
483
+ ]
484
+ elif self.rotation_tolerance > 0:
485
+ # Balanced mode with weighted randomness
486
+ selection_method = "weighted-random"
487
  if tier1_keys:
488
  selected_key = self._select_weighted_random(
489
  tier1_keys, self.rotation_tolerance
 
500
  ]
501
  else:
502
  # Deterministic: sort by usage within each tier
503
+ selection_method = "least-used"
504
  tier1_keys.sort(key=lambda x: x[1])
505
  tier2_keys.sort(key=lambda x: x[1])
506
 
 
592
  elif key_state["models_in_use"].get(model, 0) < max_concurrent:
593
  tier2_keys.append((key, usage_count))
594
 
595
+ # Determine selection method based on provider's rotation mode
596
+ provider = model.split("/")[0] if "/" in model else ""
597
+ rotation_mode = self._get_rotation_mode(provider)
 
598
 
599
+ if rotation_mode == "sequential":
600
+ # Sequential mode: stick with same credential until exhausted
601
+ selection_method = "sequential"
602
+ if tier1_keys:
603
+ selected_key = self._select_sequential(
604
+ tier1_keys, credential_priorities
605
+ )
606
+ tier1_keys = [
607
+ (k, u) for k, u in tier1_keys if k == selected_key
608
+ ]
609
+ if tier2_keys:
610
+ selected_key = self._select_sequential(
611
+ tier2_keys, credential_priorities
612
+ )
613
+ tier2_keys = [
614
+ (k, u) for k, u in tier2_keys if k == selected_key
615
+ ]
616
+ elif self.rotation_tolerance > 0:
617
+ # Balanced mode with weighted randomness
618
+ selection_method = "weighted-random"
619
  if tier1_keys:
620
  selected_key = self._select_weighted_random(
621
  tier1_keys, self.rotation_tolerance
 
632
  ]
633
  else:
634
  # Deterministic: sort by usage within each tier
635
+ selection_method = "least-used"
636
  tier1_keys.sort(key=lambda x: x[1])
637
  tier2_keys.sort(key=lambda x: x[1])
638
 
 
884
  if classified_error.error_type in ["rate_limit", "quota_exceeded"]:
885
  # Rate limit / Quota errors: use retry_after if available, otherwise default to 60s
886
  cooldown_seconds = classified_error.retry_after or 60
887
+ if classified_error.retry_after:
888
+ # Log with human-readable duration for provider-parsed cooldowns
889
+ hours = cooldown_seconds / 3600
890
+ if hours >= 1:
891
+ lib_logger.info(
892
+ f"Quota/rate limit on key {mask_credential(key)} for model {model}. "
893
+ f"Applying provider-specified cooldown: {cooldown_seconds}s ({hours:.1f}h)"
894
+ )
895
+ else:
896
+ lib_logger.info(
897
+ f"Rate limit on key {mask_credential(key)} for model {model}. "
898
+ f"Applying provider-specified cooldown: {cooldown_seconds}s"
899
+ )
900
+ else:
901
+ lib_logger.info(
902
+ f"Rate limit on key {mask_credential(key)} for model {model}. "
903
+ f"Using default cooldown: {cooldown_seconds}s"
904
+ )
905
  elif classified_error.error_type == "authentication":
906
  # Apply a 5-minute key-level lockout for auth errors
907
  key_data["key_cooldown_until"] = time.time() + 300