Explicit org selector in codegen profile and codegen login (#1208)

Author: jayhack

src/codegen/cli/auth/login.py

@@ -4,9 +4,10 @@
 import typer
 
 from codegen.cli.api.webapp_routes import USER_SECRETS_ROUTE
-from codegen.cli.auth.token_manager import TokenManager
+from codegen.cli.auth.token_manager import TokenManager, get_cached_organizations, set_default_organization
 from codegen.cli.env.global_env import global_env
 from codegen.cli.errors import AuthError
+from codegen.cli.utils.simple_selector import simple_org_selector
 
 
 def login_routine(token: str | None = None) -> str:
@@ -27,20 +28,47 @@ def login_routine(token: str | None = None) -> str:
 
     # If no token provided, guide user through browser flow
     if not token:
-        rich.print(f"Opening {USER_SECRETS_ROUTE} to get your authentication token...")
         webbrowser.open_new(USER_SECRETS_ROUTE)
-        token = typer.prompt("Please enter your authentication token from the browser", hide_input=False)
+        token = typer.prompt(f"Enter your token from {USER_SECRETS_ROUTE}", hide_input=False)
 
     if not token:
         rich.print("[red]Error:[/red] Token must be provided via CODEGEN_USER_ACCESS_TOKEN environment variable or manual input")
         raise typer.Exit(1)
 
     # 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 and profile to:[/green] {token_manager.token_file}")
+
+        # Show organization selector if multiple organizations available
+        organizations = get_cached_organizations()
+        if organizations and len(organizations) > 1:
+            rich.print("\n[blue]Multiple organizations found. Please select your default:[/blue]")
+            selected_org = simple_org_selector(organizations, title="🏢 Select Default Organization")
+
+            if selected_org:
+                org_id = selected_org.get("id")
+                org_name = selected_org.get("name")
+                try:
+                    set_default_organization(org_id, org_name)
+                    rich.print(f"[green]✓ Set default organization:[/green] {org_name}")
+                except Exception as e:
+                    rich.print(f"[yellow]Warning: Could not set default organization: {e}[/yellow]")
+                    rich.print("[yellow]You can set it later with 'codegen profile'[/yellow]")
+            else:
+                rich.print("[yellow]No organization selected. You can set it later with 'codegen profile'[/yellow]")
+        elif organizations and len(organizations) == 1:
+            # Single organization - set it automatically
+            org = organizations[0]
+            org_id = org.get("id")
+            org_name = org.get("name")
+            try:
+                set_default_organization(org_id, org_name)
+                rich.print(f"[green]✓ Set default organization:[/green] {org_name}")
+            except Exception as e:
+                rich.print(f"[yellow]Warning: Could not set default organization: {e}[/yellow]")
+
         return token
     except AuthError as e:
         rich.print(f"[red]Error:[/red] {e!s}")

src/codegen/cli/auth/token_manager.py

@@ -61,11 +61,7 @@ def save_token_with_org_info(self, token: str) -> None:
                 # Store ALL organizations in cache for local resolution
                 all_orgs = [{"id": org.get("id"), "name": org.get("name")} for org in orgs]
                 primary_org = orgs[0]  # Use first org as primary/default
-                auth_data["organization"] = {
-                    "id": primary_org.get("id"), 
-                    "name": primary_org.get("name"), 
-                    "all_orgs": all_orgs
-                }
+                auth_data["organization"] = {"id": primary_org.get("id"), "name": primary_org.get("name"), "all_orgs": all_orgs}
                 auth_data["organizations_cache"] = all_orgs  # Separate cache for easy access
 
         except requests.RequestException as e:
@@ -180,7 +176,7 @@ def get_user_info(self) -> dict | None:
 
     def get_cached_organizations(self) -> list[dict] | None:
         """Get all cached organizations.
-        
+
         Returns:
             List of organization dictionaries with 'id' and 'name' keys, or None if no cache.
         """
@@ -194,37 +190,75 @@ def get_cached_organizations(self) -> list[dict] | None:
 
     def is_org_id_in_cache(self, org_id: int) -> bool:
         """Check if an organization ID exists in the local cache.
-        
+
         Args:
             org_id: The organization ID to check
-            
+
         Returns:
             True if the organization ID is found in cache, False otherwise.
         """
         cached_orgs = self.get_cached_organizations()
         if not cached_orgs:
             return False
-        
+
         return any(org.get("id") == org_id for org in cached_orgs)
 
     def get_org_name_from_cache(self, org_id: int) -> str | None:
         """Get organization name from cache by ID.
-        
+
         Args:
             org_id: The organization ID to look up
-            
+
         Returns:
             Organization name if found in cache, None otherwise.
         """
         cached_orgs = self.get_cached_organizations()
         if not cached_orgs:
             return None
-            
+
         for org in cached_orgs:
             if org.get("id") == org_id:
                 return org.get("name")
         return None
 
+    def set_default_organization(self, org_id: int, org_name: str) -> None:
+        """Set the default organization in auth.json.
+
+        Args:
+            org_id: The organization ID to set as default
+            org_name: The organization name
+        """
+        auth_data = self.get_auth_data()
+        if not auth_data:
+            msg = "No authentication data found. Please run 'codegen login' first."
+            raise ValueError(msg)
+
+        # Verify the org exists in cache
+        if not self.is_org_id_in_cache(org_id):
+            msg = f"Organization {org_id} not found in cache. Please run 'codegen login' to refresh."
+            raise ValueError(msg)
+
+        # Update the organization info
+        auth_data["organization"] = {"id": org_id, "name": org_name, "all_orgs": auth_data.get("organization", {}).get("all_orgs", [])}
+
+        # Save to file
+        try:
+            import json
+
+            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
+            global _token_cache, _cache_mtime
+            _token_cache = None
+            _cache_mtime = None
+        except Exception as e:
+            msg = f"Error saving default organization: {e}"
+            raise ValueError(msg)
+
 
 def get_current_token() -> str | None:
     """Get the current authentication token if one exists.
@@ -289,7 +323,7 @@ def get_current_org_name() -> str | None:
 
 def get_cached_organizations() -> list[dict] | None:
     """Get all cached organizations.
-    
+
     Returns:
         List of organization dictionaries with 'id' and 'name' keys, or None if no cache.
     """
@@ -299,10 +333,10 @@ def get_cached_organizations() -> list[dict] | None:
 
 def is_org_id_cached(org_id: int) -> bool:
     """Check if an organization ID exists in the local cache.
-    
+
     Args:
         org_id: The organization ID to check
-        
+
     Returns:
         True if the organization ID is found in cache, False otherwise.
     """
@@ -312,10 +346,10 @@ def is_org_id_cached(org_id: int) -> bool:
 
 def get_org_name_from_cache(org_id: int) -> str | None:
     """Get organization name from cache by ID.
-    
+
     Args:
         org_id: The organization ID to look up
-        
+
     Returns:
         Organization name if found in cache, None otherwise.
     """
@@ -335,9 +369,10 @@ def get_current_user_info() -> dict | None:
 
 # Repository caching functions (similar to organization caching)
 
+
 def get_cached_repositories() -> list[dict] | None:
     """Get all cached repositories.
-    
+
     Returns:
         List of repository dictionaries with 'id' and 'name' keys, or None if no cache.
     """
@@ -350,7 +385,7 @@ def get_cached_repositories() -> list[dict] | None:
 
 def cache_repositories(repositories: list[dict]) -> None:
     """Cache repositories to local storage.
-    
+
     Args:
         repositories: List of repository dictionaries to cache
     """
@@ -361,53 +396,65 @@ def cache_repositories(repositories: list[dict]) -> None:
         # Save back to file
         try:
             import json
-            with open(token_manager.token_file, 'w') as f:
+
+            with open(token_manager.token_file, "w") as f:
                 json.dump(auth_data, f, indent=2)
         except Exception:
             pass  # Fail silently
 
 
 def is_repo_id_cached(repo_id: int) -> bool:
     """Check if a repository ID exists in the local cache.
-    
+
     Args:
         repo_id: The repository ID to check
-        
+
     Returns:
         True if the repository ID is found in cache, False otherwise.
     """
     cached_repos = get_cached_repositories()
     if not cached_repos:
         return False
-    
+
     return any(repo.get("id") == repo_id for repo in cached_repos)
 
 
 def get_repo_name_from_cache(repo_id: int) -> str | None:
     """Get repository name from cache by ID.
-    
+
     Args:
         repo_id: The repository ID to look up
-        
+
     Returns:
         Repository name if found in cache, None otherwise.
     """
     cached_repos = get_cached_repositories()
     if not cached_repos:
         return None
-    
+
     for repo in cached_repos:
         if repo.get("id") == repo_id:
             return repo.get("name")
-    
+
     return None
 
 
 def get_current_repo_name() -> str | None:
     """Get the current repository name from environment or cache."""
     from codegen.cli.utils.repo import get_current_repo_id
-    
+
     repo_id = get_current_repo_id()
     if repo_id:
         return get_repo_name_from_cache(repo_id)
     return None
+
+
+def set_default_organization(org_id: int, org_name: str) -> None:
+    """Set the default organization in auth.json.
+
+    Args:
+        org_id: The organization ID to set as default
+        org_name: The organization name
+    """
+    token_manager = TokenManager()
+    return token_manager.set_default_organization(org_id, org_name)

src/codegen/cli/cli.py

@@ -15,7 +15,7 @@
 from codegen.cli.commands.login.main import login
 from codegen.cli.commands.logout.main import logout
 from codegen.cli.commands.org.main import org
-from codegen.cli.commands.profile.main import profile
+from codegen.cli.commands.profile.main import profile_app
 from codegen.cli.commands.repo.main import repo
 from codegen.cli.commands.style_debug.main import style_debug
 from codegen.cli.commands.tools.main import tools
@@ -42,7 +42,7 @@ def version_callback(value: bool):
 main.command("login", help="Store authentication token.")(login)
 main.command("logout", help="Clear stored authentication token.")(logout)
 main.command("org", help="Manage and switch between organizations.")(org)
-main.command("profile", help="Display information about the currently authenticated user.")(profile)
+# Profile is now a Typer app
 main.command("repo", help="Manage repository configuration and environment variables.")(repo)
 main.command("style-debug", help="Debug command to visualize CLI styling (spinners, etc).")(style_debug)
 main.command("tools", help="List available tools from the Codegen API.")(tools)
@@ -53,6 +53,7 @@ def version_callback(value: bool):
 main.add_typer(agents_app, name="agents")
 main.add_typer(config_command, name="config")
 main.add_typer(integrations_app, name="integrations")
+main.add_typer(profile_app, name="profile")
 
 
 @main.callback(invoke_without_command=True)

src/codegen/cli/commands/login/main.py

@@ -1,4 +1,3 @@
-import rich
 import typer
 
 from codegen.cli.auth.login import login_routine
@@ -9,6 +8,6 @@ def login(token: str | None = typer.Option(None, help="API token for authenticat
     """Store authentication token."""
     # Check if already authenticated
     if get_current_token():
-        rich.print("[yellow]Info:[/yellow] You already have a token stored. Proceeding with re-authentication...")
+        pass  # Just proceed silently with re-authentication
 
     login_routine(token)