Mirrowel commited on
Commit
2b065a1
·
1 Parent(s): f7eb788

feat(auth): improve credential tool startup UX and lazy-load providers

Browse files

- Defer heavy provider imports via a new _ensure_providers_loaded() helper to avoid the 6-7s blank startup wait and load providers lazily.
- Enhance launcher TUI to show a full loading UI with spinner, header, elapsed time and provider count; call run_credential_tool(from_launcher=True) to avoid duplicate loading screens.
- Add from_launcher flag to run_credential_tool and clear_on_start parameter to main() so standalone vs launcher flows behave correctly.
- Swap direct provider_factory and PROVIDER_PLUGINS usage to go through the lazy loader in credential tool and callers.
- Improve interactive UX: preserve/clear screen appropriately, add short pauses and prompts after flows, and keep console status messages during initialization.
- Fix output formatting in qwen/iflow providers (replace escaped newlines with actual newlines) and enhance headless detection to log and print user-visible guidance when running without a GUI.

src/proxy_app/launcher_tui.py CHANGED
@@ -512,8 +512,34 @@ class LauncherTUI:
512
 
513
  def launch_credential_tool(self):
514
  """Launch credential management tool"""
515
- from rotator_library.credential_tool import run_credential_tool
516
- run_credential_tool()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
517
  # Reload environment after credential tool
518
  load_dotenv(dotenv_path=Path.cwd() / ".env", override=True)
519
 
 
512
 
513
  def launch_credential_tool(self):
514
  """Launch credential management tool"""
515
+ import time
516
+
517
+ # CRITICAL: Show full loading UI to replace the 6-7 second blank wait
518
+ self.console.clear()
519
+
520
+ _start_time = time.time()
521
+
522
+ # Show the same header as standalone mode
523
+ self.console.print("━" * 70)
524
+ self.console.print("Interactive Credential Setup Tool")
525
+ self.console.print("GitHub: https://github.com/Mirrowel/LLM-API-Key-Proxy")
526
+ self.console.print("━" * 70)
527
+ self.console.print("Loading credential management components...")
528
+
529
+ # Now import with spinner (this is where the 6-7 second delay happens)
530
+ with self.console.status("Initializing credential tool...", spinner="dots"):
531
+ from rotator_library.credential_tool import run_credential_tool, _ensure_providers_loaded
532
+ _, PROVIDER_PLUGINS = _ensure_providers_loaded()
533
+ self.console.print("✓ Credential tool initialized")
534
+
535
+ _elapsed = time.time() - _start_time
536
+ self.console.print(f"✓ Tool ready in {_elapsed:.2f}s ({len(PROVIDER_PLUGINS)} providers available)")
537
+
538
+ # Small delay to let user see the ready message
539
+ time.sleep(0.5)
540
+
541
+ # Run the tool with from_launcher=True to skip duplicate loading screen
542
+ run_credential_tool(from_launcher=True)
543
  # Reload environment after credential tool
544
  load_dotenv(dotenv_path=Path.cwd() / ".env", override=True)
545
 
src/rotator_library/credential_tool.py CHANGED
@@ -2,13 +2,14 @@
2
 
3
  import asyncio
4
  import json
 
5
  import re
6
  import time
7
  from pathlib import Path
8
  from dotenv import set_key, get_key
9
 
10
- from .provider_factory import get_provider_auth_class, get_available_providers
11
- from .providers import PROVIDER_PLUGINS
12
  from rich.console import Console
13
  from rich.panel import Panel
14
  from rich.prompt import Prompt
@@ -19,9 +20,22 @@ OAUTH_BASE_DIR.mkdir(exist_ok=True)
19
  # Use a direct path to the .env file in the project root
20
  ENV_FILE = Path.cwd() / ".env"
21
 
22
-
23
  console = Console()
24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
  def ensure_env_defaults():
26
  """
27
  Ensures the .env file exists and contains essential default values like PROXY_API_KEY.
@@ -83,6 +97,7 @@ async def setup_api_key():
83
 
84
  # Discover custom providers and add them to the list
85
  # Note: gemini_cli is OAuth-only, but qwen_code and iflow support both OAuth and API keys
 
86
  oauth_only_providers = {'gemini_cli'}
87
  discovered_providers = {
88
  p.replace('_', ' ').title(): p.upper() + "_API_KEY"
@@ -172,7 +187,8 @@ async def setup_new_credential(provider_name: str):
172
  Interactively sets up a new OAuth credential for a given provider.
173
  """
174
  try:
175
- auth_class = get_provider_auth_class(provider_name)
 
176
  auth_instance = auth_class()
177
 
178
  # Build display name for better user experience
@@ -516,15 +532,25 @@ async def export_iflow_to_env():
516
  console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error"))
517
 
518
 
519
- async def main():
520
  """
521
  An interactive CLI tool to add new credentials.
 
 
 
 
522
  """
523
- console.clear() # Clear terminal when credential tool starts
524
  ensure_env_defaults()
525
- console.print(Panel("[bold cyan]Interactive Credential Setup[/bold cyan]", title="--- API Key Proxy ---", expand=False))
 
 
 
526
 
527
  while True:
 
 
 
 
528
  console.print(Panel(
529
  Text.from_markup(
530
  "1. Add OAuth Credential\n"
@@ -547,7 +573,8 @@ async def main():
547
  break
548
 
549
  if setup_type == "1":
550
- available_providers = get_available_providers()
 
551
  oauth_friendly_names = {
552
  "gemini_cli": "Gemini CLI (OAuth)",
553
  "qwen_code": "Qwen Code (OAuth - also supports API keys)",
@@ -577,28 +604,77 @@ async def main():
577
  display_name = oauth_friendly_names.get(provider_name, provider_name.replace('_', ' ').title())
578
  console.print(f"\nStarting OAuth setup for [bold cyan]{display_name}[/bold cyan]...")
579
  await setup_new_credential(provider_name)
 
 
 
580
  else:
581
  console.print("[bold red]Invalid choice. Please try again.[/bold red]")
 
582
  except ValueError:
583
  console.print("[bold red]Invalid input. Please enter a number or 'b'.[/bold red]")
 
584
 
585
  elif setup_type == "2":
586
  await setup_api_key()
 
 
587
 
588
  elif setup_type == "3":
589
  await export_gemini_cli_to_env()
 
 
590
 
591
  elif setup_type == "4":
592
  await export_qwen_code_to_env()
 
 
593
 
594
  elif setup_type == "5":
595
  await export_iflow_to_env()
 
 
596
 
597
- console.print("\n" + "="*50 + "\n")
598
-
599
- def run_credential_tool():
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
600
  try:
601
- asyncio.run(main())
602
  console.clear() # Clear terminal when credential tool exits
603
  except KeyboardInterrupt:
604
  console.print("\n[bold yellow]Exiting setup.[/bold yellow]")
 
2
 
3
  import asyncio
4
  import json
5
+ import os
6
  import re
7
  import time
8
  from pathlib import Path
9
  from dotenv import set_key, get_key
10
 
11
+ # NOTE: Heavy imports (provider_factory, PROVIDER_PLUGINS) are deferred
12
+ # to avoid 6-7 second delay before showing loading screen
13
  from rich.console import Console
14
  from rich.panel import Panel
15
  from rich.prompt import Prompt
 
20
  # Use a direct path to the .env file in the project root
21
  ENV_FILE = Path.cwd() / ".env"
22
 
 
23
  console = Console()
24
 
25
+ # Global variables for lazily loaded modules
26
+ _provider_factory = None
27
+ _provider_plugins = None
28
+
29
+ def _ensure_providers_loaded():
30
+ """Lazy load provider modules only when needed"""
31
+ global _provider_factory, _provider_plugins
32
+ if _provider_factory is None:
33
+ from . import provider_factory as pf
34
+ from .providers import PROVIDER_PLUGINS as pp
35
+ _provider_factory = pf
36
+ _provider_plugins = pp
37
+ return _provider_factory, _provider_plugins
38
+
39
  def ensure_env_defaults():
40
  """
41
  Ensures the .env file exists and contains essential default values like PROXY_API_KEY.
 
97
 
98
  # Discover custom providers and add them to the list
99
  # Note: gemini_cli is OAuth-only, but qwen_code and iflow support both OAuth and API keys
100
+ _, PROVIDER_PLUGINS = _ensure_providers_loaded()
101
  oauth_only_providers = {'gemini_cli'}
102
  discovered_providers = {
103
  p.replace('_', ' ').title(): p.upper() + "_API_KEY"
 
187
  Interactively sets up a new OAuth credential for a given provider.
188
  """
189
  try:
190
+ provider_factory, _ = _ensure_providers_loaded()
191
+ auth_class = provider_factory.get_provider_auth_class(provider_name)
192
  auth_instance = auth_class()
193
 
194
  # Build display name for better user experience
 
532
  console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error"))
533
 
534
 
535
+ async def main(clear_on_start=True):
536
  """
537
  An interactive CLI tool to add new credentials.
538
+
539
+ Args:
540
+ clear_on_start: If False, skip initial screen clear (used when called from launcher
541
+ to preserve the loading screen)
542
  """
 
543
  ensure_env_defaults()
544
+
545
+ # Only show header if we're clearing (standalone mode)
546
+ if clear_on_start:
547
+ console.print(Panel("[bold cyan]Interactive Credential Setup[/bold cyan]", title="--- API Key Proxy ---", expand=False))
548
 
549
  while True:
550
+ # Clear screen between menu selections for cleaner UX
551
+ console.clear()
552
+ console.print(Panel("[bold cyan]Interactive Credential Setup[/bold cyan]", title="--- API Key Proxy ---", expand=False))
553
+
554
  console.print(Panel(
555
  Text.from_markup(
556
  "1. Add OAuth Credential\n"
 
573
  break
574
 
575
  if setup_type == "1":
576
+ provider_factory, _ = _ensure_providers_loaded()
577
+ available_providers = provider_factory.get_available_providers()
578
  oauth_friendly_names = {
579
  "gemini_cli": "Gemini CLI (OAuth)",
580
  "qwen_code": "Qwen Code (OAuth - also supports API keys)",
 
604
  display_name = oauth_friendly_names.get(provider_name, provider_name.replace('_', ' ').title())
605
  console.print(f"\nStarting OAuth setup for [bold cyan]{display_name}[/bold cyan]...")
606
  await setup_new_credential(provider_name)
607
+ # Don't clear after OAuth - user needs to see full flow
608
+ console.print("\n[dim]Press Enter to return to main menu...[/dim]")
609
+ input()
610
  else:
611
  console.print("[bold red]Invalid choice. Please try again.[/bold red]")
612
+ await asyncio.sleep(1.5)
613
  except ValueError:
614
  console.print("[bold red]Invalid input. Please enter a number or 'b'.[/bold red]")
615
+ await asyncio.sleep(1.5)
616
 
617
  elif setup_type == "2":
618
  await setup_api_key()
619
+ #console.print("\n[dim]Press Enter to return to main menu...[/dim]")
620
+ input()
621
 
622
  elif setup_type == "3":
623
  await export_gemini_cli_to_env()
624
+ console.print("\n[dim]Press Enter to return to main menu...[/dim]")
625
+ input()
626
 
627
  elif setup_type == "4":
628
  await export_qwen_code_to_env()
629
+ console.print("\n[dim]Press Enter to return to main menu...[/dim]")
630
+ input()
631
 
632
  elif setup_type == "5":
633
  await export_iflow_to_env()
634
+ console.print("\n[dim]Press Enter to return to main menu...[/dim]")
635
+ input()
636
 
637
+ def run_credential_tool(from_launcher=False):
638
+ """
639
+ Entry point for credential tool.
640
+
641
+ Args:
642
+ from_launcher: If True, skip loading screen (launcher already showed it)
643
+ """
644
+ # Check if we need to show loading screen
645
+ if not from_launcher:
646
+ # Standalone mode - show full loading UI
647
+ os.system('cls' if os.name == 'nt' else 'clear')
648
+
649
+ _start_time = time.time()
650
+
651
+ # Phase 1: Show initial message
652
+ print("━" * 70)
653
+ print("Interactive Credential Setup Tool")
654
+ print("GitHub: https://github.com/Mirrowel/LLM-API-Key-Proxy")
655
+ print("━" * 70)
656
+ print("Loading credential management components...")
657
+
658
+ # Phase 2: Load dependencies with spinner
659
+ with console.status("Loading authentication providers...", spinner="dots"):
660
+ _ensure_providers_loaded()
661
+ console.print("✓ Authentication providers loaded")
662
+
663
+ with console.status("Initializing credential tool...", spinner="dots"):
664
+ time.sleep(0.2) # Brief pause for UI consistency
665
+ console.print("✓ Credential tool initialized")
666
+
667
+ _elapsed = time.time() - _start_time
668
+ _, PROVIDER_PLUGINS = _ensure_providers_loaded()
669
+ print(f"✓ Tool ready in {_elapsed:.2f}s ({len(PROVIDER_PLUGINS)} providers available)")
670
+
671
+ # Small delay to let user see the ready message
672
+ time.sleep(0.5)
673
+
674
+ # Run the main async event loop
675
+ # If from launcher, don't clear screen at start to preserve loading messages
676
  try:
677
+ asyncio.run(main(clear_on_start=not from_launcher))
678
  console.clear() # Clear terminal when credential tool exits
679
  except KeyboardInterrupt:
680
  console.print("\n[bold yellow]Exiting setup.[/bold yellow]")
src/rotator_library/providers/iflow_auth_base.py CHANGED
@@ -682,21 +682,21 @@ class IFlowAuthBase:
682
  # [HEADLESS SUPPORT] Display appropriate instructions
683
  if is_headless:
684
  auth_panel_text = Text.from_markup(
685
- "Running in headless environment (no GUI detected).\\n"
686
- "Please open the URL below in a browser on another machine to authorize:\\n"
687
- "1. Visit the URL below to sign in with your phone number.\\n"
688
- "2. [bold]Authorize the application[/bold] to access your account.\\n"
689
  "3. You will be automatically redirected after authorization."
690
  )
691
  else:
692
  auth_panel_text = Text.from_markup(
693
- "1. Visit the URL below to sign in with your phone number.\\n"
694
- "2. [bold]Authorize the application[/bold] to access your account.\\n"
695
  "3. You will be automatically redirected after authorization."
696
  )
697
 
698
  console.print(Panel(auth_panel_text, title=f"iFlow OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue"))
699
- console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\\n")
700
 
701
  # [HEADLESS SUPPORT] Only attempt browser open if NOT headless
702
  if not is_headless:
 
682
  # [HEADLESS SUPPORT] Display appropriate instructions
683
  if is_headless:
684
  auth_panel_text = Text.from_markup(
685
+ "Running in headless environment (no GUI detected).\n"
686
+ "Please open the URL below in a browser on another machine to authorize:\n"
687
+ "1. Visit the URL below to sign in with your phone number.\n"
688
+ "2. [bold]Authorize the application[/bold] to access your account.\n"
689
  "3. You will be automatically redirected after authorization."
690
  )
691
  else:
692
  auth_panel_text = Text.from_markup(
693
+ "1. Visit the URL below to sign in with your phone number.\n"
694
+ "2. [bold]Authorize the application[/bold] to access your account.\n"
695
  "3. You will be automatically redirected after authorization."
696
  )
697
 
698
  console.print(Panel(auth_panel_text, title=f"iFlow OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue"))
699
+ console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n")
700
 
701
  # [HEADLESS SUPPORT] Only attempt browser open if NOT headless
702
  if not is_headless:
src/rotator_library/providers/qwen_auth_base.py CHANGED
@@ -431,21 +431,21 @@ class QwenAuthBase:
431
  # [HEADLESS SUPPORT] Display appropriate instructions
432
  if is_headless:
433
  auth_panel_text = Text.from_markup(
434
- "Running in headless environment (no GUI detected).\\n"
435
- "Please open the URL below in a browser on another machine to authorize:\\n"
436
- "1. Visit the URL below to sign in.\\n"
437
- "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\\n"
438
  "3. You will be prompted to enter your identifier after authorization."
439
  )
440
  else:
441
  auth_panel_text = Text.from_markup(
442
- "1. Visit the URL below to sign in.\\n"
443
- "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\\n"
444
  "3. You will be prompted to enter your identifier after authorization."
445
  )
446
 
447
  console.print(Panel(auth_panel_text, title=f"Qwen OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue"))
448
- console.print(f"[bold]URL:[/bold] [link={dev_data['verification_uri_complete']}]{dev_data['verification_uri_complete']}[/link]\\n")
449
 
450
  # [HEADLESS SUPPORT] Only attempt browser open if NOT headless
451
  if not is_headless:
 
431
  # [HEADLESS SUPPORT] Display appropriate instructions
432
  if is_headless:
433
  auth_panel_text = Text.from_markup(
434
+ "Running in headless environment (no GUI detected).\n"
435
+ "Please open the URL below in a browser on another machine to authorize:\n"
436
+ "1. Visit the URL below to sign in.\n"
437
+ "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n"
438
  "3. You will be prompted to enter your identifier after authorization."
439
  )
440
  else:
441
  auth_panel_text = Text.from_markup(
442
+ "1. Visit the URL below to sign in.\n"
443
+ "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n"
444
  "3. You will be prompted to enter your identifier after authorization."
445
  )
446
 
447
  console.print(Panel(auth_panel_text, title=f"Qwen OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue"))
448
+ console.print(f"[bold]URL:[/bold] [link={dev_data['verification_uri_complete']}]{dev_data['verification_uri_complete']}[/link]\n")
449
 
450
  # [HEADLESS SUPPORT] Only attempt browser open if NOT headless
451
  if not is_headless:
src/rotator_library/utils/headless_detection.py CHANGED
@@ -5,6 +5,13 @@ import logging
5
 
6
  lib_logger = logging.getLogger('rotator_library')
7
 
 
 
 
 
 
 
 
8
  def is_headless_environment() -> bool:
9
  """
10
  Detects if the current environment is headless (no GUI available).
@@ -20,10 +27,11 @@ def is_headless_environment() -> bool:
20
  """
21
  headless_indicators = []
22
 
23
- # Check DISPLAY for Linux/Unix GUI availability
24
- display = os.getenv("DISPLAY")
25
- if display is None or display.strip() == "":
26
- headless_indicators.append("No DISPLAY variable (Linux/Unix headless)")
 
27
 
28
  # Check for SSH connection
29
  if os.getenv("SSH_CONNECTION") or os.getenv("SSH_CLIENT") or os.getenv("SSH_TTY"):
@@ -62,8 +70,15 @@ def is_headless_environment() -> bool:
62
  is_headless = len(headless_indicators) > 0
63
 
64
  if is_headless:
 
65
  lib_logger.info(f"Headless environment detected: {'; '.join(headless_indicators)}")
 
 
 
 
 
66
  else:
 
67
  lib_logger.debug("GUI environment detected, browser auto-open will be attempted")
68
 
69
  return is_headless
 
5
 
6
  lib_logger = logging.getLogger('rotator_library')
7
 
8
+ # Import console for user-visible output
9
+ try:
10
+ from rich.console import Console
11
+ console = Console()
12
+ except ImportError:
13
+ console = None
14
+
15
  def is_headless_environment() -> bool:
16
  """
17
  Detects if the current environment is headless (no GUI available).
 
27
  """
28
  headless_indicators = []
29
 
30
+ # Check DISPLAY for Linux/Unix GUI availability (skip on Windows)
31
+ if os.name != 'nt': # Only check DISPLAY on non-Windows systems
32
+ display = os.getenv("DISPLAY")
33
+ if display is None or display.strip() == "":
34
+ headless_indicators.append("No DISPLAY variable (Linux/Unix headless)")
35
 
36
  # Check for SSH connection
37
  if os.getenv("SSH_CONNECTION") or os.getenv("SSH_CLIENT") or os.getenv("SSH_TTY"):
 
70
  is_headless = len(headless_indicators) > 0
71
 
72
  if is_headless:
73
+ # Log to logger
74
  lib_logger.info(f"Headless environment detected: {'; '.join(headless_indicators)}")
75
+
76
+ # Print to console for user visibility
77
+ if console:
78
+ console.print(f"[yellow]ℹ Headless environment detected:[/yellow] {'; '.join(headless_indicators)}")
79
+ console.print("[yellow]→ Browser will NOT open automatically. Please use the URL below.[/yellow]\n")
80
  else:
81
+ # Only log to debug, no console output
82
  lib_logger.debug("GUI environment detected, browser auto-open will be attempted")
83
 
84
  return is_headless