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,23 @@ 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)
+
+
 @app.command()
 def view() -> None:
     """View current configuration"""
@@ -127,6 +154,7 @@ def set_api_key(
 
     config = Config()
     config.set_api_key(api_key)
+    config.update_current_environment_file()
 
     if api_key:
         masked_key = f"{api_key[:6]}***{api_key[-4:]}" if len(api_key) > 10 else "***"
@@ -187,6 +215,7 @@ def set_team_id(
             pass
 
     config.set_team(team_id, team_name=team_name, team_role=team_role)
+    config.update_current_environment_file()
     if team_id:
         if team_name:
             console.print(f"[green]Team '{team_name}' ({team_id}) configured successfully![/green]")
@@ -201,6 +230,7 @@ def remove_team_id() -> None:
     """Remove team ID to use personal account"""
     config = Config()
     config.set_team(None)
+    config.update_current_environment_file()
     console.print("[green]Team ID removed. Using personal account.[/green]")
 
 
@@ -278,6 +308,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)

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

@@ -355,27 +355,30 @@ def load_environment(self, name: str, persist: bool = True) -> bool:
         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
+        """Update the active saved environment with the persisted config values."""
+        if self.current_environment == "production":
+            return
+
+        try:
+            sanitized_name = self._sanitize_environment_name(self.current_environment)
+        except ValueError:
+            return
+
+        env_file = self.environments_dir / f"{sanitized_name}.json"
+        if not env_file.exists():
+            return
+
+        env_config = {
+            "api_key": self.config.get("api_key", ""),
+            "team_id": self.config.get("team_id"),
+            "team_name": self.config.get("team_name"),
+            "team_role": self.config.get("team_role"),
+            "user_id": self.config.get("user_id"),
+            "base_url": self.config.get("base_url", self.DEFAULT_BASE_URL),
+            "frontend_url": self.config.get("frontend_url", self.DEFAULT_FRONTEND_URL),
+            "inference_url": self.config.get("inference_url", self.DEFAULT_INFERENCE_URL),
+        }
+        env_file.write_text(json.dumps(env_config, indent=2))
 
     def list_environments(self) -> list[str]:
         """List all saved environment names"""

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