diff --git a/src/codegen/cli/auth/login.py b/src/codegen/cli/auth/login.py index 49fddd5c8..4205cd8d2 100644 --- a/src/codegen/cli/auth/login.py +++ b/src/codegen/cli/auth/login.py @@ -37,9 +37,10 @@ def login_routine(token: str | None = None) -> str: # Validate and store token try: + rich.print("[blue]Validating token and fetching account info...[/blue]") token_manager = TokenManager() token_manager.authenticate_token(token) - rich.print(f"[green]✓ Stored token to:[/green] {token_manager.token_file}") + rich.print(f"[green]✓ Stored token and profile to:[/green] {token_manager.token_file}") return token except AuthError as e: rich.print(f"[red]Error:[/red] {e!s}") diff --git a/src/codegen/cli/auth/token_manager.py b/src/codegen/cli/auth/token_manager.py index 11e7dbb16..e25e22b52 100644 --- a/src/codegen/cli/auth/token_manager.py +++ b/src/codegen/cli/auth/token_manager.py @@ -4,6 +4,10 @@ from codegen.cli.auth.constants import AUTH_FILE, CONFIG_DIR +# Simple cache to avoid repeated file I/O +_token_cache = None +_cache_mtime = None + class TokenManager: # Simple token manager to store and retrieve tokens. @@ -20,17 +24,79 @@ def _ensure_config_dir(self): Path(self.config_dir).mkdir(parents=True, exist_ok=True) def authenticate_token(self, token: str) -> None: - """Store the token locally.""" - self.save_token(token) + """Store the token locally and fetch organization info.""" + self.save_token_with_org_info(token) + + def save_token_with_org_info(self, token: str) -> None: + """Save api token to disk along with organization info.""" + global _token_cache, _cache_mtime + + # First fetch organization info using the token + try: + import requests + + from codegen.cli.api.endpoints import API_ENDPOINT + + headers = {"Authorization": f"Bearer {token}"} + + # Test token by getting user info + user_response = requests.get(f"{API_ENDPOINT.rstrip('/')}/v1/users/me", headers=headers, timeout=10) + user_response.raise_for_status() + user_data = user_response.json() + + # Get organizations + org_response = requests.get(f"{API_ENDPOINT.rstrip('/')}/v1/organizations", headers=headers, timeout=10) + org_response.raise_for_status() + org_data = org_response.json() + + # Prepare auth data with org info + auth_data = { + "token": token, + "user": {"id": user_data.get("id"), "email": user_data.get("email"), "full_name": user_data.get("full_name"), "github_username": user_data.get("github_username")}, + } + + # Add organization info if available + orgs = org_data.get("items", []) + if orgs and len(orgs) > 0: + primary_org = orgs[0] # Use first org as primary + auth_data["organization"] = {"id": primary_org.get("id"), "name": primary_org.get("name"), "all_orgs": [{"id": org.get("id"), "name": org.get("name")} for org in orgs]} + + except requests.RequestException as e: + # If we can't fetch org info, still save the token but without org data + print(f"Warning: Could not fetch organization info: {e}") + auth_data = {"token": token} + except Exception as e: + print(f"Warning: Error fetching user/org info: {e}") + auth_data = {"token": token} + + # Save to file + try: + with open(self.token_file, "w") as f: + json.dump(auth_data, f, indent=2) + + # Secure the file permissions (read/write for owner only) + os.chmod(self.token_file, 0o600) + + # Invalidate cache + _token_cache = None + _cache_mtime = None + except Exception as e: + print(f"Error saving token: {e!s}") + raise def save_token(self, token: str) -> None: - """Save api token to disk.""" + """Save api token to disk (legacy method - just saves token).""" + global _token_cache, _cache_mtime try: with open(self.token_file, "w") as f: json.dump({"token": token}, f) # Secure the file permissions (read/write for owner only) os.chmod(self.token_file, 0o600) + + # Invalidate cache + _token_cache = None + _cache_mtime = None except Exception as e: print(f"Error saving token: {e!s}") raise @@ -58,8 +124,52 @@ def get_token(self) -> str | None: def clear_token(self) -> None: """Remove stored token.""" + global _token_cache, _cache_mtime if os.path.exists(self.token_file): os.remove(self.token_file) + # Invalidate cache + _token_cache = None + _cache_mtime = None + + def get_auth_data(self) -> dict | None: + """Retrieve complete auth data from disk.""" + try: + if not os.access(self.config_dir, os.R_OK): + return None + + if not os.path.exists(self.token_file): + return None + + with open(self.token_file) as f: + return json.load(f) + except Exception: + return None + + def get_org_id(self) -> int | None: + """Get the stored organization ID.""" + auth_data = self.get_auth_data() + if auth_data and "organization" in auth_data: + org_id = auth_data["organization"].get("id") + if org_id: + try: + return int(org_id) + except (ValueError, TypeError): + return None + return None + + def get_org_name(self) -> str | None: + """Get the stored organization name.""" + auth_data = self.get_auth_data() + if auth_data and "organization" in auth_data: + return auth_data["organization"].get("name") + return None + + def get_user_info(self) -> dict | None: + """Get the stored user info.""" + auth_data = self.get_auth_data() + if auth_data and "user" in auth_data: + return auth_data["user"] + return None def get_current_token() -> str | None: @@ -67,11 +177,67 @@ def get_current_token() -> str | None: This is a helper function that creates a TokenManager instance and retrieves the stored token. The token is validated before being returned. + Uses a simple cache to avoid repeated file I/O. Returns: Optional[str]: The current valid api token if one exists. Returns None if no token exists. """ + global _token_cache, _cache_mtime + + try: + # Check if token file exists + if not os.path.exists(AUTH_FILE): + return None + + # Get file modification time + current_mtime = os.path.getmtime(AUTH_FILE) + + # Use cache if file hasn't changed + if _token_cache is not None and _cache_mtime == current_mtime: + return _token_cache + + # Read token from file + token_manager = TokenManager() + token = token_manager.get_token() + + # Update cache + _token_cache = token + _cache_mtime = current_mtime + + return token + except Exception: + # Fall back to uncached version on any error + token_manager = TokenManager() + return token_manager.get_token() + + +def get_current_org_id() -> int | None: + """Get the stored organization ID if available. + + Returns: + Optional[int]: The organization ID if stored, None otherwise. + """ + token_manager = TokenManager() + return token_manager.get_org_id() + + +def get_current_org_name() -> str | None: + """Get the stored organization name if available. + + Returns: + Optional[str]: The organization name if stored, None otherwise. + """ + token_manager = TokenManager() + return token_manager.get_org_name() + + +def get_current_user_info() -> dict | None: + """Get the stored user info if available. + + Returns: + Optional[dict]: The user info if stored, None otherwise. + """ token_manager = TokenManager() - return token_manager.get_token() + return token_manager.get_user_info() diff --git a/src/codegen/cli/commands/agents/main.py b/src/codegen/cli/commands/agents/main.py index bf1381bc6..ebdca3427 100644 --- a/src/codegen/cli/commands/agents/main.py +++ b/src/codegen/cli/commands/agents/main.py @@ -26,18 +26,19 @@ def list_agents(org_id: int | None = typer.Option(None, help="Organization ID (d raise typer.Exit(1) try: - # Resolve org id + # Resolve org id (now fast, uses stored data) resolved_org_id = resolve_org_id(org_id) if resolved_org_id is None: console.print("[red]Error:[/red] Organization ID not provided. Pass --org-id, set CODEGEN_ORG_ID, or REPOSITORY_ORG_ID.") raise typer.Exit(1) - # Make API request to list agent runs with spinner + # Start spinner for API calls only spinner = create_spinner("Fetching your recent API agent runs...") spinner.start() try: headers = {"Authorization": f"Bearer {token}"} + # Filter to only API source type and current user's agent runs params = { "source_type": "API", diff --git a/src/codegen/cli/commands/profile/main.py b/src/codegen/cli/commands/profile/main.py index d615c04f5..d5c6d4795 100644 --- a/src/codegen/cli/commands/profile/main.py +++ b/src/codegen/cli/commands/profile/main.py @@ -7,7 +7,7 @@ from rich.panel import Panel from codegen.cli.api.endpoints import API_ENDPOINT -from codegen.cli.auth.token_manager import get_current_token +from codegen.cli.auth.token_manager import get_current_org_name, get_current_token, get_current_user_info from codegen.cli.rich.spinners import create_spinner from codegen.cli.utils.org import resolve_org_id @@ -22,89 +22,89 @@ def profile(): console.print("[red]Error:[/red] Not authenticated. Please run 'codegen login' first.") raise typer.Exit(1) - try: - # Make API request to get current user info with spinner - spinner = create_spinner("Fetching user profile and organization info...") + # Try to get stored user and org info first (fast, no API calls) + user_info = get_current_user_info() + org_name = get_current_org_name() + org_id = resolve_org_id() # This now uses stored data first + + # If we have stored data, use it directly + if user_info and user_info.get("id"): + user_id = user_info.get("id", "Unknown") + full_name = user_info.get("full_name", "") + email = user_info.get("email", "") + github_username = user_info.get("github_username", "") + role = "Member" # Default role for stored data + else: + # Fall back to API call if no stored data + spinner = create_spinner("Fetching user profile info...") spinner.start() - try: headers = {"Authorization": f"Bearer {token}"} - url = f"{API_ENDPOINT.rstrip('/')}/v1/users/me" - response = requests.get(url, headers=headers) - response.raise_for_status() - user_data = response.json() + user_response = requests.get(f"{API_ENDPOINT.rstrip('/')}/v1/users/me", headers=headers) + user_response.raise_for_status() + user_data = user_response.json() + + user_id = user_data.get("id", "Unknown") + full_name = user_data.get("full_name", "") + email = user_data.get("email", "") + github_username = user_data.get("github_username", "") + role = user_data.get("role", "Member") + except requests.RequestException as e: + spinner.stop() + console.print(f"[red]Error:[/red] Failed to fetch profile information: {e}") + raise typer.Exit(1) finally: spinner.stop() - # Extract user information - user_id = user_data.get("id", "Unknown") - full_name = user_data.get("full_name", "") - email = user_data.get("email", "") - github_username = user_data.get("github_username", "") - github_user_id = user_data.get("github_user_id", "") - avatar_url = user_data.get("avatar_url", "") - role = user_data.get("role", "") - - # Get organization information - org_id = resolve_org_id() - org_name = None - - if org_id: - try: - # Fetch organizations to get the name - orgs_url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations" - orgs_response = requests.get(orgs_url, headers=headers) - orgs_response.raise_for_status() - orgs_data = orgs_response.json() - - # Find the matching organization - organizations = orgs_data.get("items", []) - for org in organizations: - if org.get("id") == org_id: - org_name = org.get("name", f"Organization {org_id}") - break - - if not org_name: - org_name = f"Organization {org_id}" # Fallback if not found - - except Exception: - # If we can't fetch org name, fall back to showing ID - org_name = f"Organization {org_id}" - - # Create profile display - profile_info = [] - if user_id != "Unknown": - profile_info.append(f"[cyan]User ID:[/cyan] {user_id}") - if full_name: - profile_info.append(f"[cyan]Name:[/cyan] {full_name}") - if email: - profile_info.append(f"[cyan]Email:[/cyan] {email}") - if github_username: - profile_info.append(f"[cyan]GitHub:[/cyan] {github_username}") - if org_name: - profile_info.append(f"[cyan]Organization:[/cyan] {org_name}") - elif org_id: - profile_info.append(f"[cyan]Organization:[/cyan] Organization {org_id}") - else: - profile_info.append("[cyan]Organization:[/cyan] [yellow]Not configured[/yellow] (set CODEGEN_ORG_ID or REPOSITORY_ORG_ID)") - if role: - profile_info.append(f"[cyan]Role:[/cyan] {role}") - - profile_text = "\n".join(profile_info) if profile_info else "No profile information available" + # If no stored org name but we have an org_id, try to fetch it + if org_id and not org_name: + spinner = create_spinner("Fetching organization info...") + spinner.start() + try: + headers = {"Authorization": f"Bearer {token}"} + orgs_response = requests.get(f"{API_ENDPOINT.rstrip('/')}/v1/organizations", headers=headers) + orgs_response.raise_for_status() + orgs_data = orgs_response.json() + + # Find the organization by ID + orgs = orgs_data.get("items", []) + for org in orgs: + if org.get("id") == org_id: + org_name = org.get("name") + break + except requests.RequestException: + # Ignore errors for org name lookup - not critical + pass + finally: + spinner.stop() - console.print( - Panel( - profile_text, - title="👤 [bold]User Profile[/bold]", - border_style="cyan", - box=box.ROUNDED, - padding=(1, 2), - ) + # Build profile information + profile_info = [] + if user_id != "Unknown": + profile_info.append(f"[cyan]User ID:[/cyan] {user_id}") + if full_name: + profile_info.append(f"[cyan]Name:[/cyan] {full_name}") + if email: + profile_info.append(f"[cyan]Email:[/cyan] {email}") + if github_username: + profile_info.append(f"[cyan]GitHub:[/cyan] {github_username}") + if org_name: + profile_info.append(f"[cyan]Organization:[/cyan] {org_name}") + elif org_id: + profile_info.append(f"[cyan]Organization:[/cyan] Organization {org_id}") + else: + profile_info.append("[cyan]Organization:[/cyan] [yellow]Not configured[/yellow] (set CODEGEN_ORG_ID or REPOSITORY_ORG_ID)") + if role: + profile_info.append(f"[cyan]Role:[/cyan] {role}") + + profile_text = "\n".join(profile_info) if profile_info else "No profile information available" + + console.print( + Panel( + profile_text, + title="👤 [bold]User Profile[/bold]", + border_style="cyan", + box=box.ROUNDED, + padding=(1, 2), ) - - except requests.RequestException as e: - console.print(f"[red]Error fetching profile:[/red] {e}", style="bold red") - raise typer.Exit(1) - except Exception as e: - console.print(f"[red]Unexpected error:[/red] {e}", style="bold red") - raise typer.Exit(1) + ) diff --git a/src/codegen/cli/utils/org.py b/src/codegen/cli/utils/org.py index 06c2f7f57..bbb800751 100644 --- a/src/codegen/cli/utils/org.py +++ b/src/codegen/cli/utils/org.py @@ -1,15 +1,20 @@ """Organization resolution utilities for CLI commands.""" import os -from typing import Optional +import time import requests from codegen.cli.api.endpoints import API_ENDPOINT -from codegen.cli.auth.token_manager import get_current_token +from codegen.cli.auth.token_manager import get_current_org_id, get_current_token from codegen.cli.commands.claude.quiet_console import console -def resolve_org_id(explicit_org_id: Optional[int] = None) -> Optional[int]: +# Cache for org resolution to avoid repeated API calls +_org_cache = {} +_cache_timeout = 300 # 5 minutes + + +def resolve_org_id(explicit_org_id: int | None = None) -> int | None: """Resolve the organization id from CLI input or environment. Order of precedence: @@ -18,6 +23,7 @@ def resolve_org_id(explicit_org_id: Optional[int] = None) -> Optional[int]: Returns None if not found. """ + global _org_cache if explicit_org_id is not None: return explicit_org_id @@ -32,27 +38,47 @@ def resolve_org_id(explicit_org_id: Optional[int] = None) -> Optional[int]: except ValueError: pass + # Try stored org ID from auth data (fast, no API call) + stored_org_id = get_current_org_id() + if stored_org_id: + return stored_org_id + # Attempt auto-detection via API: if user belongs to organizations, use the first try: token = get_current_token() if not token: print("No token found") return None + + # Check cache first + cache_key = f"org_auto_detect_{token[:10]}" # Use first 10 chars as key + current_time = time.time() + + if cache_key in _org_cache: + cached_data, cache_time = _org_cache[cache_key] + if current_time - cache_time < _cache_timeout: + return cached_data + headers = {"Authorization": f"Bearer {token}"} url = f"{API_ENDPOINT.rstrip('/')}/v1/organizations" resp = requests.get(url, headers=headers, timeout=10) resp.raise_for_status() data = resp.json() items = data.get("items") or [] + + org_id = None if isinstance(items, list) and len(items) >= 1: org = items[0] - org_id = org.get("id") + org_id_raw = org.get("id") try: - return int(org_id) + org_id = int(org_id_raw) except Exception: - return None - # None returned - return None + org_id = None + + # Cache the result + _org_cache[cache_key] = (org_id, current_time) + return org_id + except Exception as e: console.print(f"Exception: {e}") return None @@ -61,4 +87,3 @@ def resolve_org_id(explicit_org_id: Optional[int] = None) -> Optional[int]: return int(env_val) except ValueError: return None -