Refine profile handling in prime config

Author: willccbb

packages/prime/src/prime_cli/commands/config.py

@@ -16,6 +16,16 @@
 
 # Team ID validation pattern: CUID (v1)
 TEAM_ID_PATTERN = re.compile(r"^c[a-z0-9]{24}$")
+PROFILE_OVERRIDE_ENV_VARS = (
+    "PRIME_API_KEY",
+    "PRIME_TEAM_ID",
+    "PRIME_USER_ID",
+    "PRIME_API_BASE_URL",
+    "PRIME_BASE_URL",
+    "PRIME_FRONTEND_URL",
+    "PRIME_INFERENCE_URL",
+    "PRIME_CONTEXT",
+)
 
 
 def validate_team_id(team_id: str) -> bool:
@@ -32,6 +42,34 @@ def validate_team_id(team_id: str) -> bool:
     return bool(TEAM_ID_PATTERN.match(team_id))
 
 
+def _active_profile_override_env_vars() -> list[str]:
+    return [name for name in PROFILE_OVERRIDE_ENV_VARS if (os.getenv(name) or "").strip()]
+
+
+def _require_profile_env_unset(command: str) -> None:
+    names = _active_profile_override_env_vars()
+    if not names:
+        return
+    joined = ", ".join(names)
+    console.print(
+        f"[red]Error:[/red] {joined} {'is' if len(names) == 1 else 'are'} set in your "
+        f"environment, so [bold]prime config {command}[/bold] cannot make a saved profile "
+        "active. Unset the environment override and rerun the command."
+    )
+    raise typer.Exit(1)
+
+
+def _require_no_context_override(command: str) -> None:
+    context = Config.context_override()
+    if not context:
+        return
+    console.print(
+        f"[red]Error:[/red] [bold]prime config {command}[/bold] cannot update auth while "
+        f"PRIME_CONTEXT is set to '{context}'. Run the command without a temporary context."
+    )
+    raise typer.Exit(1)
+
+
 @app.command()
 def view() -> None:
     """View current configuration"""
@@ -116,6 +154,8 @@ def set_api_key(
     ),
 ) -> None:
     """Set your API key (prompts securely if not provided)"""
+    _require_no_context_override("set-api-key")
+
     if api_key is None:
         # Interactive mode with secure prompt
         api_key = typer.prompt(
@@ -126,6 +166,7 @@ def set_api_key(
         )
 
     config = Config()
+    config.use_production_environment()
     config.set_api_key(api_key)
 
     if api_key:
@@ -140,7 +181,6 @@ def set_api_key(
                 user_id = data.get("id")
                 if user_id:
                     config.set_user_id(user_id)
-                    config.update_current_environment_file()
         except (APIError, Exception):
             pass
 
@@ -161,6 +201,7 @@ def set_team_id(
     ),
 ) -> None:
     """Set your team ID."""
+    _require_no_context_override("set-team-id")
     config = Config()
 
     # Validate team ID format
@@ -186,6 +227,7 @@ def set_team_id(
         except (APIError, Exception):
             pass
 
+    config.use_production_environment()
     config.set_team(team_id, team_name=team_name, team_role=team_role)
     if team_id:
         if team_name:
@@ -199,7 +241,9 @@ def set_team_id(
 @app.command()
 def remove_team_id() -> None:
     """Remove team ID to use personal account"""
+    _require_no_context_override("remove-team-id")
     config = Config()
+    config.use_production_environment()
     config.set_team(None)
     console.print("[green]Team ID removed. Using personal account.[/green]")
 
@@ -278,6 +322,7 @@ def _set_environment(
     env: str,
 ) -> None:
     """Set URLs for a specific environment"""
+    _require_profile_env_unset(f"use {env}")
     config = Config()
 
     # Try to load the environment (handles both built-in and custom)
@@ -384,7 +429,7 @@ def reset(
         config.set_base_url(Config.DEFAULT_BASE_URL)
         config.set_frontend_url(Config.DEFAULT_FRONTEND_URL)
         config.set_ssh_key_path(Config.DEFAULT_SSH_KEY_PATH)
-        config.set_current_environment("production")
+        config.use_production_environment()
         console.print("[green]Configuration reset to defaults![/green]")
 
 

packages/prime/src/prime_cli/commands/login.py

@@ -20,14 +20,25 @@
 console = get_console()
 
 
+def _require_no_context_override() -> None:
+    context = Config.context_override()
+    if not context:
+        return
+    console.print(
+        f"[red]Error:[/red] [bold]prime login[/bold] cannot update auth while "
+        f"PRIME_CONTEXT is set to '{context}'. Run the command without a temporary context."
+    )
+    raise typer.Exit(1)
+
+
 def fetch_and_select_team(client: APIClient, config: Config) -> None:
     """Fetch user's teams and prompt for selection."""
     try:
         teams = fetch_teams(client)
 
         if not teams:
+            config.use_production_environment()
             config.set_team(None)
-            config.update_current_environment_file()
             return
 
         console.print("\n[bold]Select:[/bold]\n")
@@ -49,8 +60,8 @@ def fetch_and_select_team(client: APIClient, config: Config) -> None:
                 selection = typer.prompt("Select", type=int, default=1)
 
                 if selection == 1:
+                    config.use_production_environment()
                     config.set_team(None)
-                    config.update_current_environment_file()
                     console.print("[green]Using personal account.[/green]")
                     return
 
@@ -62,27 +73,27 @@ def fetch_and_select_team(client: APIClient, config: Config) -> None:
 
                     if not team_id:
                         console.print("[yellow]Invalid team. Using personal account.[/yellow]")
+                        config.use_production_environment()
                         config.set_team(None)
-                        config.update_current_environment_file()
                         return
 
+                    config.use_production_environment()
                     config.set_team(team_id, team_name=team_name, team_role=team_role)
-                    config.update_current_environment_file()
                     console.print(f"[green]Using team '{team_name}'.[/green]")
                     return
 
                 console.print(f"[red]Invalid selection. Enter 1-{len(teams) + 1}.[/red]")
             except Abort:
+                config.use_production_environment()
                 config.set_team(None)
-                config.update_current_environment_file()
                 return
 
     except Abort:
+        config.use_production_environment()
         config.set_team(None)
-        config.update_current_environment_file()
     except (APIError, Exception):
+        config.use_production_environment()
         config.set_team(None)
-        config.update_current_environment_file()
 
 
 def generate_ephemeral_keypair() -> tuple[rsa.RSAPrivateKey, str]:
@@ -127,6 +138,7 @@ def login(
     headless: bool = typer.Option(False, "--headless", help="Don't attempt to open browser"),
 ) -> None:
     """Login to Prime Intellect"""
+    _require_no_context_override()
     config = Config()
     settings = config.view()
 
@@ -204,9 +216,8 @@ def login(
                     if decrypted_result:
                         # Update config with decrypted token
                         api_key = decrypted_result.decode()
+                        config.use_production_environment()
                         config.set_api_key(api_key)
-                        # Also update the current environment's saved file
-                        config.update_current_environment_file()
 
                         # Attempt to fetch the current user id
                         client = APIClient(api_key=api_key)
@@ -219,7 +230,6 @@ def login(
                                 user_id = data.get("id")
                                 if user_id:
                                     config.set_user_id(user_id)
-                                    config.update_current_environment_file()
                         except (APIError, Exception):
                             console.print("[yellow]Logged in, but failed to fetch user id[/yellow]")
 

packages/prime/src/prime_cli/commands/switch.py

@@ -18,6 +18,17 @@
 PERSONAL_TARGET = "personal"
 
 
+def _require_no_context_override() -> None:
+    context = Config.context_override()
+    if not context:
+        return
+    console.print(
+        f"[red]Error:[/red] [bold]prime switch[/bold] cannot update auth while "
+        f"PRIME_CONTEXT is set to '{context}'. Run the command without a temporary context."
+    )
+    raise typer.Exit(1)
+
+
 def _select_team_by_target(teams: list[dict], target: str) -> Optional[dict]:
     normalized_target = target.strip().lower()
 
@@ -35,8 +46,8 @@ def _select_team_by_target(teams: list[dict], target: str) -> Optional[dict]:
 
 
 def _switch_to_personal(config: Config) -> None:
+    config.use_production_environment()
     config.set_team(None)
-    config.update_current_environment_file()
     console.print("[green]Switched to personal account.[/green]")
 
 
@@ -49,8 +60,8 @@ def _switch_to_team(config: Config, team: dict) -> None:
         console.print("[red]Error:[/red] Selected team is missing a team ID.")
         raise typer.Exit(1)
 
+    config.use_production_environment()
     config.set_team(team_id, team_name=team_name, team_role=team_role)
-    config.update_current_environment_file()
     console.print(f"[green]Switched to team '{team_name}'.[/green]")
 
 
@@ -67,6 +78,7 @@ def switch(
     ),
 ) -> None:
     """Switch the active account context."""
+    _require_no_context_override()
     config = Config()
 
     if config.team_id_from_env:

packages/prime/src/prime_cli/commands/whoami.py

@@ -31,9 +31,9 @@ def whoami() -> None:
 
         # Update config
         config = Config()
-        if user_id:
+        if user_id and not Config.context_override():
+            config.use_production_environment()
             config.set_user_id(user_id)
-            config.update_current_environment_file()
 
         # Display account info table
         table = Table(title="Account")

packages/prime/src/prime_cli/core/config.py

@@ -41,6 +41,13 @@ def __init__(self) -> None:
         if context:
             self.load_environment(context, persist=False)
 
+    @staticmethod
+    def context_override() -> str | None:
+        context = os.getenv("PRIME_CONTEXT")
+        if context and context.strip():
+            return context.strip()
+        return None
+
     @staticmethod
     def _strip_api_v1(url: str) -> str:
         # make base_url consistent even if user passed a /api/v1 variant
@@ -214,6 +221,15 @@ def set_current_environment(self, value: str) -> None:
         self.config["current_environment"] = value
         self._save_config(self.config)
 
+    def use_production_environment(self) -> None:
+        """Use the built-in production endpoint configuration."""
+        if self.current_environment != "production":
+            self.config["base_url"] = self.DEFAULT_BASE_URL
+            self.config["frontend_url"] = self.DEFAULT_FRONTEND_URL
+            self.config["inference_url"] = self.DEFAULT_INFERENCE_URL
+        self.config["current_environment"] = "production"
+        self._save_config(self.config)
+
     def _sanitize_environment_name(self, name: str) -> str:
         """Sanitize environment name to prevent path traversal"""
         # Only allow alphanumeric characters, hyphens, and underscores
@@ -354,29 +370,6 @@ def load_environment(self, name: str, persist: bool = True) -> bool:
             raise
         return False
 
-    def update_current_environment_file(self) -> None:
-        """Update the current environment's saved file with current config"""
-        if self.current_environment != "production":
-            # Only update custom environments, not the built-in production
-            try:
-                sanitized_name = self._sanitize_environment_name(self.current_environment)
-                env_file = self.environments_dir / f"{sanitized_name}.json"
-                if env_file.exists():
-                    env_config = {
-                        "api_key": self.api_key,
-                        "team_id": self.team_id,
-                        "team_name": None if self.team_id_from_env else self.team_name,
-                        "team_role": None if self.team_id_from_env else self.team_role,
-                        "user_id": self.user_id,
-                        "base_url": self.base_url,
-                        "frontend_url": self.frontend_url,
-                        "inference_url": self.inference_url,
-                    }
-                    env_file.write_text(json.dumps(env_config, indent=2))
-            except ValueError:
-                # Skip updating if environment name is invalid
-                pass
-
     def list_environments(self) -> list[str]:
         """List all saved environment names"""
         environments = ["production"]  # Built-in environment

packages/prime/src/prime_cli/main.py

@@ -107,7 +107,17 @@ def callback(
                 typer.echo(f"  - {env_name}", err=True)
             raise typer.Exit(1)
 
-        # Set environment variable so Config instances in subcommands pick it up
+        previous_context = os.environ.get("PRIME_CONTEXT")
+
+        def restore_context() -> None:
+            if previous_context is None:
+                os.environ.pop("PRIME_CONTEXT", None)
+            else:
+                os.environ["PRIME_CONTEXT"] = previous_context
+
+        ctx.call_on_close(restore_context)
+
+        # Set environment variable so Config instances in subcommands pick it up.
         os.environ["PRIME_CONTEXT"] = context
 
     # Check for updates (only when a subcommand is being executed)

packages/prime/tests/test_config_delete.py

@@ -11,13 +11,32 @@
     "COLUMNS": "200",
     "LINES": "50",
     "PRIME_DISABLE_VERSION_CHECK": "1",
+    "PRIME_API_KEY": "",
+    "PRIME_TEAM_ID": "",
+    "PRIME_USER_ID": "",
+    "PRIME_API_BASE_URL": "",
+    "PRIME_BASE_URL": "",
+    "PRIME_FRONTEND_URL": "",
+    "PRIME_INFERENCE_URL": "",
+    "PRIME_CONTEXT": "",
 }
 
 
 @pytest.fixture
 def temp_home(tmp_path: Any, monkeypatch: pytest.MonkeyPatch) -> Path:
     monkeypatch.setenv("HOME", str(tmp_path))
     monkeypatch.setattr("prime_cli.main.check_for_update", lambda: (False, None))
+    for name in (
+        "PRIME_API_KEY",
+        "PRIME_TEAM_ID",
+        "PRIME_USER_ID",
+        "PRIME_API_BASE_URL",
+        "PRIME_BASE_URL",
+        "PRIME_FRONTEND_URL",
+        "PRIME_INFERENCE_URL",
+        "PRIME_CONTEXT",
+    ):
+        monkeypatch.delenv(name, raising=False)
     return tmp_path