Skip to content

Commit 19c3309

Browse files
committed
Refine profile handling in prime config
1 parent 8dfec09 commit 19c3309

7 files changed

Lines changed: 285 additions & 38 deletions

File tree

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

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,16 @@
1616

1717
# Team ID validation pattern: CUID (v1)
1818
TEAM_ID_PATTERN = re.compile(r"^c[a-z0-9]{24}$")
19+
PROFILE_OVERRIDE_ENV_VARS = (
20+
"PRIME_API_KEY",
21+
"PRIME_TEAM_ID",
22+
"PRIME_USER_ID",
23+
"PRIME_API_BASE_URL",
24+
"PRIME_BASE_URL",
25+
"PRIME_FRONTEND_URL",
26+
"PRIME_INFERENCE_URL",
27+
"PRIME_CONTEXT",
28+
)
1929

2030

2131
def validate_team_id(team_id: str) -> bool:
@@ -32,6 +42,23 @@ def validate_team_id(team_id: str) -> bool:
3242
return bool(TEAM_ID_PATTERN.match(team_id))
3343

3444

45+
def _active_profile_override_env_vars() -> list[str]:
46+
return [name for name in PROFILE_OVERRIDE_ENV_VARS if (os.getenv(name) or "").strip()]
47+
48+
49+
def _require_profile_env_unset(command: str) -> None:
50+
names = _active_profile_override_env_vars()
51+
if not names:
52+
return
53+
joined = ", ".join(names)
54+
console.print(
55+
f"[red]Error:[/red] {joined} {'is' if len(names) == 1 else 'are'} set in your "
56+
f"environment, so [bold]prime config {command}[/bold] cannot make a saved profile "
57+
"active. Unset the environment override and rerun the command."
58+
)
59+
raise typer.Exit(1)
60+
61+
3562
@app.command()
3663
def view() -> None:
3764
"""View current configuration"""
@@ -126,6 +153,7 @@ def set_api_key(
126153
)
127154

128155
config = Config()
156+
config.use_production_environment()
129157
config.set_api_key(api_key)
130158

131159
if api_key:
@@ -140,7 +168,6 @@ def set_api_key(
140168
user_id = data.get("id")
141169
if user_id:
142170
config.set_user_id(user_id)
143-
config.update_current_environment_file()
144171
except (APIError, Exception):
145172
pass
146173

@@ -186,6 +213,7 @@ def set_team_id(
186213
except (APIError, Exception):
187214
pass
188215

216+
config.use_production_environment()
189217
config.set_team(team_id, team_name=team_name, team_role=team_role)
190218
if team_id:
191219
if team_name:
@@ -200,6 +228,7 @@ def set_team_id(
200228
def remove_team_id() -> None:
201229
"""Remove team ID to use personal account"""
202230
config = Config()
231+
config.use_production_environment()
203232
config.set_team(None)
204233
console.print("[green]Team ID removed. Using personal account.[/green]")
205234

@@ -278,6 +307,7 @@ def _set_environment(
278307
env: str,
279308
) -> None:
280309
"""Set URLs for a specific environment"""
310+
_require_profile_env_unset(f"use {env}")
281311
config = Config()
282312

283313
# Try to load the environment (handles both built-in and custom)
@@ -384,7 +414,7 @@ def reset(
384414
config.set_base_url(Config.DEFAULT_BASE_URL)
385415
config.set_frontend_url(Config.DEFAULT_FRONTEND_URL)
386416
config.set_ssh_key_path(Config.DEFAULT_SSH_KEY_PATH)
387-
config.set_current_environment("production")
417+
config.use_production_environment()
388418
console.print("[green]Configuration reset to defaults![/green]")
389419

390420

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

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ def fetch_and_select_team(client: APIClient, config: Config) -> None:
2626
teams = fetch_teams(client)
2727

2828
if not teams:
29+
config.use_production_environment()
2930
config.set_team(None)
30-
config.update_current_environment_file()
3131
return
3232

3333
console.print("\n[bold]Select:[/bold]\n")
@@ -49,8 +49,8 @@ def fetch_and_select_team(client: APIClient, config: Config) -> None:
4949
selection = typer.prompt("Select", type=int, default=1)
5050

5151
if selection == 1:
52+
config.use_production_environment()
5253
config.set_team(None)
53-
config.update_current_environment_file()
5454
console.print("[green]Using personal account.[/green]")
5555
return
5656

@@ -62,27 +62,27 @@ def fetch_and_select_team(client: APIClient, config: Config) -> None:
6262

6363
if not team_id:
6464
console.print("[yellow]Invalid team. Using personal account.[/yellow]")
65+
config.use_production_environment()
6566
config.set_team(None)
66-
config.update_current_environment_file()
6767
return
6868

69+
config.use_production_environment()
6970
config.set_team(team_id, team_name=team_name, team_role=team_role)
70-
config.update_current_environment_file()
7171
console.print(f"[green]Using team '{team_name}'.[/green]")
7272
return
7373

7474
console.print(f"[red]Invalid selection. Enter 1-{len(teams) + 1}.[/red]")
7575
except Abort:
76+
config.use_production_environment()
7677
config.set_team(None)
77-
config.update_current_environment_file()
7878
return
7979

8080
except Abort:
81+
config.use_production_environment()
8182
config.set_team(None)
82-
config.update_current_environment_file()
8383
except (APIError, Exception):
84+
config.use_production_environment()
8485
config.set_team(None)
85-
config.update_current_environment_file()
8686

8787

8888
def generate_ephemeral_keypair() -> tuple[rsa.RSAPrivateKey, str]:
@@ -204,9 +204,8 @@ def login(
204204
if decrypted_result:
205205
# Update config with decrypted token
206206
api_key = decrypted_result.decode()
207+
config.use_production_environment()
207208
config.set_api_key(api_key)
208-
# Also update the current environment's saved file
209-
config.update_current_environment_file()
210209

211210
# Attempt to fetch the current user id
212211
client = APIClient(api_key=api_key)
@@ -219,7 +218,6 @@ def login(
219218
user_id = data.get("id")
220219
if user_id:
221220
config.set_user_id(user_id)
222-
config.update_current_environment_file()
223221
except (APIError, Exception):
224222
console.print("[yellow]Logged in, but failed to fetch user id[/yellow]")
225223

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ def _select_team_by_target(teams: list[dict], target: str) -> Optional[dict]:
3535

3636

3737
def _switch_to_personal(config: Config) -> None:
38+
config.use_production_environment()
3839
config.set_team(None)
39-
config.update_current_environment_file()
4040
console.print("[green]Switched to personal account.[/green]")
4141

4242

@@ -49,8 +49,8 @@ def _switch_to_team(config: Config, team: dict) -> None:
4949
console.print("[red]Error:[/red] Selected team is missing a team ID.")
5050
raise typer.Exit(1)
5151

52+
config.use_production_environment()
5253
config.set_team(team_id, team_name=team_name, team_role=team_role)
53-
config.update_current_environment_file()
5454
console.print(f"[green]Switched to team '{team_name}'.[/green]")
5555

5656

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@ def whoami() -> None:
3232
# Update config
3333
config = Config()
3434
if user_id:
35+
config.use_production_environment()
3536
config.set_user_id(user_id)
36-
config.update_current_environment_file()
3737

3838
# Display account info table
3939
table = Table(title="Account")

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

Lines changed: 8 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,14 @@ def set_current_environment(self, value: str) -> None:
214214
self.config["current_environment"] = value
215215
self._save_config(self.config)
216216

217+
def use_production_environment(self) -> None:
218+
"""Use the built-in production endpoint configuration."""
219+
self.config["base_url"] = self.DEFAULT_BASE_URL
220+
self.config["frontend_url"] = self.DEFAULT_FRONTEND_URL
221+
self.config["inference_url"] = self.DEFAULT_INFERENCE_URL
222+
self.config["current_environment"] = "production"
223+
self._save_config(self.config)
224+
217225
def _sanitize_environment_name(self, name: str) -> str:
218226
"""Sanitize environment name to prevent path traversal"""
219227
# Only allow alphanumeric characters, hyphens, and underscores
@@ -354,29 +362,6 @@ def load_environment(self, name: str, persist: bool = True) -> bool:
354362
raise
355363
return False
356364

357-
def update_current_environment_file(self) -> None:
358-
"""Update the current environment's saved file with current config"""
359-
if self.current_environment != "production":
360-
# Only update custom environments, not the built-in production
361-
try:
362-
sanitized_name = self._sanitize_environment_name(self.current_environment)
363-
env_file = self.environments_dir / f"{sanitized_name}.json"
364-
if env_file.exists():
365-
env_config = {
366-
"api_key": self.api_key,
367-
"team_id": self.team_id,
368-
"team_name": None if self.team_id_from_env else self.team_name,
369-
"team_role": None if self.team_id_from_env else self.team_role,
370-
"user_id": self.user_id,
371-
"base_url": self.base_url,
372-
"frontend_url": self.frontend_url,
373-
"inference_url": self.inference_url,
374-
}
375-
env_file.write_text(json.dumps(env_config, indent=2))
376-
except ValueError:
377-
# Skip updating if environment name is invalid
378-
pass
379-
380365
def list_environments(self) -> list[str]:
381366
"""List all saved environment names"""
382367
environments = ["production"] # Built-in environment

packages/prime/tests/test_config_delete.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,32 @@
1111
"COLUMNS": "200",
1212
"LINES": "50",
1313
"PRIME_DISABLE_VERSION_CHECK": "1",
14+
"PRIME_API_KEY": "",
15+
"PRIME_TEAM_ID": "",
16+
"PRIME_USER_ID": "",
17+
"PRIME_API_BASE_URL": "",
18+
"PRIME_BASE_URL": "",
19+
"PRIME_FRONTEND_URL": "",
20+
"PRIME_INFERENCE_URL": "",
21+
"PRIME_CONTEXT": "",
1422
}
1523

1624

1725
@pytest.fixture
1826
def temp_home(tmp_path: Any, monkeypatch: pytest.MonkeyPatch) -> Path:
1927
monkeypatch.setenv("HOME", str(tmp_path))
2028
monkeypatch.setattr("prime_cli.main.check_for_update", lambda: (False, None))
29+
for name in (
30+
"PRIME_API_KEY",
31+
"PRIME_TEAM_ID",
32+
"PRIME_USER_ID",
33+
"PRIME_API_BASE_URL",
34+
"PRIME_BASE_URL",
35+
"PRIME_FRONTEND_URL",
36+
"PRIME_INFERENCE_URL",
37+
"PRIME_CONTEXT",
38+
):
39+
monkeypatch.delenv(name, raising=False)
2140
return tmp_path
2241

2342

0 commit comments

Comments
 (0)