From eddc800ab8d78083da4d5f38648aa5b37cc2021f Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Tue, 21 Jul 2026 14:55:26 -0700 Subject: [PATCH 1/8] feat(cli): connect setup to remote platforms Let collaborative users configure the local CLI against an existing deployment while preserving the default local setup path. Reuse the established OIDC flow and persist the selected deployment in the active context. AIRCORE-934 Signed-off-by: Tyler Bray --- docs/cli/configuration.mdx | 2 + docs/get-started/setup.mdx | 23 +- .../nemo_platform_ext/cli/commands/auth.py | 331 +++++++++--------- .../nemo_platform_ext/cli/commands/setup.py | 175 +++++++-- .../tests/cli/commands/test_setup.py | 242 +++++++++++++ .../tests/cli/commands/test_setup_cli.py | 110 ++++++ .../src/nemo_platform/cli/commands/auth.py | 331 +++++++++--------- .../src/nemo_platform/cli/commands/setup.py | 175 +++++++-- .../cli/commands/test_setup.py | 242 +++++++++++++ .../cli/commands/test_setup_cli.py | 110 ++++++ 10 files changed, 1366 insertions(+), 375 deletions(-) create mode 100644 packages/nemo_platform_ext/tests/cli/commands/test_setup_cli.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup_cli.py diff --git a/docs/cli/configuration.mdx b/docs/cli/configuration.mdx index eb403c55dc..957b675793 100644 --- a/docs/cli/configuration.mdx +++ b/docs/cli/configuration.mdx @@ -13,6 +13,8 @@ The quickest way to connect to an existing deployment is: nemo auth login --base-url https://nmp.example.com ``` +During interactive onboarding, `nemo setup` also offers to connect to a remote deployment when the currently configured platform is unreachable. It verifies and saves the new URL in the active context, then runs the same authentication flow. + To configure a named context: ```bash diff --git a/docs/get-started/setup.mdx b/docs/get-started/setup.mdx index 995d6bdf03..569ab3a63a 100644 --- a/docs/get-started/setup.mdx +++ b/docs/get-started/setup.mdx @@ -74,7 +74,7 @@ nemo setup The wizard walks through each stage: -1. **Start services** — launches the platform locally +1. **Connect to the platform** — uses a running platform, starts services locally, or connects to a remote instance 2. **Choose a provider** — select your model provider and enter your API key 3. **Register the provider** — connects the provider to the platform 4. **Discover models** — finds available models from your provider @@ -84,6 +84,21 @@ The wizard walks through each stage: Each stage is idempotent — you can re-run `nemo setup` at any time to add more providers or update your configuration. +### Connect to an existing deployment + +When the configured platform is not reachable, the setup wizard offers three choices: + +```text +Platform not reachable at http://localhost:8080. Start local services? + 1. Yes, start services now + 2. No, I want to connect to a remote Platform instance + 3. No, I'll start them myself +``` + +Choose the remote option and enter the deployment's base URL. Setup verifies the URL, saves it to the active CLI context, and starts the existing OIDC login flow when the deployment has authentication enabled. It then continues with provider, model, skills, and demo-agent setup. + +This connects the local CLI to an existing deployment; it does not install the platform on a remote host. + ### Auth and OIDC source installs `nemo setup` and `nemo services run` do not create authorization role bindings by themselves. If you enable auth in source config with `auth.enabled: true` and set `auth.admin_email`, run the platform seed task so the bootstrap IAM bindings exist. @@ -183,6 +198,12 @@ nemo setup --auto --start-services --install-skills --deploy-agent export OPENAI_API_KEY=sk-... nemo setup --auto +# Connect to an authenticated remote platform without prompts +export NMP_BASE_URL=https://nmp.example.com +export NMP_ACCESS_TOKEN=... +export OPENAI_API_KEY=sk-... +nemo setup --auto --no-start-services + # Skip skills and agent nemo setup --no-install-skills --no-deploy-agent ``` diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py index 235f2e520c..93ed4f13cb 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/auth.py @@ -11,11 +11,13 @@ import asyncio import logging +import os import time from typing import Annotated, cast import httpx import typer +from rich.console import Console from nemo_platform_ext.auth.helpers import ( AuthError, @@ -31,6 +33,7 @@ from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors from nemo_platform_ext.cli.core.help_formatter import create_typer_app +from nemo_platform_ext.config.config import Config from nemo_platform_ext.config.models import ConfigParams, Context app = create_typer_app( @@ -156,6 +159,168 @@ def auth_callback(ctx: typer.Context) -> None: typer.echo(ctx.get_help()) +def _login_with_oidc( + cli_context: CLIContext, + *, + no_browser: bool = False, + scope: str | None = None, + username: str | None = None, + password: str | None = None, + selected_context: str | None = None, +) -> bool: + """Authenticate the selected context using the cluster's OIDC configuration. + + Returns ``False`` when cluster authentication is disabled. + """ + from nemo_platform_ext.auth.device_flow import ( + DeviceFlowError, + authenticate_with_device_flow, + authenticate_with_password_grant, + ) + + console = Console() + context = cli_context.get_sdk_context() + base_url = str(context.cluster.base_url).rstrip("/") + + console.print(f"\nDiscovering auth configuration from {base_url}...") + + try: + oidc_config = discover_nmp_config(base_url) + except httpx.HTTPError as exc: + raise AuthError(f"Failed to discover auth configuration: {exc}") from exc + + if not oidc_config.auth_enabled: + console.print("[yellow]Authentication is not enabled on this cluster.[/]") + console.print("You can use the API without authentication.") + return False + + if not oidc_config.token_endpoint: + raise AuthError( + "This cluster does not have OIDC token endpoint configured.\n" + "Use OIDC configuration for device/password login, or for local testing use:\n" + "nemo auth login --unsigned-token --email " + ) + + login_username = username or os.environ.get("NMP_OIDC_USERNAME") + login_password = password or os.environ.get("NMP_OIDC_PASSWORD") + use_password_grant = bool(login_username and login_password) + + if use_password_grant: + if not oidc_config.client_id: + raise AuthError("OIDC client_id is required for password grant.") + elif not oidc_config.device_authorization_endpoint: + raise AuthError( + "This cluster does not support device flow authentication.\n" + "For non-interactive login use: nemo auth login --username --password \n" + "Or set NMP_OIDC_USERNAME and NMP_OIDC_PASSWORD (e.g. in CI)." + ) + + console.print(f"[green]Found OIDC configuration[/] (issuer: {oidc_config.issuer})") + + raw_defaults = oidc_config.default_scopes + default_baseline = " ".join(item for item in raw_defaults.split() if ":" not in item) + if scope: + seen: set[str] = set() + parts: list[str] = [] + for item in default_baseline.split(): + if item not in seen: + seen.add(item) + parts.append(item) + for item in scope.split(): + if item not in seen: + seen.add(item) + parts.append(item) + requested_scopes = " ".join(parts) + else: + requested_scopes = raw_defaults + + scope_prefix = normalize_scope_prefix(oidc_config.scope_prefix) + effective_scope = build_effective_scope(requested_scopes, oidc_config.scope_prefix) + + console.print("\n[bold]Requesting scopes:[/]") + for requested_scope in requested_scopes.split(): + if scope_prefix and (":" in requested_scope or requested_scope.endswith(".default")): + console.print(f" [cyan]{requested_scope}[/] [dim]({scope_prefix}{requested_scope})[/]") + else: + console.print(f" [cyan]{requested_scope}[/]") + console.print() + + if use_password_grant: + if login_username is None or login_password is None: + raise AuthError("Username and password are required for password grant.") + client_id = cast(str, oidc_config.client_id) + try: + token_response = authenticate_with_password_grant( + token_endpoint=oidc_config.token_endpoint, + client_id=client_id, + username=login_username, + password=login_password, + scope=effective_scope, + ) + except DeviceFlowError as exc: + raise AuthError(f"Authentication failed: {exc}") from exc + else: + if oidc_config.device_authorization_endpoint is None: + raise AuthError("This cluster does not support device flow authentication.") + client_id = cast(str, oidc_config.client_id) + try: + token_response = asyncio.run( + authenticate_with_device_flow( + device_authorization_endpoint=oidc_config.device_authorization_endpoint, + token_endpoint=oidc_config.token_endpoint, + client_id=client_id, + scope=effective_scope, + open_browser=not no_browser, + ) + ) + except DeviceFlowError as exc: + raise AuthError(f"Authentication failed: {exc}") from exc + + token = token_response.token_for_nmp + claims = decode_jwt_claims(token) + user_email = claims.get("upn") or claims.get("email") or claims.get("preferred_username") + raw_granted_scopes = claims.get("scp") or claims.get("scope") + granted_scopes: list[str] = [] + if isinstance(raw_granted_scopes, str): + granted_scopes = raw_granted_scopes.split() + elif isinstance(raw_granted_scopes, list): + granted_scopes = [item for item in raw_granted_scopes if isinstance(item, str)] + + validate_requested_scopes_granted(effective_scope, granted_scopes, scope_prefix) + + config_params: ConfigParams = {"access_token": token} + if token_response.refresh_token: + config_params["refresh_token"] = token_response.refresh_token + if selected_context is not None: + config_params["current_context"] = context.context_name + Config.write(config_params, context_name=context.context_name) + + console.print("\n[bold green]Authentication successful![/]") + if user_email: + console.print(f" Logged in as: [cyan]{user_email}[/]") + + if granted_scopes: + display_scopes = [ + item[len(scope_prefix) :] if scope_prefix and item.startswith(scope_prefix) else item + for item in granted_scopes + ] + console.print(f" Granted scopes: [cyan]{' '.join(display_scopes)}[/]") + + if token_response.refresh_token: + console.print(" Refresh token: [green]saved[/] (enables automatic token renewal)") + else: + console.print(" Refresh token: [yellow]not available[/] (add 'offline_access' scope to enable)") + + console.print("\n[bold green]Credentials saved to config file.[/]") + if runtime_token_source := _runtime_token_source_label(): + console.print( + f"[yellow]Warning:[/] {runtime_token_source} is active and will override these saved credentials. " + "Unset the runtime token override to use this login for future commands." + ) + console.print("\n[dim]Run 'nemo workspaces list' to verify your access.[/]") + return True + + @app.command("login") @handle_errors def login( @@ -274,17 +439,6 @@ def login( # Device flow, show code only nemo auth login --no-browser """ - import os - - from rich.console import Console - - from nemo_platform_ext.auth.device_flow import ( - DeviceFlowError, - authenticate_with_device_flow, - authenticate_with_password_grant, - ) - from nemo_platform_ext.config.config import Config - cli_context: CLIContext = ctx.obj selected_context = cli_context.overrides.get("current_context") @@ -376,155 +530,16 @@ def login( console.print("\n[dim]Run 'nemo auth status' to inspect the token.[/]") return - context = cli_context.get_sdk_context() - base_url = str(context.cluster.base_url).rstrip("/") - - console.print(f"\nDiscovering auth configuration from {base_url}...") - - try: - oidc_config = discover_nmp_config(base_url) - except httpx.HTTPError as exc: - raise AuthError(f"Failed to discover auth configuration: {exc}") from exc - - if not oidc_config.auth_enabled: - console.print("[yellow]Authentication is not enabled on this cluster.[/]") - console.print("You can use the API without authentication.") + if not _login_with_oidc( + cli_context, + no_browser=no_browser, + scope=scope, + username=username, + password=password, + selected_context=selected_context, + ): raise typer.Exit(0) - if not oidc_config.token_endpoint: - raise AuthError( - "This cluster does not have OIDC token endpoint configured.\n" - "Use OIDC configuration for device/password login, or for local testing use:\n" - "nemo auth login --unsigned-token --email " - ) - - login_username = username or os.environ.get("NMP_OIDC_USERNAME") - login_password = password or os.environ.get("NMP_OIDC_PASSWORD") - use_password_grant = bool(login_username and login_password) - - if use_password_grant: - if not oidc_config.client_id: - raise AuthError("OIDC client_id is required for password grant.") - else: - if not oidc_config.device_authorization_endpoint: - raise AuthError( - "This cluster does not support device flow authentication.\n" - "For non-interactive login use: nemo auth login --username --password \n" - "Or set NMP_OIDC_USERNAME and NMP_OIDC_PASSWORD (e.g. in CI)." - ) - - console.print(f"[green]Found OIDC configuration[/] (issuer: {oidc_config.issuer})") - - # Use only generic scopes from cluster defaults (exclude platform/custom scopes like platform:read) - # so platform scopes come only from --scope. Merge with --scope if provided. - raw_defaults = oidc_config.default_scopes - default_baseline = " ".join(s for s in raw_defaults.split() if ":" not in s) - if scope: - seen: set[str] = set() - parts: list[str] = [] - for s in default_baseline.split(): - if s not in seen: - seen.add(s) - parts.append(s) - for s in scope.split(): - if s not in seen: - seen.add(s) - parts.append(s) - requested_scopes = " ".join(parts) - else: - requested_scopes = raw_defaults - - scope_prefix = normalize_scope_prefix(oidc_config.scope_prefix) - effective_scope = build_effective_scope(requested_scopes, oidc_config.scope_prefix) - - # Display the scopes being requested - console.print("\n[bold]Requesting scopes:[/]") - for s in requested_scopes.split(): - if scope_prefix and (":" in s or s.endswith(".default")): - console.print(f" [cyan]{s}[/] [dim]({scope_prefix}{s})[/]") - else: - console.print(f" [cyan]{s}[/]") - console.print() - - if use_password_grant: - if login_username is None or login_password is None: - raise AuthError("Username and password are required for password grant.") - client_id = cast(str, oidc_config.client_id) - try: - token_response = authenticate_with_password_grant( - token_endpoint=oidc_config.token_endpoint, - client_id=client_id, - username=login_username, - password=login_password, - scope=effective_scope, - ) - except DeviceFlowError as exc: - raise AuthError(f"Authentication failed: {exc}") from exc - else: - if oidc_config.device_authorization_endpoint is None: - raise AuthError("This cluster does not support device flow authentication.") - client_id = cast(str, oidc_config.client_id) - device_authorization_endpoint = oidc_config.device_authorization_endpoint - try: - token_response = asyncio.run( - authenticate_with_device_flow( - device_authorization_endpoint=device_authorization_endpoint, - token_endpoint=oidc_config.token_endpoint, - client_id=client_id, - scope=effective_scope, - open_browser=not no_browser, - ) - ) - except DeviceFlowError as exc: - raise AuthError(f"Authentication failed: {exc}") from exc - - token = token_response.token_for_nmp - - claims = decode_jwt_claims(token) - user_email = claims.get("upn") or claims.get("email") or claims.get("preferred_username") - raw_granted_scopes = claims.get("scp") or claims.get("scope") - granted_scopes: list[str] = [] - if isinstance(raw_granted_scopes, str): - granted_scopes = raw_granted_scopes.split() - elif isinstance(raw_granted_scopes, list): - granted_scopes = [scope for scope in raw_granted_scopes if isinstance(scope, str)] - - validate_requested_scopes_granted(effective_scope, granted_scopes, scope_prefix) - - config_params: ConfigParams = {"access_token": token} - if token_response.refresh_token: - config_params["refresh_token"] = token_response.refresh_token - if selected_context is not None: - config_params["current_context"] = context.context_name - Config.write(config_params, context_name=context.context_name) - - console.print("\n[bold green]Authentication successful![/]") - if user_email: - console.print(f" Logged in as: [cyan]{user_email}[/]") - - if granted_scopes: - # Normalize scopes by stripping prefix for display - display_scopes = [] - for s in granted_scopes: - if scope_prefix and s.startswith(scope_prefix): - display_scopes.append(s[len(scope_prefix) :]) - else: - display_scopes.append(s) - console.print(f" Granted scopes: [cyan]{' '.join(display_scopes)}[/]") - - if token_response.refresh_token: - console.print(" Refresh token: [green]saved[/] (enables automatic token renewal)") - else: - console.print(" Refresh token: [yellow]not available[/] (add 'offline_access' scope to enable)") - - console.print("\n[bold green]Credentials saved to config file.[/]") - if runtime_token_source := _runtime_token_source_label(): - console.print( - f"[yellow]Warning:[/] {runtime_token_source} is active and will override these saved credentials. " - "Unset the runtime token override to use this login for future commands." - ) - console.print("\n[dim]Run 'nemo workspaces list' to verify your access.[/]") - @app.command("logout") @handle_errors diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index 626d56bd07..b70e744562 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -19,7 +19,7 @@ from importlib.resources import files from importlib.resources.abc import Traversable from pathlib import Path -from typing import Annotated +from typing import Annotated, Literal from urllib.parse import urlparse import httpx @@ -42,7 +42,7 @@ from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors from nemo_platform_ext.config.config import Config -from nemo_platform_ext.config.models import ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform_ext.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig from nemo_platform_ext.local.process import ( check_port_available_for_start, compute_scope, @@ -313,6 +313,65 @@ def _check_platform_reachable_with_retries( return False +def _prompt_remote_base_url() -> str: + """Prompt until the user provides a reachable remote Platform URL.""" + while True: + base_url = prompt_text( + "Enter the remote Platform base URL: ", + validator=non_empty_validator("Base URL"), + ).strip() + parsed = urlparse(base_url) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + console.print(f"{CROSS} Enter a valid HTTP or HTTPS URL.") + continue + + base_url = base_url.rstrip("/") + if _check_platform_reachable_with_retries(base_url): + return base_url + + console.print(f"{CROSS} Unable to connect to NeMo Platform at {base_url}.") + + +def _configure_remote_connection(cli_context: CLIContext, base_url: str, workspace: str) -> None: + """Persist a remote Platform URL in the active CLI context.""" + context_name = cli_context.get_sdk_context().context_name + Config.write( + {"base_url": base_url, "workspace": workspace}, + context_name=context_name, + ) + cli_context.overrides["base_url"] = base_url + cli_context.reset_sdk_context() + + +def _ensure_platform_auth(cli_context: CLIContext) -> None: + """Authenticate the active context when it lacks usable credentials.""" + from nemo_platform_ext.cli.commands.auth import _login_with_oidc, _runtime_token_source_label + + context = cli_context.get_sdk_context() + if runtime_token_source := _runtime_token_source_label(): + console.print(f"{CHECK} Using {runtime_token_source}\n") + return + + authenticated = _login_with_oidc(cli_context, selected_context=context.context_name) + if not authenticated: + Config.write( + {"access_token": None, "refresh_token": None}, + context_name=context.context_name, + ) + cli_context.reset_sdk_context() + + +def _platform_request_headers(cli_context: CLIContext) -> dict[str, str] | None: + """Return authentication headers for direct Platform HTTP requests.""" + context = cli_context.get_sdk_context() + if context.user is None: + return None + headers = context.user.get_client_config().get("default_headers") + if not isinstance(headers, dict): + return None + return {key: value for key, value in headers.items() if isinstance(key, str) and isinstance(value, str)} + + def _check_controller_health(base_url: str, timeout: float = 5.0) -> tuple[bool, str]: """Query ``/status`` and assess controller health. @@ -626,6 +685,12 @@ def _resolve_services_port(base_url: str) -> int: return parsed.port or 8080 +def _is_local_base_url(base_url: str) -> bool: + """Return whether *base_url* points at the local machine.""" + parsed = urlparse(base_url) + return parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"} + + def _start_services_background(base_url: str, data_dir: str | None = None) -> subprocess.Popen: """Launch ``nemo services run`` as a background process. @@ -709,7 +774,7 @@ def _maybe_start_services( auto: bool, start_services: bool | None, timeout: int = _SERVICE_STARTUP_TIMEOUT_SECONDS, -) -> None: +) -> Literal["ready", "connect_remote"]: """Start services if requested, restarting if already running. In interactive mode (auto=False), prompts the user if start_services is None. @@ -720,11 +785,17 @@ def _maybe_start_services( (including any newly installed plugins) is picked up. Data lives in SQLite so nothing is lost across restarts. """ + if start_services is True and not _is_local_base_url(base_url): + raise typer.BadParameter( + "--start-services requires a local Platform URL", + param_hint="--start-services", + ) + already_running = _check_platform_reachable(base_url) if already_running and start_services is not True: console.print(f"{CHECK} Platform already running at {base_url}\n") - return + return "ready" should_start = start_services if should_start is None: @@ -734,14 +805,20 @@ def _maybe_start_services( console.print(" [cyan]nemo setup --auto --start-services[/cyan]") console.print(" [cyan]nemo services run[/cyan]") raise typer.Exit(1) - should_start = ( - prompt_choice( - message=f"Platform not reachable at {base_url}. Start local services?", - options=[("yes", "Yes, start services now"), ("no", "No, I'll start them myself")], - default="yes", - ) - == "yes" + if not _is_local_base_url(base_url): + return "connect_remote" + action = prompt_choice( + message=f"Platform not reachable at {base_url}. Start local services?", + options=[ + ("yes", "Yes, start services now"), + ("remote", "No, I want to connect to a remote Platform instance"), + ("manual", "No, I'll start them myself"), + ], + default="yes", ) + if action == "remote": + return "connect_remote" + should_start = action == "yes" if not should_start: console.print(f"{CROSS} Cannot reach platform at {base_url}") @@ -789,6 +866,7 @@ def _maybe_start_services( raise typer.Exit(1) console.print(f"{CHECK} Platform running at {base_url} (pid {proc.pid})\n") + return "ready" # --------------------------------------------------------------------------- @@ -1142,11 +1220,12 @@ def _agent_config_path() -> Traversable | None: return None -def _agent_exists(base_url: str, workspace: str) -> bool: +def _agent_exists(base_url: str, workspace: str, headers: dict[str, str] | None = None) -> bool: """Return True if the demo agent already exists on the platform.""" try: resp = httpx.get( f"{base_url.rstrip('/')}/apis/agents/v2/workspaces/{workspace}/agents/{_DEMO_AGENT_NAME}", + headers=headers, timeout=10.0, ) return resp.status_code == 200 @@ -1154,11 +1233,12 @@ def _agent_exists(base_url: str, workspace: str) -> bool: return False -def _agents_api_ready(base_url: str, workspace: str) -> bool: +def _agents_api_ready(base_url: str, workspace: str, headers: dict[str, str] | None = None) -> bool: """Return True if the agents API is responding.""" try: resp = httpx.get( f"{base_url.rstrip('/')}/apis/agents/v2/workspaces/{workspace}/agents", + headers=headers, timeout=3.0, ) return resp.status_code == 200 @@ -1166,19 +1246,26 @@ def _agents_api_ready(base_url: str, workspace: str) -> bool: return False -def _deploy_demo_agent(base_url: str, workspace: str, config_path: Traversable, default_model: str) -> bool: +def _deploy_demo_agent( + base_url: str, + workspace: str, + config_path: Traversable, + default_model: str, + headers: dict[str, str] | None = None, +) -> bool: """Create and deploy the demo calculator agent. Returns True on success.""" # Optional plugin: import here so ``nemo setup`` works without nemo-agents installed. from nemo_agents_plugin.utils import expand_env_vars api_base = base_url.rstrip("/") - if not _agent_exists(base_url, workspace): + if not _agent_exists(base_url, workspace, headers=headers): config_dict = _yaml.safe_load(config_path.read_text(encoding="utf-8")) config_dict = expand_env_vars(config_dict, vars_dict={"NEMO_DEFAULT_MODEL": default_model}) payload = {"name": _DEMO_AGENT_NAME, "description": "Demo calculator agent", "config": config_dict} resp = httpx.post( f"{api_base}/apis/agents/v2/workspaces/{workspace}/agents", + headers=headers, json=payload, timeout=30.0, ) @@ -1189,6 +1276,7 @@ def _deploy_demo_agent(base_url: str, workspace: str, config_path: Traversable, resp = httpx.post( f"{api_base}/apis/agents/v2/workspaces/{workspace}/deployments", + headers=headers, json={"agent": _DEMO_AGENT_NAME}, timeout=30.0, ) @@ -1212,6 +1300,7 @@ def _deploy_demo_agent(base_url: str, workspace: str, config_path: Traversable, try: dep_resp = httpx.get( f"{api_base}/apis/agents/v2/workspaces/{workspace}/deployments/{deployment_name}", + headers=headers, timeout=3.0, ) if dep_resp.status_code == 200: @@ -1235,6 +1324,7 @@ def _maybe_deploy_agent( auto: bool, deploy_agent: bool | None, default_model: str | None = None, + headers: dict[str, str] | None = None, ) -> bool: """Optionally deploy the demo calculator agent. @@ -1289,7 +1379,7 @@ def _maybe_deploy_agent( while time.monotonic() < deadline: elapsed = int(time.monotonic() - start) spinner.update(f"[bold cyan]Waiting for agents API... ({elapsed}s)") - if _agents_api_ready(base_url, workspace): + if _agents_api_ready(base_url, workspace, headers=headers): api_ready = True break _pause(_AGENT_API_READINESS_POLL_INTERVAL) @@ -1299,7 +1389,13 @@ def _maybe_deploy_agent( return False try: - return _deploy_demo_agent(base_url, workspace, config_path, default_model=default_model) + return _deploy_demo_agent( + base_url, + workspace, + config_path, + default_model=default_model, + headers=headers, + ) except Exception as exc: console.print(f" {WARN} Agent deployment failed: {exc}") return False @@ -1373,9 +1469,9 @@ def _register_provider_interactive( default_extra_headers: dict[str, str] | None = None, ) -> None: """Create or update secret + provider for idempotent re-runs.""" - secret_name = f"{provider_name}-api-key" if api_key else None - - if secret_name: + secret_name: str | None = None + if api_key: + secret_name = f"{provider_name}-api-key" if _secret_exists(client, secret_name, workspace): _update_secret(client, secret_name, api_key, workspace) console.print(f" {CHECK} Updated secret '{secret_name}'") @@ -1649,12 +1745,12 @@ def setup_command( ), ] = None, ) -> None: - """Set up NeMo Platform: start services, configure a provider, install skills. + """Set up NeMo Platform: connect or start services, configure a provider, install skills. - Walks through starting local services, selecting a provider, entering - credentials, registering the provider with the platform, picking a - default model, installing coding agent skills, and optionally deploying - a demo agent. + Uses an already-running platform, starts local services, or connects the + CLI to an existing remote deployment. Then selects and registers an + inference provider, picks a default model, installs coding agent skills, + and optionally deploys a demo agent. Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -1669,11 +1765,12 @@ def setup_command( nemo setup --auto nemo setup --auto --start-services --install-skills --deploy-agent nemo setup --auto --start-services --ready-timeout 360 + NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent """ cli_context: CLIContext = ctx.obj - base_url = cli_context.get_base_url() + base_url = cli_context.get_base_url() or DEFAULT_BASE_URL console.print("\n[bold cyan]NeMo Platform Setup[/bold cyan]\n") @@ -1684,7 +1781,15 @@ def setup_command( effective_timeout = _SERVICE_STARTUP_TIMEOUT_SECONDS if ready_timeout is None else ready_timeout if effective_timeout <= 0: raise typer.BadParameter("--ready-timeout must be greater than 0", param_hint="--ready-timeout") - _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) + try: + service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) + if service_result == "connect_remote": + base_url = _prompt_remote_base_url() + _configure_remote_connection(cli_context, base_url, workspace) + _ensure_platform_auth(cli_context) + except UserCancelled: + console.print(f"\n{WARN} Setup cancelled.") + raise typer.Exit(0) from None if not _check_platform_reachable_with_retries(base_url): console.print(f"\n{CROSS} Cannot reach platform at {base_url}") @@ -1797,7 +1902,14 @@ def _run_auto_mode( skills_scope=skills_scope, skills_from=skills_from, ) - _maybe_deploy_agent(base_url, workspace, auto=True, deploy_agent=deploy_agent, default_model=default_model) + _maybe_deploy_agent( + base_url, + workspace, + auto=True, + deploy_agent=deploy_agent, + default_model=default_model, + headers=_platform_request_headers(cli_context), + ) if _verify_platform_health(base_url): console.print(f"\n{CHECK} [green]Setup complete![/green]") @@ -1882,7 +1994,12 @@ def _run_interactive_mode( console.print("\n[bold]Step 7: Demo agent (optional)[/bold]\n") demo_deployed = _maybe_deploy_agent( - base_url, workspace, auto=False, deploy_agent=deploy_agent, default_model=default_model + base_url, + workspace, + auto=False, + deploy_agent=deploy_agent, + default_model=default_model, + headers=_platform_request_headers(cli_context), ) _print_onboarding(base_url, provider_name, default_model, demo_deployed=demo_deployed) diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index bfc66f71e9..720fb3bb79 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import httpx +import nemo_platform_ext.cli.commands.setup as setup_commands import pytest import typer from click.exceptions import Exit as ClickExit @@ -69,16 +70,21 @@ from nemo_platform_ext.cli.commands.skills import registry as skills_registry from nemo_platform_ext.cli.commands.skills.base import Scope, Skill from nemo_platform_ext.cli.commands.skills.registry import UnsupportedAgentError +from nemo_platform_ext.config.config import Config from nemo_platform_ext.config.models import ( Cluster, ConfigFile, ConfigParams, Context, ContextDefinition, + NoAuthUser, + OAuthUser, ) from nemo_platform_ext.local.process import PortConflict +from nemo_platform_ext.ui.prompts import UserCancelled from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest +from pydantic import SecretStr SETUP_MOD = "nemo_platform_ext.cli.commands.setup" @@ -545,6 +551,52 @@ def test_skips_when_running_and_explicitly_false(self): with patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True): _maybe_start_services("http://localhost:8080", auto=False, start_services=False) + def test_returns_remote_choice_without_starting_services(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="remote") as mock_prompt, + patch(f"{SETUP_MOD}._start_services_background") as mock_start, + ): + result = _maybe_start_services("http://localhost:8080", auto=False, start_services=None) + + assert result == "connect_remote" + assert mock_prompt.call_args.kwargs["options"] == [ + ("yes", "Yes, start services now"), + ("remote", "No, I want to connect to a remote Platform instance"), + ("manual", "No, I'll start them myself"), + ] + mock_start.assert_not_called() + + def test_start_myself_choice_keeps_existing_exit(self, capsys): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="manual"), + pytest.raises(ClickExit), + ): + _maybe_start_services("http://localhost:8080", auto=False, start_services=None) + + assert "Start the platform first" in capsys.readouterr().err + + def test_unreachable_remote_url_selects_remote_connection_without_local_prompt(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice") as mock_prompt, + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "connect_remote" + mock_prompt.assert_not_called() + + def test_rejects_start_services_for_remote_url(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}._start_services_background") as mock_start, + pytest.raises(typer.BadParameter, match="local Platform URL"), + ): + _maybe_start_services("https://remote.example.com", auto=False, start_services=True) + + mock_start.assert_not_called() + def test_restarts_when_running_and_start_services_true(self): reachable_calls = [True, True, False, True] @@ -622,6 +674,160 @@ def test_allows_start_when_port_free(self, maybe_start_preflight_mocks): maybe_start_preflight_mocks.assert_called_once() +class TestRemoteConnection: + def test_prompts_again_until_platform_is_reachable(self, capsys): + with ( + patch( + f"{SETUP_MOD}.prompt_text", + side_effect=["https://unreachable.example.com", "https://remote.example.com/"], + ), + patch( + f"{SETUP_MOD}._check_platform_reachable_with_retries", + side_effect=[False, True], + ), + ): + base_url = setup_commands._prompt_remote_base_url() + + assert base_url == "https://remote.example.com" + assert "Unable to connect" in capsys.readouterr().err + + def test_persists_remote_url_and_workspace_in_active_context(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + cli_context = MagicMock() + cli_context.overrides = {} + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) + + setup_commands._configure_remote_connection( + cli_context, + "https://remote.example.com", + "shared-workspace", + ) + + context = Config.load(config_path=config_path).resolve() + assert str(context.cluster.base_url) == "https://remote.example.com/" + assert context.workspace == "shared-workspace" + cli_context.reset_sdk_context.assert_called_once_with() + + def test_updates_selected_context_and_runtime_url_override(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + Config.write({"base_url": "https://default.example.com"}) + Config.write({"base_url": "https://old-dev.example.com"}, context_name="dev") + + cli_context = MagicMock() + cli_context.overrides = { + "current_context": "dev", + "base_url": "https://stale-override.example.com", + } + cli_context.get_sdk_context.return_value = Context( + context_name="dev", + cluster=Cluster(name="dev-cluster", base_url="https://stale-override.example.com"), + user=NoAuthUser(name="dev-user"), + workspace="default", + preferences={}, + ) + + setup_commands._configure_remote_connection( + cli_context, + "https://new-dev.example.com", + "shared-workspace", + ) + + config_file = Config.load(config_path=config_path).get_config_file() + default_cluster = next(cluster for cluster in config_file.clusters if cluster.name == "default-cluster") + dev_cluster = next(cluster for cluster in config_file.clusters if cluster.name == "dev-cluster") + assert str(default_cluster.base_url) == "https://default.example.com/" + assert str(dev_cluster.base_url) == "https://new-dev.example.com/" + assert cli_context.overrides["base_url"] == "https://new-dev.example.com" + + def test_authenticates_when_context_has_no_credentials(self): + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) + + with patch("nemo_platform_ext.cli.commands.auth._login_with_oidc", return_value=True) as mock_login: + setup_commands._ensure_platform_auth(cli_context) + + mock_login.assert_called_once_with(cli_context, selected_context="default") + cli_context.reset_sdk_context.assert_called_once_with() + + def test_reauthenticates_stored_context_credentials_for_new_remote(self): + cli_context = MagicMock() + context = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=OAuthUser(name="default-user", token=SecretStr("token")), + workspace="default", + preferences={}, + ) + cli_context.get_sdk_context.return_value = context + + with patch("nemo_platform_ext.cli.commands.auth._login_with_oidc", return_value=True) as mock_login: + setup_commands._ensure_platform_auth(cli_context) + + mock_login.assert_called_once_with(cli_context, selected_context="default") + cli_context.reset_sdk_context.assert_called_once_with() + + def test_reuses_runtime_access_token_override(self, monkeypatch): + monkeypatch.setenv("NMP_ACCESS_TOKEN", "runtime-token") + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=OAuthUser(name="default-user", token=SecretStr("runtime-token")), + workspace="default", + preferences={}, + ) + + with patch("nemo_platform_ext.cli.commands.auth._login_with_oidc") as mock_login: + setup_commands._ensure_platform_auth(cli_context) + + mock_login.assert_not_called() + + def test_clears_stale_credentials_when_remote_auth_is_disabled(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) + Config.write( + { + "base_url": "https://remote.example.com", + "access_token": "old-cluster-token", + } + ) + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Config.load(config_path=config_path).resolve() + + with patch("nemo_platform_ext.cli.commands.auth._login_with_oidc", return_value=False): + setup_commands._ensure_platform_auth(cli_context) + + assert isinstance(Config.load(config_path=config_path).resolve().user, NoAuthUser) + cli_context.reset_sdk_context.assert_called_once_with() + + def test_platform_request_headers_include_context_token(self): + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=OAuthUser(name="default-user", token=SecretStr("remote-token")), + workspace="default", + preferences={}, + ) + + assert setup_commands._platform_request_headers(cli_context) == {"Authorization": "Bearer remote-token"} + + class TestLocalDataDirHelpers: """Tests for the XDG-default data-dir helpers used by `nemo setup`.""" @@ -2297,6 +2503,42 @@ def test_proceeds_when_tty_without_auto(self): ): self._invoke(auto=False) + def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", side_effect=UserCancelled), + pytest.raises(typer.Exit) as exc_info, + ): + self._invoke(auto=False) + + assert exc_info.value.exit_code == 0 + assert "Setup cancelled" in capsys.readouterr().err + + +class TestSetupCommandRemoteFlow: + def test_remote_choice_connects_before_continuing_setup(self): + ctx = MagicMock(spec=typer.Context) + cli_context = MagicMock() + cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, + patch(f"{SETUP_MOD}._ensure_platform_auth") as mock_auth, + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx) + + mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "default") + mock_auth.assert_called_once_with(cli_context) + assert mock_run.call_args.args[3] == "https://remote.example.com" + # --------------------------------------------------------------------------- # Controller health check diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup_cli.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup_cli.py new file mode 100644 index 0000000000..9ffb653682 --- /dev/null +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup_cli.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI integration tests for ``nemo setup`` connection selection.""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform_ext.cli.app import app +from nemo_platform_ext.cli.core.context import CLIContext +from nemo_platform_ext.config.config import Config +from typer.testing import CliRunner + +SETUP_MOD = "nemo_platform_ext.cli.commands.setup" + + +def test_remote_choice_retries_and_persists_connection( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.touch() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + + client = MagicMock() + client.workspaces.retrieve.return_value = MagicMock() + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="remote"), + patch( + f"{SETUP_MOD}.prompt_text", + side_effect=["https://unreachable.example.com", "https://remote.example.com/"], + ), + patch( + f"{SETUP_MOD}._check_platform_reachable_with_retries", + side_effect=[False, True, True], + ), + patch(f"{SETUP_MOD}._ensure_platform_auth") as ensure_auth, + patch(f"{SETUP_MOD}._start_services_background") as start_services, + patch.object(CLIContext, "get_client", return_value=client), + patch(f"{SETUP_MOD}._run_interactive_mode") as run_interactive, + ): + result = CliRunner().invoke(app, ["setup"]) + + assert result.exit_code == 0, result.output + assert "Unable to connect to NeMo Platform at https://unreachable.example.com" in result.output + context = Config.load(config_path=config_path).resolve() + assert str(context.cluster.base_url) == "https://remote.example.com/" + start_services.assert_not_called() + ensure_auth.assert_called_once() + assert run_interactive.call_args.args[3] == "https://remote.example.com" + + +def test_local_choice_starts_services_and_keeps_local_connection( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.touch() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + + client = MagicMock() + client.workspaces.retrieve.return_value = MagicMock() + process = MagicMock(pid=1234) + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="yes"), + patch(f"{SETUP_MOD}._prompt_data_dir", return_value="/tmp/nemo-demo"), + patch(f"{SETUP_MOD}.importlib.util.find_spec", return_value=MagicMock()), + patch(f"{SETUP_MOD}._ensure_port_available_for_start"), + patch(f"{SETUP_MOD}._start_services_background", return_value=process) as start_services, + patch(f"{SETUP_MOD}._wait_for_platform", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch.object(CLIContext, "get_client", return_value=client), + patch(f"{SETUP_MOD}._run_interactive_mode") as run_interactive, + ): + result = CliRunner().invoke(app, ["setup"]) + + assert result.exit_code == 0, result.output + context = Config.load(config_path=config_path).resolve() + assert str(context.cluster.base_url) == "http://localhost:8080/" + start_services.assert_called_once() + assert run_interactive.call_args.args[3].rstrip("/") == "http://localhost:8080" + + +def test_start_myself_exits_without_starting_or_mutating_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.touch() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="manual"), + patch(f"{SETUP_MOD}._start_services_background") as start_services, + ): + result = CliRunner().invoke(app, ["setup"]) + + assert result.exit_code == 1 + assert "Start the platform first" in result.output + assert config_path.read_text() == "" + start_services.assert_not_called() diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py index 4bf510bb51..b8b04cff48 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/auth.py @@ -11,11 +11,13 @@ import asyncio import logging +import os import time from typing import Annotated, cast import httpx import typer +from rich.console import Console from nemo_platform.auth.helpers import ( AuthError, @@ -31,6 +33,7 @@ from nemo_platform.cli.core.context import CLIContext from nemo_platform.cli.core.errors import handle_errors from nemo_platform.cli.core.help_formatter import create_typer_app +from nemo_platform.config.config import Config from nemo_platform.config.models import ConfigParams, Context app = create_typer_app( @@ -156,6 +159,168 @@ def auth_callback(ctx: typer.Context) -> None: typer.echo(ctx.get_help()) +def _login_with_oidc( + cli_context: CLIContext, + *, + no_browser: bool = False, + scope: str | None = None, + username: str | None = None, + password: str | None = None, + selected_context: str | None = None, +) -> bool: + """Authenticate the selected context using the cluster's OIDC configuration. + + Returns ``False`` when cluster authentication is disabled. + """ + from nemo_platform.auth.device_flow import ( + DeviceFlowError, + authenticate_with_device_flow, + authenticate_with_password_grant, + ) + + console = Console() + context = cli_context.get_sdk_context() + base_url = str(context.cluster.base_url).rstrip("/") + + console.print(f"\nDiscovering auth configuration from {base_url}...") + + try: + oidc_config = discover_nmp_config(base_url) + except httpx.HTTPError as exc: + raise AuthError(f"Failed to discover auth configuration: {exc}") from exc + + if not oidc_config.auth_enabled: + console.print("[yellow]Authentication is not enabled on this cluster.[/]") + console.print("You can use the API without authentication.") + return False + + if not oidc_config.token_endpoint: + raise AuthError( + "This cluster does not have OIDC token endpoint configured.\n" + "Use OIDC configuration for device/password login, or for local testing use:\n" + "nemo auth login --unsigned-token --email " + ) + + login_username = username or os.environ.get("NMP_OIDC_USERNAME") + login_password = password or os.environ.get("NMP_OIDC_PASSWORD") + use_password_grant = bool(login_username and login_password) + + if use_password_grant: + if not oidc_config.client_id: + raise AuthError("OIDC client_id is required for password grant.") + elif not oidc_config.device_authorization_endpoint: + raise AuthError( + "This cluster does not support device flow authentication.\n" + "For non-interactive login use: nemo auth login --username --password \n" + "Or set NMP_OIDC_USERNAME and NMP_OIDC_PASSWORD (e.g. in CI)." + ) + + console.print(f"[green]Found OIDC configuration[/] (issuer: {oidc_config.issuer})") + + raw_defaults = oidc_config.default_scopes + default_baseline = " ".join(item for item in raw_defaults.split() if ":" not in item) + if scope: + seen: set[str] = set() + parts: list[str] = [] + for item in default_baseline.split(): + if item not in seen: + seen.add(item) + parts.append(item) + for item in scope.split(): + if item not in seen: + seen.add(item) + parts.append(item) + requested_scopes = " ".join(parts) + else: + requested_scopes = raw_defaults + + scope_prefix = normalize_scope_prefix(oidc_config.scope_prefix) + effective_scope = build_effective_scope(requested_scopes, oidc_config.scope_prefix) + + console.print("\n[bold]Requesting scopes:[/]") + for requested_scope in requested_scopes.split(): + if scope_prefix and (":" in requested_scope or requested_scope.endswith(".default")): + console.print(f" [cyan]{requested_scope}[/] [dim]({scope_prefix}{requested_scope})[/]") + else: + console.print(f" [cyan]{requested_scope}[/]") + console.print() + + if use_password_grant: + if login_username is None or login_password is None: + raise AuthError("Username and password are required for password grant.") + client_id = cast(str, oidc_config.client_id) + try: + token_response = authenticate_with_password_grant( + token_endpoint=oidc_config.token_endpoint, + client_id=client_id, + username=login_username, + password=login_password, + scope=effective_scope, + ) + except DeviceFlowError as exc: + raise AuthError(f"Authentication failed: {exc}") from exc + else: + if oidc_config.device_authorization_endpoint is None: + raise AuthError("This cluster does not support device flow authentication.") + client_id = cast(str, oidc_config.client_id) + try: + token_response = asyncio.run( + authenticate_with_device_flow( + device_authorization_endpoint=oidc_config.device_authorization_endpoint, + token_endpoint=oidc_config.token_endpoint, + client_id=client_id, + scope=effective_scope, + open_browser=not no_browser, + ) + ) + except DeviceFlowError as exc: + raise AuthError(f"Authentication failed: {exc}") from exc + + token = token_response.token_for_nmp + claims = decode_jwt_claims(token) + user_email = claims.get("upn") or claims.get("email") or claims.get("preferred_username") + raw_granted_scopes = claims.get("scp") or claims.get("scope") + granted_scopes: list[str] = [] + if isinstance(raw_granted_scopes, str): + granted_scopes = raw_granted_scopes.split() + elif isinstance(raw_granted_scopes, list): + granted_scopes = [item for item in raw_granted_scopes if isinstance(item, str)] + + validate_requested_scopes_granted(effective_scope, granted_scopes, scope_prefix) + + config_params: ConfigParams = {"access_token": token} + if token_response.refresh_token: + config_params["refresh_token"] = token_response.refresh_token + if selected_context is not None: + config_params["current_context"] = context.context_name + Config.write(config_params, context_name=context.context_name) + + console.print("\n[bold green]Authentication successful![/]") + if user_email: + console.print(f" Logged in as: [cyan]{user_email}[/]") + + if granted_scopes: + display_scopes = [ + item[len(scope_prefix) :] if scope_prefix and item.startswith(scope_prefix) else item + for item in granted_scopes + ] + console.print(f" Granted scopes: [cyan]{' '.join(display_scopes)}[/]") + + if token_response.refresh_token: + console.print(" Refresh token: [green]saved[/] (enables automatic token renewal)") + else: + console.print(" Refresh token: [yellow]not available[/] (add 'offline_access' scope to enable)") + + console.print("\n[bold green]Credentials saved to config file.[/]") + if runtime_token_source := _runtime_token_source_label(): + console.print( + f"[yellow]Warning:[/] {runtime_token_source} is active and will override these saved credentials. " + "Unset the runtime token override to use this login for future commands." + ) + console.print("\n[dim]Run 'nemo workspaces list' to verify your access.[/]") + return True + + @app.command("login") @handle_errors def login( @@ -274,17 +439,6 @@ def login( # Device flow, show code only nemo auth login --no-browser """ - import os - - from rich.console import Console - - from nemo_platform.auth.device_flow import ( - DeviceFlowError, - authenticate_with_device_flow, - authenticate_with_password_grant, - ) - from nemo_platform.config.config import Config - cli_context: CLIContext = ctx.obj selected_context = cli_context.overrides.get("current_context") @@ -376,155 +530,16 @@ def login( console.print("\n[dim]Run 'nemo auth status' to inspect the token.[/]") return - context = cli_context.get_sdk_context() - base_url = str(context.cluster.base_url).rstrip("/") - - console.print(f"\nDiscovering auth configuration from {base_url}...") - - try: - oidc_config = discover_nmp_config(base_url) - except httpx.HTTPError as exc: - raise AuthError(f"Failed to discover auth configuration: {exc}") from exc - - if not oidc_config.auth_enabled: - console.print("[yellow]Authentication is not enabled on this cluster.[/]") - console.print("You can use the API without authentication.") + if not _login_with_oidc( + cli_context, + no_browser=no_browser, + scope=scope, + username=username, + password=password, + selected_context=selected_context, + ): raise typer.Exit(0) - if not oidc_config.token_endpoint: - raise AuthError( - "This cluster does not have OIDC token endpoint configured.\n" - "Use OIDC configuration for device/password login, or for local testing use:\n" - "nemo auth login --unsigned-token --email " - ) - - login_username = username or os.environ.get("NMP_OIDC_USERNAME") - login_password = password or os.environ.get("NMP_OIDC_PASSWORD") - use_password_grant = bool(login_username and login_password) - - if use_password_grant: - if not oidc_config.client_id: - raise AuthError("OIDC client_id is required for password grant.") - else: - if not oidc_config.device_authorization_endpoint: - raise AuthError( - "This cluster does not support device flow authentication.\n" - "For non-interactive login use: nemo auth login --username --password \n" - "Or set NMP_OIDC_USERNAME and NMP_OIDC_PASSWORD (e.g. in CI)." - ) - - console.print(f"[green]Found OIDC configuration[/] (issuer: {oidc_config.issuer})") - - # Use only generic scopes from cluster defaults (exclude platform/custom scopes like platform:read) - # so platform scopes come only from --scope. Merge with --scope if provided. - raw_defaults = oidc_config.default_scopes - default_baseline = " ".join(s for s in raw_defaults.split() if ":" not in s) - if scope: - seen: set[str] = set() - parts: list[str] = [] - for s in default_baseline.split(): - if s not in seen: - seen.add(s) - parts.append(s) - for s in scope.split(): - if s not in seen: - seen.add(s) - parts.append(s) - requested_scopes = " ".join(parts) - else: - requested_scopes = raw_defaults - - scope_prefix = normalize_scope_prefix(oidc_config.scope_prefix) - effective_scope = build_effective_scope(requested_scopes, oidc_config.scope_prefix) - - # Display the scopes being requested - console.print("\n[bold]Requesting scopes:[/]") - for s in requested_scopes.split(): - if scope_prefix and (":" in s or s.endswith(".default")): - console.print(f" [cyan]{s}[/] [dim]({scope_prefix}{s})[/]") - else: - console.print(f" [cyan]{s}[/]") - console.print() - - if use_password_grant: - if login_username is None or login_password is None: - raise AuthError("Username and password are required for password grant.") - client_id = cast(str, oidc_config.client_id) - try: - token_response = authenticate_with_password_grant( - token_endpoint=oidc_config.token_endpoint, - client_id=client_id, - username=login_username, - password=login_password, - scope=effective_scope, - ) - except DeviceFlowError as exc: - raise AuthError(f"Authentication failed: {exc}") from exc - else: - if oidc_config.device_authorization_endpoint is None: - raise AuthError("This cluster does not support device flow authentication.") - client_id = cast(str, oidc_config.client_id) - device_authorization_endpoint = oidc_config.device_authorization_endpoint - try: - token_response = asyncio.run( - authenticate_with_device_flow( - device_authorization_endpoint=device_authorization_endpoint, - token_endpoint=oidc_config.token_endpoint, - client_id=client_id, - scope=effective_scope, - open_browser=not no_browser, - ) - ) - except DeviceFlowError as exc: - raise AuthError(f"Authentication failed: {exc}") from exc - - token = token_response.token_for_nmp - - claims = decode_jwt_claims(token) - user_email = claims.get("upn") or claims.get("email") or claims.get("preferred_username") - raw_granted_scopes = claims.get("scp") or claims.get("scope") - granted_scopes: list[str] = [] - if isinstance(raw_granted_scopes, str): - granted_scopes = raw_granted_scopes.split() - elif isinstance(raw_granted_scopes, list): - granted_scopes = [scope for scope in raw_granted_scopes if isinstance(scope, str)] - - validate_requested_scopes_granted(effective_scope, granted_scopes, scope_prefix) - - config_params: ConfigParams = {"access_token": token} - if token_response.refresh_token: - config_params["refresh_token"] = token_response.refresh_token - if selected_context is not None: - config_params["current_context"] = context.context_name - Config.write(config_params, context_name=context.context_name) - - console.print("\n[bold green]Authentication successful![/]") - if user_email: - console.print(f" Logged in as: [cyan]{user_email}[/]") - - if granted_scopes: - # Normalize scopes by stripping prefix for display - display_scopes = [] - for s in granted_scopes: - if scope_prefix and s.startswith(scope_prefix): - display_scopes.append(s[len(scope_prefix) :]) - else: - display_scopes.append(s) - console.print(f" Granted scopes: [cyan]{' '.join(display_scopes)}[/]") - - if token_response.refresh_token: - console.print(" Refresh token: [green]saved[/] (enables automatic token renewal)") - else: - console.print(" Refresh token: [yellow]not available[/] (add 'offline_access' scope to enable)") - - console.print("\n[bold green]Credentials saved to config file.[/]") - if runtime_token_source := _runtime_token_source_label(): - console.print( - f"[yellow]Warning:[/] {runtime_token_source} is active and will override these saved credentials. " - "Unset the runtime token override to use this login for future commands." - ) - console.print("\n[dim]Run 'nemo workspaces list' to verify your access.[/]") - @app.command("logout") @handle_errors diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 78849e4181..52257ed768 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -19,7 +19,7 @@ from importlib.resources import files from importlib.resources.abc import Traversable from pathlib import Path -from typing import Annotated +from typing import Annotated, Literal from urllib.parse import urlparse import httpx @@ -42,7 +42,7 @@ from nemo_platform.cli.core.context import CLIContext from nemo_platform.cli.core.errors import handle_errors from nemo_platform.config.config import Config -from nemo_platform.config.models import ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig from nemo_platform.local.process import ( check_port_available_for_start, compute_scope, @@ -313,6 +313,65 @@ def _check_platform_reachable_with_retries( return False +def _prompt_remote_base_url() -> str: + """Prompt until the user provides a reachable remote Platform URL.""" + while True: + base_url = prompt_text( + "Enter the remote Platform base URL: ", + validator=non_empty_validator("Base URL"), + ).strip() + parsed = urlparse(base_url) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + console.print(f"{CROSS} Enter a valid HTTP or HTTPS URL.") + continue + + base_url = base_url.rstrip("/") + if _check_platform_reachable_with_retries(base_url): + return base_url + + console.print(f"{CROSS} Unable to connect to NeMo Platform at {base_url}.") + + +def _configure_remote_connection(cli_context: CLIContext, base_url: str, workspace: str) -> None: + """Persist a remote Platform URL in the active CLI context.""" + context_name = cli_context.get_sdk_context().context_name + Config.write( + {"base_url": base_url, "workspace": workspace}, + context_name=context_name, + ) + cli_context.overrides["base_url"] = base_url + cli_context.reset_sdk_context() + + +def _ensure_platform_auth(cli_context: CLIContext) -> None: + """Authenticate the active context when it lacks usable credentials.""" + from nemo_platform.cli.commands.auth import _login_with_oidc, _runtime_token_source_label + + context = cli_context.get_sdk_context() + if runtime_token_source := _runtime_token_source_label(): + console.print(f"{CHECK} Using {runtime_token_source}\n") + return + + authenticated = _login_with_oidc(cli_context, selected_context=context.context_name) + if not authenticated: + Config.write( + {"access_token": None, "refresh_token": None}, + context_name=context.context_name, + ) + cli_context.reset_sdk_context() + + +def _platform_request_headers(cli_context: CLIContext) -> dict[str, str] | None: + """Return authentication headers for direct Platform HTTP requests.""" + context = cli_context.get_sdk_context() + if context.user is None: + return None + headers = context.user.get_client_config().get("default_headers") + if not isinstance(headers, dict): + return None + return {key: value for key, value in headers.items() if isinstance(key, str) and isinstance(value, str)} + + def _check_controller_health(base_url: str, timeout: float = 5.0) -> tuple[bool, str]: """Query ``/status`` and assess controller health. @@ -626,6 +685,12 @@ def _resolve_services_port(base_url: str) -> int: return parsed.port or 8080 +def _is_local_base_url(base_url: str) -> bool: + """Return whether *base_url* points at the local machine.""" + parsed = urlparse(base_url) + return parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"} + + def _start_services_background(base_url: str, data_dir: str | None = None) -> subprocess.Popen: """Launch ``nemo services run`` as a background process. @@ -709,7 +774,7 @@ def _maybe_start_services( auto: bool, start_services: bool | None, timeout: int = _SERVICE_STARTUP_TIMEOUT_SECONDS, -) -> None: +) -> Literal["ready", "connect_remote"]: """Start services if requested, restarting if already running. In interactive mode (auto=False), prompts the user if start_services is None. @@ -720,11 +785,17 @@ def _maybe_start_services( (including any newly installed plugins) is picked up. Data lives in SQLite so nothing is lost across restarts. """ + if start_services is True and not _is_local_base_url(base_url): + raise typer.BadParameter( + "--start-services requires a local Platform URL", + param_hint="--start-services", + ) + already_running = _check_platform_reachable(base_url) if already_running and start_services is not True: console.print(f"{CHECK} Platform already running at {base_url}\n") - return + return "ready" should_start = start_services if should_start is None: @@ -734,14 +805,20 @@ def _maybe_start_services( console.print(" [cyan]nemo setup --auto --start-services[/cyan]") console.print(" [cyan]nemo services run[/cyan]") raise typer.Exit(1) - should_start = ( - prompt_choice( - message=f"Platform not reachable at {base_url}. Start local services?", - options=[("yes", "Yes, start services now"), ("no", "No, I'll start them myself")], - default="yes", - ) - == "yes" + if not _is_local_base_url(base_url): + return "connect_remote" + action = prompt_choice( + message=f"Platform not reachable at {base_url}. Start local services?", + options=[ + ("yes", "Yes, start services now"), + ("remote", "No, I want to connect to a remote Platform instance"), + ("manual", "No, I'll start them myself"), + ], + default="yes", ) + if action == "remote": + return "connect_remote" + should_start = action == "yes" if not should_start: console.print(f"{CROSS} Cannot reach platform at {base_url}") @@ -789,6 +866,7 @@ def _maybe_start_services( raise typer.Exit(1) console.print(f"{CHECK} Platform running at {base_url} (pid {proc.pid})\n") + return "ready" # --------------------------------------------------------------------------- @@ -1142,11 +1220,12 @@ def _agent_config_path() -> Traversable | None: return None -def _agent_exists(base_url: str, workspace: str) -> bool: +def _agent_exists(base_url: str, workspace: str, headers: dict[str, str] | None = None) -> bool: """Return True if the demo agent already exists on the platform.""" try: resp = httpx.get( f"{base_url.rstrip('/')}/apis/agents/v2/workspaces/{workspace}/agents/{_DEMO_AGENT_NAME}", + headers=headers, timeout=10.0, ) return resp.status_code == 200 @@ -1154,11 +1233,12 @@ def _agent_exists(base_url: str, workspace: str) -> bool: return False -def _agents_api_ready(base_url: str, workspace: str) -> bool: +def _agents_api_ready(base_url: str, workspace: str, headers: dict[str, str] | None = None) -> bool: """Return True if the agents API is responding.""" try: resp = httpx.get( f"{base_url.rstrip('/')}/apis/agents/v2/workspaces/{workspace}/agents", + headers=headers, timeout=3.0, ) return resp.status_code == 200 @@ -1166,19 +1246,26 @@ def _agents_api_ready(base_url: str, workspace: str) -> bool: return False -def _deploy_demo_agent(base_url: str, workspace: str, config_path: Traversable, default_model: str) -> bool: +def _deploy_demo_agent( + base_url: str, + workspace: str, + config_path: Traversable, + default_model: str, + headers: dict[str, str] | None = None, +) -> bool: """Create and deploy the demo calculator agent. Returns True on success.""" # Optional plugin: import here so ``nemo setup`` works without nemo-agents installed. from nemo_agents_plugin.utils import expand_env_vars api_base = base_url.rstrip("/") - if not _agent_exists(base_url, workspace): + if not _agent_exists(base_url, workspace, headers=headers): config_dict = _yaml.safe_load(config_path.read_text(encoding="utf-8")) config_dict = expand_env_vars(config_dict, vars_dict={"NEMO_DEFAULT_MODEL": default_model}) payload = {"name": _DEMO_AGENT_NAME, "description": "Demo calculator agent", "config": config_dict} resp = httpx.post( f"{api_base}/apis/agents/v2/workspaces/{workspace}/agents", + headers=headers, json=payload, timeout=30.0, ) @@ -1189,6 +1276,7 @@ def _deploy_demo_agent(base_url: str, workspace: str, config_path: Traversable, resp = httpx.post( f"{api_base}/apis/agents/v2/workspaces/{workspace}/deployments", + headers=headers, json={"agent": _DEMO_AGENT_NAME}, timeout=30.0, ) @@ -1212,6 +1300,7 @@ def _deploy_demo_agent(base_url: str, workspace: str, config_path: Traversable, try: dep_resp = httpx.get( f"{api_base}/apis/agents/v2/workspaces/{workspace}/deployments/{deployment_name}", + headers=headers, timeout=3.0, ) if dep_resp.status_code == 200: @@ -1235,6 +1324,7 @@ def _maybe_deploy_agent( auto: bool, deploy_agent: bool | None, default_model: str | None = None, + headers: dict[str, str] | None = None, ) -> bool: """Optionally deploy the demo calculator agent. @@ -1289,7 +1379,7 @@ def _maybe_deploy_agent( while time.monotonic() < deadline: elapsed = int(time.monotonic() - start) spinner.update(f"[bold cyan]Waiting for agents API... ({elapsed}s)") - if _agents_api_ready(base_url, workspace): + if _agents_api_ready(base_url, workspace, headers=headers): api_ready = True break _pause(_AGENT_API_READINESS_POLL_INTERVAL) @@ -1299,7 +1389,13 @@ def _maybe_deploy_agent( return False try: - return _deploy_demo_agent(base_url, workspace, config_path, default_model=default_model) + return _deploy_demo_agent( + base_url, + workspace, + config_path, + default_model=default_model, + headers=headers, + ) except Exception as exc: console.print(f" {WARN} Agent deployment failed: {exc}") return False @@ -1373,9 +1469,9 @@ def _register_provider_interactive( default_extra_headers: dict[str, str] | None = None, ) -> None: """Create or update secret + provider for idempotent re-runs.""" - secret_name = f"{provider_name}-api-key" if api_key else None - - if secret_name: + secret_name: str | None = None + if api_key: + secret_name = f"{provider_name}-api-key" if _secret_exists(client, secret_name, workspace): _update_secret(client, secret_name, api_key, workspace) console.print(f" {CHECK} Updated secret '{secret_name}'") @@ -1649,12 +1745,12 @@ def setup_command( ), ] = None, ) -> None: - """Set up NeMo Platform: start services, configure a provider, install skills. + """Set up NeMo Platform: connect or start services, configure a provider, install skills. - Walks through starting local services, selecting a provider, entering - credentials, registering the provider with the platform, picking a - default model, installing coding agent skills, and optionally deploying - a demo agent. + Uses an already-running platform, starts local services, or connects the + CLI to an existing remote deployment. Then selects and registers an + inference provider, picks a default model, installs coding agent skills, + and optionally deploys a demo agent. Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -1669,11 +1765,12 @@ def setup_command( nemo setup --auto nemo setup --auto --start-services --install-skills --deploy-agent nemo setup --auto --start-services --ready-timeout 360 + NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent """ cli_context: CLIContext = ctx.obj - base_url = cli_context.get_base_url() + base_url = cli_context.get_base_url() or DEFAULT_BASE_URL console.print("\n[bold cyan]NeMo Platform Setup[/bold cyan]\n") @@ -1684,7 +1781,15 @@ def setup_command( effective_timeout = _SERVICE_STARTUP_TIMEOUT_SECONDS if ready_timeout is None else ready_timeout if effective_timeout <= 0: raise typer.BadParameter("--ready-timeout must be greater than 0", param_hint="--ready-timeout") - _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) + try: + service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) + if service_result == "connect_remote": + base_url = _prompt_remote_base_url() + _configure_remote_connection(cli_context, base_url, workspace) + _ensure_platform_auth(cli_context) + except UserCancelled: + console.print(f"\n{WARN} Setup cancelled.") + raise typer.Exit(0) from None if not _check_platform_reachable_with_retries(base_url): console.print(f"\n{CROSS} Cannot reach platform at {base_url}") @@ -1797,7 +1902,14 @@ def _run_auto_mode( skills_scope=skills_scope, skills_from=skills_from, ) - _maybe_deploy_agent(base_url, workspace, auto=True, deploy_agent=deploy_agent, default_model=default_model) + _maybe_deploy_agent( + base_url, + workspace, + auto=True, + deploy_agent=deploy_agent, + default_model=default_model, + headers=_platform_request_headers(cli_context), + ) if _verify_platform_health(base_url): console.print(f"\n{CHECK} [green]Setup complete![/green]") @@ -1882,7 +1994,12 @@ def _run_interactive_mode( console.print("\n[bold]Step 7: Demo agent (optional)[/bold]\n") demo_deployed = _maybe_deploy_agent( - base_url, workspace, auto=False, deploy_agent=deploy_agent, default_model=default_model + base_url, + workspace, + auto=False, + deploy_agent=deploy_agent, + default_model=default_model, + headers=_platform_request_headers(cli_context), ) _print_onboarding(base_url, provider_name, default_model, demo_deployed=demo_deployed) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index 84a99ee90d..66c68ea9b6 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import httpx +import nemo_platform.cli.commands.setup as setup_commands import pytest import typer from click.exceptions import Exit as ClickExit @@ -69,16 +70,21 @@ from nemo_platform.cli.commands.skills import registry as skills_registry from nemo_platform.cli.commands.skills.base import Scope, Skill from nemo_platform.cli.commands.skills.registry import UnsupportedAgentError +from nemo_platform.config.config import Config from nemo_platform.config.models import ( Cluster, ConfigFile, ConfigParams, Context, ContextDefinition, + NoAuthUser, + OAuthUser, ) from nemo_platform.local.process import PortConflict +from nemo_platform.ui.prompts import UserCancelled from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest +from pydantic import SecretStr SETUP_MOD = "nemo_platform.cli.commands.setup" @@ -545,6 +551,52 @@ def test_skips_when_running_and_explicitly_false(self): with patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True): _maybe_start_services("http://localhost:8080", auto=False, start_services=False) + def test_returns_remote_choice_without_starting_services(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="remote") as mock_prompt, + patch(f"{SETUP_MOD}._start_services_background") as mock_start, + ): + result = _maybe_start_services("http://localhost:8080", auto=False, start_services=None) + + assert result == "connect_remote" + assert mock_prompt.call_args.kwargs["options"] == [ + ("yes", "Yes, start services now"), + ("remote", "No, I want to connect to a remote Platform instance"), + ("manual", "No, I'll start them myself"), + ] + mock_start.assert_not_called() + + def test_start_myself_choice_keeps_existing_exit(self, capsys): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="manual"), + pytest.raises(ClickExit), + ): + _maybe_start_services("http://localhost:8080", auto=False, start_services=None) + + assert "Start the platform first" in capsys.readouterr().err + + def test_unreachable_remote_url_selects_remote_connection_without_local_prompt(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice") as mock_prompt, + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "connect_remote" + mock_prompt.assert_not_called() + + def test_rejects_start_services_for_remote_url(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}._start_services_background") as mock_start, + pytest.raises(typer.BadParameter, match="local Platform URL"), + ): + _maybe_start_services("https://remote.example.com", auto=False, start_services=True) + + mock_start.assert_not_called() + def test_restarts_when_running_and_start_services_true(self): reachable_calls = [True, True, False, True] @@ -622,6 +674,160 @@ def test_allows_start_when_port_free(self, maybe_start_preflight_mocks): maybe_start_preflight_mocks.assert_called_once() +class TestRemoteConnection: + def test_prompts_again_until_platform_is_reachable(self, capsys): + with ( + patch( + f"{SETUP_MOD}.prompt_text", + side_effect=["https://unreachable.example.com", "https://remote.example.com/"], + ), + patch( + f"{SETUP_MOD}._check_platform_reachable_with_retries", + side_effect=[False, True], + ), + ): + base_url = setup_commands._prompt_remote_base_url() + + assert base_url == "https://remote.example.com" + assert "Unable to connect" in capsys.readouterr().err + + def test_persists_remote_url_and_workspace_in_active_context(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + cli_context = MagicMock() + cli_context.overrides = {} + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) + + setup_commands._configure_remote_connection( + cli_context, + "https://remote.example.com", + "shared-workspace", + ) + + context = Config.load(config_path=config_path).resolve() + assert str(context.cluster.base_url) == "https://remote.example.com/" + assert context.workspace == "shared-workspace" + cli_context.reset_sdk_context.assert_called_once_with() + + def test_updates_selected_context_and_runtime_url_override(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + Config.write({"base_url": "https://default.example.com"}) + Config.write({"base_url": "https://old-dev.example.com"}, context_name="dev") + + cli_context = MagicMock() + cli_context.overrides = { + "current_context": "dev", + "base_url": "https://stale-override.example.com", + } + cli_context.get_sdk_context.return_value = Context( + context_name="dev", + cluster=Cluster(name="dev-cluster", base_url="https://stale-override.example.com"), + user=NoAuthUser(name="dev-user"), + workspace="default", + preferences={}, + ) + + setup_commands._configure_remote_connection( + cli_context, + "https://new-dev.example.com", + "shared-workspace", + ) + + config_file = Config.load(config_path=config_path).get_config_file() + default_cluster = next(cluster for cluster in config_file.clusters if cluster.name == "default-cluster") + dev_cluster = next(cluster for cluster in config_file.clusters if cluster.name == "dev-cluster") + assert str(default_cluster.base_url) == "https://default.example.com/" + assert str(dev_cluster.base_url) == "https://new-dev.example.com/" + assert cli_context.overrides["base_url"] == "https://new-dev.example.com" + + def test_authenticates_when_context_has_no_credentials(self): + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) + + with patch("nemo_platform.cli.commands.auth._login_with_oidc", return_value=True) as mock_login: + setup_commands._ensure_platform_auth(cli_context) + + mock_login.assert_called_once_with(cli_context, selected_context="default") + cli_context.reset_sdk_context.assert_called_once_with() + + def test_reauthenticates_stored_context_credentials_for_new_remote(self): + cli_context = MagicMock() + context = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=OAuthUser(name="default-user", token=SecretStr("token")), + workspace="default", + preferences={}, + ) + cli_context.get_sdk_context.return_value = context + + with patch("nemo_platform.cli.commands.auth._login_with_oidc", return_value=True) as mock_login: + setup_commands._ensure_platform_auth(cli_context) + + mock_login.assert_called_once_with(cli_context, selected_context="default") + cli_context.reset_sdk_context.assert_called_once_with() + + def test_reuses_runtime_access_token_override(self, monkeypatch): + monkeypatch.setenv("NMP_ACCESS_TOKEN", "runtime-token") + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=OAuthUser(name="default-user", token=SecretStr("runtime-token")), + workspace="default", + preferences={}, + ) + + with patch("nemo_platform.cli.commands.auth._login_with_oidc") as mock_login: + setup_commands._ensure_platform_auth(cli_context) + + mock_login.assert_not_called() + + def test_clears_stale_credentials_when_remote_auth_is_disabled(self, tmp_path, monkeypatch): + config_path = tmp_path / "config.yaml" + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + monkeypatch.delenv("NMP_ACCESS_TOKEN", raising=False) + Config.write( + { + "base_url": "https://remote.example.com", + "access_token": "old-cluster-token", + } + ) + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Config.load(config_path=config_path).resolve() + + with patch("nemo_platform.cli.commands.auth._login_with_oidc", return_value=False): + setup_commands._ensure_platform_auth(cli_context) + + assert isinstance(Config.load(config_path=config_path).resolve().user, NoAuthUser) + cli_context.reset_sdk_context.assert_called_once_with() + + def test_platform_request_headers_include_context_token(self): + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="remote", base_url="https://remote.example.com"), + user=OAuthUser(name="default-user", token=SecretStr("remote-token")), + workspace="default", + preferences={}, + ) + + assert setup_commands._platform_request_headers(cli_context) == {"Authorization": "Bearer remote-token"} + + class TestLocalDataDirHelpers: """Tests for the XDG-default data-dir helpers used by `nemo setup`.""" @@ -2297,6 +2503,42 @@ def test_proceeds_when_tty_without_auto(self): ): self._invoke(auto=False) + def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", side_effect=UserCancelled), + pytest.raises(typer.Exit) as exc_info, + ): + self._invoke(auto=False) + + assert exc_info.value.exit_code == 0 + assert "Setup cancelled" in capsys.readouterr().err + + +class TestSetupCommandRemoteFlow: + def test_remote_choice_connects_before_continuing_setup(self): + ctx = MagicMock(spec=typer.Context) + cli_context = MagicMock() + cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, + patch(f"{SETUP_MOD}._ensure_platform_auth") as mock_auth, + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx) + + mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "default") + mock_auth.assert_called_once_with(cli_context) + assert mock_run.call_args.args[3] == "https://remote.example.com" + # --------------------------------------------------------------------------- # Controller health check diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup_cli.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup_cli.py new file mode 100644 index 0000000000..c63bd9fdfb --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup_cli.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI integration tests for ``nemo setup`` connection selection.""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform.cli.app import app +from nemo_platform.cli.core.context import CLIContext +from nemo_platform.config.config import Config +from typer.testing import CliRunner + +SETUP_MOD = "nemo_platform.cli.commands.setup" + + +def test_remote_choice_retries_and_persists_connection( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.touch() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + + client = MagicMock() + client.workspaces.retrieve.return_value = MagicMock() + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="remote"), + patch( + f"{SETUP_MOD}.prompt_text", + side_effect=["https://unreachable.example.com", "https://remote.example.com/"], + ), + patch( + f"{SETUP_MOD}._check_platform_reachable_with_retries", + side_effect=[False, True, True], + ), + patch(f"{SETUP_MOD}._ensure_platform_auth") as ensure_auth, + patch(f"{SETUP_MOD}._start_services_background") as start_services, + patch.object(CLIContext, "get_client", return_value=client), + patch(f"{SETUP_MOD}._run_interactive_mode") as run_interactive, + ): + result = CliRunner().invoke(app, ["setup"]) + + assert result.exit_code == 0, result.output + assert "Unable to connect to NeMo Platform at https://unreachable.example.com" in result.output + context = Config.load(config_path=config_path).resolve() + assert str(context.cluster.base_url) == "https://remote.example.com/" + start_services.assert_not_called() + ensure_auth.assert_called_once() + assert run_interactive.call_args.args[3] == "https://remote.example.com" + + +def test_local_choice_starts_services_and_keeps_local_connection( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.touch() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + + client = MagicMock() + client.workspaces.retrieve.return_value = MagicMock() + process = MagicMock(pid=1234) + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="yes"), + patch(f"{SETUP_MOD}._prompt_data_dir", return_value="/tmp/nemo-demo"), + patch(f"{SETUP_MOD}.importlib.util.find_spec", return_value=MagicMock()), + patch(f"{SETUP_MOD}._ensure_port_available_for_start"), + patch(f"{SETUP_MOD}._start_services_background", return_value=process) as start_services, + patch(f"{SETUP_MOD}._wait_for_platform", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch.object(CLIContext, "get_client", return_value=client), + patch(f"{SETUP_MOD}._run_interactive_mode") as run_interactive, + ): + result = CliRunner().invoke(app, ["setup"]) + + assert result.exit_code == 0, result.output + context = Config.load(config_path=config_path).resolve() + assert str(context.cluster.base_url) == "http://localhost:8080/" + start_services.assert_called_once() + assert run_interactive.call_args.args[3].rstrip("/") == "http://localhost:8080" + + +def test_start_myself_exits_without_starting_or_mutating_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config_path = tmp_path / "config.yaml" + config_path.touch() + monkeypatch.setenv("NMP_CONFIG_FILE", str(config_path)) + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), + patch(f"{SETUP_MOD}.prompt_choice", return_value="manual"), + patch(f"{SETUP_MOD}._start_services_background") as start_services, + ): + result = CliRunner().invoke(app, ["setup"]) + + assert result.exit_code == 1 + assert "Start the platform first" in result.output + assert config_path.read_text() == "" + start_services.assert_not_called() From 4169615bf002360d7abec640a506b43c55acee43 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Tue, 21 Jul 2026 15:29:01 -0700 Subject: [PATCH 2/8] fix(cli): sync setup help text in manifest registry Keep TopLevelEntry.help aligned with setup_command so test_manifest_help_matches_loaded_manual_entry passes after the remote-connect wording change. AIRCORE-934 Signed-off-by: Tyler Bray --- docs/cli/reference.mdx | 11 ++++++----- .../cli/commands/manifest_registry.py | 11 ++++++----- .../tests/cli/commands/test_agent.py | 2 +- packages/nemo_platform_ext/tests/cli/test_app.py | 2 +- .../nemo_platform/cli/commands/manifest_registry.py | 11 ++++++----- .../nemo_platform_ext/cli/commands/test_agent.py | 2 +- .../tests/vendored/nemo_platform_ext/cli/test_app.py | 2 +- 7 files changed, 22 insertions(+), 19 deletions(-) diff --git a/docs/cli/reference.mdx b/docs/cli/reference.mdx index caae36c381..85a5fae77d 100644 --- a/docs/cli/reference.mdx +++ b/docs/cli/reference.mdx @@ -42,12 +42,12 @@ nemo [GLOBAL OPTIONS] COMMAND [ARGS]... ### nemo setup -Set up NeMo Platform: start services, configure a provider, install skills. +Set up NeMo Platform: connect or start services, configure a provider, install skills. -Walks through starting local services, selecting a provider, entering -credentials, registering the provider with the platform, picking a -default model, installing coding agent skills, and optionally deploying -a demo agent. +Uses an already-running platform, starts local services, or connects the +CLI to an existing remote deployment. Then selects and registers an +inference provider, picks a default model, installs coding agent skills, +and optionally deploys a demo agent. Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -64,6 +64,7 @@ nemo setup nemo setup --auto nemo setup --auto --start-services --install-skills --deploy-agent nemo setup --auto --start-services --ready-timeout 360 +NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent ``` diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py index 98639dbb3b..3cf4440931 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py @@ -34,12 +34,12 @@ TopLevelEntry( import_path="nemo_platform_ext.cli.commands.setup:setup_command", help="""\ -Set up NeMo Platform: start services, configure a provider, install skills. +Set up NeMo Platform: connect or start services, configure a provider, install skills. -Walks through starting local services, selecting a provider, entering -credentials, registering the provider with the platform, picking a -default model, installing coding agent skills, and optionally deploying -a demo agent. +Uses an already-running platform, starts local services, or connects the +CLI to an existing remote deployment. Then selects and registers an +inference provider, picks a default model, installs coding agent skills, +and optionally deploys a demo agent. Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -54,6 +54,7 @@ nemo setup --auto nemo setup --auto --start-services --install-skills --deploy-agent nemo setup --auto --start-services --ready-timeout 360 + NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent""", name="setup", diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_agent.py b/packages/nemo_platform_ext/tests/cli/commands/test_agent.py index 39d618ae7d..4fc204b1a0 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_agent.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_agent.py @@ -175,7 +175,7 @@ def test_commands_uses_visible_command_order_with_discovered_plugins(self): assert result.exit_code == 0 command_rows = [line for line in result.stdout.splitlines() if line.startswith("| nemo ")] assert command_rows == [ - "| nemo setup | Setup | Set up NeMo Platform: start services, configure a provider, install skills. |", + "| nemo setup | Setup | Set up NeMo Platform: connect or start services, configure a provider, install skills. |", "| nemo services | Setup | Run platform services locally. |", "| nemo skills | Setup | Install AI agent skill files for Nemo. |", "| nemo chat | CLI functions | Start an interactive chat session with a model. |", diff --git a/packages/nemo_platform_ext/tests/cli/test_app.py b/packages/nemo_platform_ext/tests/cli/test_app.py index 55b1bf0da6..dad582eaa2 100644 --- a/packages/nemo_platform_ext/tests/cli/test_app.py +++ b/packages/nemo_platform_ext/tests/cli/test_app.py @@ -53,7 +53,7 @@ def test_help_includes_getting_started(): assert "nemo docs --list" in result.stdout assert "nemo services run --help" in result.stdout # Help panel truncates long command descriptions; match the visible prefix. - assert "Set up NeMo Platform: start services" in result.stdout + assert "Set up NeMo Platform: connect or start services" in result.stdout assert "--help, -h" in result.stdout assert "nemo auth login --base-url" not in result.stdout assert "nemo quickstart configure" not in result.stdout diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py index 75ab429697..cadad8e139 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py @@ -34,12 +34,12 @@ TopLevelEntry( import_path="nemo_platform.cli.commands.setup:setup_command", help="""\ -Set up NeMo Platform: start services, configure a provider, install skills. +Set up NeMo Platform: connect or start services, configure a provider, install skills. -Walks through starting local services, selecting a provider, entering -credentials, registering the provider with the platform, picking a -default model, installing coding agent skills, and optionally deploying -a demo agent. +Uses an already-running platform, starts local services, or connects the +CLI to an existing remote deployment. Then selects and registers an +inference provider, picks a default model, installs coding agent skills, +and optionally deploys a demo agent. Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -54,6 +54,7 @@ nemo setup --auto nemo setup --auto --start-services --install-skills --deploy-agent nemo setup --auto --start-services --ready-timeout 360 + NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent""", name="setup", diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_agent.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_agent.py index 6060cfb346..5ed300825e 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_agent.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_agent.py @@ -175,7 +175,7 @@ def test_commands_uses_visible_command_order_with_discovered_plugins(self): assert result.exit_code == 0 command_rows = [line for line in result.stdout.splitlines() if line.startswith("| nemo ")] assert command_rows == [ - "| nemo setup | Setup | Set up NeMo Platform: start services, configure a provider, install skills. |", + "| nemo setup | Setup | Set up NeMo Platform: connect or start services, configure a provider, install skills. |", "| nemo services | Setup | Run platform services locally. |", "| nemo skills | Setup | Install AI agent skill files for Nemo. |", "| nemo chat | CLI functions | Start an interactive chat session with a model. |", diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py index be8ef844d4..b1d074e9f8 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/test_app.py @@ -53,7 +53,7 @@ def test_help_includes_getting_started(): assert "nemo docs --list" in result.stdout assert "nemo services run --help" in result.stdout # Help panel truncates long command descriptions; match the visible prefix. - assert "Set up NeMo Platform: start services" in result.stdout + assert "Set up NeMo Platform: connect or start services" in result.stdout assert "--help, -h" in result.stdout assert "nemo auth login --base-url" not in result.stdout assert "nemo quickstart configure" not in result.stdout From 2ee65056c3258e7dac2b64c80531a705aa95710b Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Tue, 21 Jul 2026 20:11:19 -0700 Subject: [PATCH 3/8] fix(cli): preserve active workspace on remote reconnect Only overwrite the configured workspace during remote setup when --workspace was passed explicitly. AIRCORE-934 Signed-off-by: Tyler Bray --- .../nemo_platform_ext/cli/commands/setup.py | 10 ++ .../tests/cli/commands/test_setup.py | 100 ++++++++++++++++++ .../src/nemo_platform/cli/commands/setup.py | 10 ++ .../cli/commands/test_setup.py | 100 ++++++++++++++++++ 4 files changed, 220 insertions(+) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index b70e744562..143542af07 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -332,6 +332,15 @@ def _prompt_remote_base_url() -> str: console.print(f"{CROSS} Unable to connect to NeMo Platform at {base_url}.") +def _resolve_setup_workspace(ctx: typer.Context, cli_context: CLIContext, workspace: str) -> str: + """Prefer an explicit ``--workspace``; otherwise keep the active context workspace.""" + from click.core import ParameterSource + + if ctx.get_parameter_source("workspace") != ParameterSource.DEFAULT: + return workspace + return cli_context.get_sdk_context().workspace or workspace + + def _configure_remote_connection(cli_context: CLIContext, base_url: str, workspace: str) -> None: """Persist a remote Platform URL in the active CLI context.""" context_name = cli_context.get_sdk_context().context_name @@ -1785,6 +1794,7 @@ def setup_command( service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) if service_result == "connect_remote": base_url = _prompt_remote_base_url() + workspace = _resolve_setup_workspace(ctx, cli_context, workspace) _configure_remote_connection(cli_context, base_url, workspace) _ensure_platform_auth(cli_context) except UserCancelled: diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index 720fb3bb79..518a65d45a 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -13,6 +13,7 @@ import nemo_platform_ext.cli.commands.setup as setup_commands import pytest import typer +from click.core import ParameterSource from click.exceptions import Exit as ClickExit from nemo_platform.resources.inference.providers import ProvidersResource from nemo_platform_ext.cli.commands.setup import ( @@ -57,6 +58,7 @@ _register_provider_interactive, _render_onboarding_card, _resolve_provider_for_url, + _resolve_setup_workspace, _run_interactive_mode, _save_data_dir, _select_default_model, @@ -747,6 +749,35 @@ def test_updates_selected_context_and_runtime_url_override(self, tmp_path, monke assert str(dev_cluster.base_url) == "https://new-dev.example.com/" assert cli_context.overrides["base_url"] == "https://new-dev.example.com" + def test_preserves_active_workspace_when_flag_not_explicit(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + + assert _resolve_setup_workspace(ctx, cli_context, "default") == "team-a" + ctx.get_parameter_source.assert_called_once_with("workspace") + + def test_uses_explicit_workspace_flag_over_active_context(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.COMMANDLINE + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + + assert _resolve_setup_workspace(ctx, cli_context, "shared-workspace") == "shared-workspace" + def test_authenticates_when_context_has_no_credentials(self): cli_context = MagicMock() cli_context.get_sdk_context.return_value = Context( @@ -2518,8 +2549,16 @@ def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): class TestSetupCommandRemoteFlow: def test_remote_choice_connects_before_continuing_setup(self): ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT cli_context = MagicMock() cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() ctx.obj = cli_context @@ -2539,6 +2578,67 @@ def test_remote_choice_connects_before_continuing_setup(self): mock_auth.assert_called_once_with(cli_context) assert mock_run.call_args.args[3] == "https://remote.example.com" + def test_remote_choice_preserves_active_workspace_when_flag_omitted(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT + cli_context = MagicMock() + cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, + patch(f"{SETUP_MOD}._ensure_platform_auth"), + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx) + + mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") + mock_bootstrap.assert_called_once_with("https://remote.example.com", "team-a") + assert mock_run.call_args.args[2] == "team-a" + + def test_remote_choice_uses_explicit_workspace_flag(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.COMMANDLINE + cli_context = MagicMock() + cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, + patch(f"{SETUP_MOD}._ensure_platform_auth"), + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx, workspace="shared-workspace") + + mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") + assert mock_run.call_args.args[2] == "shared-workspace" + # --------------------------------------------------------------------------- # Controller health check diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 52257ed768..721b92606b 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -332,6 +332,15 @@ def _prompt_remote_base_url() -> str: console.print(f"{CROSS} Unable to connect to NeMo Platform at {base_url}.") +def _resolve_setup_workspace(ctx: typer.Context, cli_context: CLIContext, workspace: str) -> str: + """Prefer an explicit ``--workspace``; otherwise keep the active context workspace.""" + from click.core import ParameterSource + + if ctx.get_parameter_source("workspace") != ParameterSource.DEFAULT: + return workspace + return cli_context.get_sdk_context().workspace or workspace + + def _configure_remote_connection(cli_context: CLIContext, base_url: str, workspace: str) -> None: """Persist a remote Platform URL in the active CLI context.""" context_name = cli_context.get_sdk_context().context_name @@ -1785,6 +1794,7 @@ def setup_command( service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) if service_result == "connect_remote": base_url = _prompt_remote_base_url() + workspace = _resolve_setup_workspace(ctx, cli_context, workspace) _configure_remote_connection(cli_context, base_url, workspace) _ensure_platform_auth(cli_context) except UserCancelled: diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index 66c68ea9b6..ac74af7fcd 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -13,6 +13,7 @@ import nemo_platform.cli.commands.setup as setup_commands import pytest import typer +from click.core import ParameterSource from click.exceptions import Exit as ClickExit from nemo_platform.resources.inference.providers import ProvidersResource from nemo_platform.cli.commands.setup import ( @@ -57,6 +58,7 @@ _register_provider_interactive, _render_onboarding_card, _resolve_provider_for_url, + _resolve_setup_workspace, _run_interactive_mode, _save_data_dir, _select_default_model, @@ -747,6 +749,35 @@ def test_updates_selected_context_and_runtime_url_override(self, tmp_path, monke assert str(dev_cluster.base_url) == "https://new-dev.example.com/" assert cli_context.overrides["base_url"] == "https://new-dev.example.com" + def test_preserves_active_workspace_when_flag_not_explicit(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + + assert _resolve_setup_workspace(ctx, cli_context, "default") == "team-a" + ctx.get_parameter_source.assert_called_once_with("workspace") + + def test_uses_explicit_workspace_flag_over_active_context(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.COMMANDLINE + cli_context = MagicMock() + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + + assert _resolve_setup_workspace(ctx, cli_context, "shared-workspace") == "shared-workspace" + def test_authenticates_when_context_has_no_credentials(self): cli_context = MagicMock() cli_context.get_sdk_context.return_value = Context( @@ -2518,8 +2549,16 @@ def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): class TestSetupCommandRemoteFlow: def test_remote_choice_connects_before_continuing_setup(self): ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT cli_context = MagicMock() cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() ctx.obj = cli_context @@ -2539,6 +2578,67 @@ def test_remote_choice_connects_before_continuing_setup(self): mock_auth.assert_called_once_with(cli_context) assert mock_run.call_args.args[3] == "https://remote.example.com" + def test_remote_choice_preserves_active_workspace_when_flag_omitted(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT + cli_context = MagicMock() + cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, + patch(f"{SETUP_MOD}._ensure_platform_auth"), + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx) + + mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") + mock_bootstrap.assert_called_once_with("https://remote.example.com", "team-a") + assert mock_run.call_args.args[2] == "team-a" + + def test_remote_choice_uses_explicit_workspace_flag(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.COMMANDLINE + cli_context = MagicMock() + cli_context.get_base_url.return_value = "http://localhost:8080" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), + user=NoAuthUser(name="default-user"), + workspace="team-a", + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, + patch(f"{SETUP_MOD}._ensure_platform_auth"), + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx, workspace="shared-workspace") + + mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") + assert mock_run.call_args.args[2] == "shared-workspace" + # --------------------------------------------------------------------------- # Controller health check From 3b98a01cd0d5031e7104fd72f51efbd0b8f16c63 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Wed, 22 Jul 2026 14:59:26 -0700 Subject: [PATCH 4/8] fix(cli): support hosted setup endpoints (#934) Signed-off-by: Tyler Bray --- .../nemo_platform_ext/cli/commands/setup.py | 36 ++++-- .../src/nemo_platform_ext/config/config.py | 12 +- .../nemo_platform_ext/quickstart/config.py | 13 ++- .../tests/cli/commands/test_setup.py | 61 +++++++++- .../tests/config/test_config.py | 22 ++++ .../src/nemo_platform/cli/commands/setup.py | 106 +++++++++++------- .../src/nemo_platform/config/config.py | 26 +++-- 7 files changed, 209 insertions(+), 67 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index 143542af07..2fa3229dc5 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -41,6 +41,7 @@ from nemo_platform_ext.cli.commands.skills.registry import get_installer, load_skills from nemo_platform_ext.cli.core.context import CLIContext from nemo_platform_ext.cli.core.errors import handle_errors +from nemo_platform_ext.client.tls import client_verify_from_env from nemo_platform_ext.config.config import Config from nemo_platform_ext.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig from nemo_platform_ext.local.process import ( @@ -285,13 +286,25 @@ def _bootstrap_config_if_missing(base_url: str, workspace: str) -> None: Config.write(params) +_PLATFORM_REACHABILITY_PATHS = ("/status", "/cluster-info") + + def _check_platform_reachable(base_url: str, timeout: float = 5.0) -> bool: - """Return True if the platform health endpoint responds.""" - try: - resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout) - return resp.status_code == 200 - except Exception: - return False + """Return True if a platform health endpoint responds. + + Local ``nemo services run`` publishes ``/status``. Hosted deployments may + only expose ``/cluster-info`` on ingress, so try both. + """ + verify = client_verify_from_env() + root = base_url.rstrip("/") + for path in _PLATFORM_REACHABILITY_PATHS: + try: + resp = httpx.get(f"{root}{path}", timeout=timeout, verify=verify) + if resp.status_code == 200: + return True + except Exception: + continue + return False def _check_platform_reachable_with_retries( @@ -385,14 +398,18 @@ def _check_controller_health(base_url: str, timeout: float = 5.0) -> tuple[bool, """Query ``/status`` and assess controller health. Returns ``(True, "")`` when controllers are populated and all healthy. + Returns ``(True, detail)`` when ``/status`` is not published (hosted ingress). Returns ``(False, detail)`` when unhealthy, unreachable, or empty after retry. If ``controllers.status`` is empty on the first call (startup timing race), waits ``_CONTROLLER_HEALTH_RETRY_DELAY`` seconds and retries once. """ + verify = client_verify_from_env() for attempt in range(2): try: - resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout) + resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout, verify=verify) + if resp.status_code == 404: + return True, "Hosted deployment does not publish /status." if resp.status_code != 200: return False, f"Unexpected status {resp.status_code} from /status endpoint." data = resp.json() @@ -428,6 +445,9 @@ def _verify_platform_health(base_url: str) -> bool: """ ok, detail = _check_controller_health(base_url) if ok: + if detail: + console.print(f"\n{WARN} [yellow]{detail}[/yellow]") + console.print(" Setup may have succeeded, but controller health could not be verified.") return True if "no controllers" in detail.lower(): @@ -1794,6 +1814,8 @@ def setup_command( service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) if service_result == "connect_remote": base_url = _prompt_remote_base_url() + _bootstrap_config_if_missing(base_url, workspace) + cli_context.reset_sdk_context() workspace = _resolve_setup_workspace(ctx, cli_context, workspace) _configure_remote_connection(cli_context, base_url, workspace) _ensure_platform_auth(cli_context) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py index 7adf87c9a6..0a94bdb418 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py @@ -30,6 +30,14 @@ logger = logging.getLogger(__name__) +def _secure_chmod(path: Path, mode: int) -> None: + """Apply filesystem permissions, ignoring failures on dirs we do not own (e.g. /tmp).""" + try: + os.chmod(path, mode) + except PermissionError: + pass + + @dataclass(frozen=True) class _RuntimeAccessTokenSource: token: str @@ -266,7 +274,7 @@ def save(self, config_path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) path.parent.mkdir(parents=True, exist_ok=True) - os.chmod(path.parent, stat.S_IRWXU) # 700 + _secure_chmod(path.parent, stat.S_IRWXU) # 700 # Serialize with secrets revealed using context config_data = self._config_file.model_dump( @@ -279,7 +287,7 @@ def save(self, config_path: Path | None = None) -> None: yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) # Set secure file permissions (owner read/write only) - os.chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _secure_chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 600 # Update stored path if we saved to a new location self._config_path = path diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py index 6a21562158..29cd4beeb5 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py @@ -18,6 +18,15 @@ from ._registry import image_registry_host from .gpu_config import parse_comma_separated_non_negative_integers + +def _secure_chmod(path: Path, mode: int) -> None: + """Apply filesystem permissions, ignoring failures on dirs we do not own (e.g. /tmp).""" + try: + os.chmod(path, mode) + except PermissionError: + pass + + InferenceProviderType = Literal["nvidia-build", "host-gpu"] # Registry and repo placeholder for SDK-stamped nightly/milestone tags. @@ -228,7 +237,7 @@ def save(self, path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) config_path.parent.mkdir(parents=True, exist_ok=True) - os.chmod(config_path.parent, stat.S_IRWXU) # 700 + _secure_chmod(config_path.parent, stat.S_IRWXU) # 700 # Serialize with secrets revealed config_data = self.model_dump( @@ -250,7 +259,7 @@ def save(self, path: Path | None = None) -> None: yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) # Set secure file permissions (owner read/write only) - os.chmod(config_path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _secure_chmod(config_path, stat.S_IRUSR | stat.S_IWUSR) # 600 @classmethod def remove(cls) -> None: diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index 518a65d45a..d2ddae5895 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -7,7 +7,7 @@ import logging import sys -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import httpx import nemo_platform_ext.cli.commands.setup as setup_commands @@ -164,17 +164,44 @@ def test_each_known_provider_resolves(self): class TestCheckPlatformReachable: - def test_reachable(self): + def test_reachable_via_status(self): mock_resp = MagicMock() mock_resp.status_code = 200 - with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp): + with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp) as mock_get: assert _check_platform_reachable("http://localhost:8080") is True + mock_get.assert_called_once_with( + "http://localhost:8080/status", + timeout=5.0, + verify=mock_get.call_args.kwargs["verify"], + ) + + def test_reachable_via_cluster_info_when_status_missing(self): + status_resp = MagicMock() + status_resp.status_code = 404 + cluster_resp = MagicMock() + cluster_resp.status_code = 200 + + def _get(url, **kwargs): + if url.endswith("/status"): + return status_resp + if url.endswith("/cluster-info"): + return cluster_resp + raise AssertionError(f"unexpected url: {url}") + + with patch("nemo_platform_ext.cli.commands.setup.httpx.get", side_effect=_get): + assert _check_platform_reachable("https://nemo-platform-freeplay.dev.aire.nvidia.com") is True + + def test_unreachable_when_all_probes_fail(self): + mock_resp = MagicMock() + mock_resp.status_code = 404 + with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp): + assert _check_platform_reachable("http://localhost:8080") is False def test_unreachable(self): with patch("nemo_platform_ext.cli.commands.setup.httpx.get", side_effect=Exception("conn refused")): assert _check_platform_reachable("http://localhost:8080") is False - def test_non_200_status(self): + def test_non_200_status_without_cluster_info_fallback(self): mock_resp = MagicMock() mock_resp.status_code = 503 with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp): @@ -2606,7 +2633,10 @@ def test_remote_choice_preserves_active_workspace_when_flag_omitted(self): setup_command(ctx) mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") - mock_bootstrap.assert_called_once_with("https://remote.example.com", "team-a") + assert mock_bootstrap.call_args_list == [ + call("https://remote.example.com", "default"), + call("https://remote.example.com", "team-a"), + ] assert mock_run.call_args.args[2] == "team-a" def test_remote_choice_uses_explicit_workspace_flag(self): @@ -2711,6 +2741,14 @@ def test_non_200_response(self): ok, _ = _check_controller_health("http://localhost:8080") assert ok is False + def test_status_not_published_on_hosted_deployment(self): + resp = MagicMock() + resp.status_code = 404 + with patch(f"{SETUP_MOD}.httpx.get", return_value=resp): + ok, msg = _check_controller_health("https://nemo-platform-freeplay.dev.aire.nvidia.com") + assert ok is True + assert "does not publish /status" in msg + def test_invalid_json_response(self): resp = MagicMock() resp.status_code = 200 @@ -2748,6 +2786,19 @@ def test_healthy_returns_true(self): result = _verify_platform_health("http://localhost:8080") assert result is True + def test_missing_status_endpoint_returns_true_with_warning(self): + with ( + patch( + f"{SETUP_MOD}._check_controller_health", + return_value=(True, "Hosted deployment does not publish /status."), + ), + patch(f"{SETUP_MOD}.console") as mock_console, + ): + result = _verify_platform_health("https://nemo-platform-freeplay.dev.aire.nvidia.com") + assert result is True + printed = " ".join(str(c) for c in mock_console.print.call_args_list) + assert "does not publish /status" in printed + def test_unhealthy_prints_red_error(self): with ( patch(f"{SETUP_MOD}._check_controller_health", return_value=(False, "models_controller unhealthy")), diff --git a/packages/nemo_platform_ext/tests/config/test_config.py b/packages/nemo_platform_ext/tests/config/test_config.py index 652c879681..7d867a778b 100644 --- a/packages/nemo_platform_ext/tests/config/test_config.py +++ b/packages/nemo_platform_ext/tests/config/test_config.py @@ -3,7 +3,9 @@ """Unit tests for the config module.""" +import os from pathlib import Path +from unittest.mock import patch import pytest import yaml @@ -865,6 +867,26 @@ def test_save_creates_file_with_600_permissions(self, tmp_path: Path): file_mode = config_path.stat().st_mode & 0o777 assert file_mode == 0o600, f"Expected 600, got {oct(file_mode)}" + def test_save_tolerates_unowned_parent_directory(self, tmp_path: Path): + """Saving under a world-writable parent (e.g. /tmp) should not fail on chmod.""" + config_path = tmp_path / "config.yaml" + + real_chmod = os.chmod + + def chmod_side_effect(path, mode): + if Path(path) == tmp_path: + raise PermissionError("Operation not permitted") + real_chmod(path, mode) + + with patch("nemo_platform_ext.config.config.os.chmod", side_effect=chmod_side_effect): + Config.write( + {"base_url": "http://test.example.com"}, + context_name="default", + config_path=config_path, + ) + + assert config_path.exists() + class TestUserTypeDiscriminator: """Test User type discriminator and deserialization.""" diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 721b92606b..908cdb24d8 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -10,59 +10,60 @@ from __future__ import annotations -import importlib.util -import logging import os -import subprocess import time +import logging +import subprocess +import importlib.util +from typing import Literal, Annotated +from pathlib import Path from dataclasses import dataclass +from urllib.parse import urlparse from importlib.resources import files from importlib.resources.abc import Traversable -from pathlib import Path -from typing import Annotated, Literal -from urllib.parse import urlparse +import yaml as _yaml import httpx import typer -import yaml as _yaml -from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.secrets.client import SecretsClient -from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest -from nmp.common.config import nmp_user_data_dir -from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig -from pydantic import SecretStr from rich import box -from rich.console import Console +from pydantic import SecretStr from rich.panel import Panel +from rich.console import Console +from nmp.common.config import nmp_user_data_dir +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig +from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.secrets.client import SecretsClient -from nemo_platform.cli.commands.skills import registry as skills_registry -from nemo_platform.cli.commands.skills.base import Scope, Skill -from nemo_platform.cli.commands.skills.registry import get_installer, load_skills -from nemo_platform.cli.core.context import CLIContext -from nemo_platform.cli.core.errors import handle_errors -from nemo_platform.config.config import Config -from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig -from nemo_platform.local.process import ( - check_port_available_for_start, - compute_scope, - format_port_conflict, - log_path_for, - start_background, - stop_instance, -) +from nemo_platform import NeMoPlatform +from nemo_platform.client.tls import client_verify_from_env from nemo_platform.ui.prompts import ( UserCancelled, - is_interactive, - non_empty_validator, + prompt_text, prompt_choice, + prompt_select, + is_interactive, prompt_confirm, - prompt_multiselect, prompt_password, - prompt_select, - prompt_text, + prompt_multiselect, + non_empty_validator, provider_name_validator, ) +from nemo_platform.config.config import Config +from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform.local.process import ( + log_path_for, + compute_scope, + stop_instance, + start_background, + format_port_conflict, + check_port_available_for_start, +) +from nemo_platform.cli.core.errors import handle_errors +from nemo_platform.cli.core.context import CLIContext +from nemo_platform.cli.commands.skills import registry as skills_registry +from nemo_platform.cli.commands.skills.base import Scope, Skill +from nemo_platform.cli.commands.skills.registry import load_skills, get_installer logger = logging.getLogger(__name__) console = Console(stderr=True) @@ -285,13 +286,25 @@ def _bootstrap_config_if_missing(base_url: str, workspace: str) -> None: Config.write(params) +_PLATFORM_REACHABILITY_PATHS = ("/status", "/cluster-info") + + def _check_platform_reachable(base_url: str, timeout: float = 5.0) -> bool: - """Return True if the platform health endpoint responds.""" - try: - resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout) - return resp.status_code == 200 - except Exception: - return False + """Return True if a platform health endpoint responds. + + Local ``nemo services run`` publishes ``/status``. Hosted deployments may + only expose ``/cluster-info`` on ingress, so try both. + """ + verify = client_verify_from_env() + root = base_url.rstrip("/") + for path in _PLATFORM_REACHABILITY_PATHS: + try: + resp = httpx.get(f"{root}{path}", timeout=timeout, verify=verify) + if resp.status_code == 200: + return True + except Exception: + continue + return False def _check_platform_reachable_with_retries( @@ -385,14 +398,18 @@ def _check_controller_health(base_url: str, timeout: float = 5.0) -> tuple[bool, """Query ``/status`` and assess controller health. Returns ``(True, "")`` when controllers are populated and all healthy. + Returns ``(True, detail)`` when ``/status`` is not published (hosted ingress). Returns ``(False, detail)`` when unhealthy, unreachable, or empty after retry. If ``controllers.status`` is empty on the first call (startup timing race), waits ``_CONTROLLER_HEALTH_RETRY_DELAY`` seconds and retries once. """ + verify = client_verify_from_env() for attempt in range(2): try: - resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout) + resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout, verify=verify) + if resp.status_code == 404: + return True, "Hosted deployment does not publish /status." if resp.status_code != 200: return False, f"Unexpected status {resp.status_code} from /status endpoint." data = resp.json() @@ -428,6 +445,9 @@ def _verify_platform_health(base_url: str) -> bool: """ ok, detail = _check_controller_health(base_url) if ok: + if detail: + console.print(f"\n{WARN} [yellow]{detail}[/yellow]") + console.print(" Setup may have succeeded, but controller health could not be verified.") return True if "no controllers" in detail.lower(): @@ -1794,6 +1814,8 @@ def setup_command( service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) if service_result == "connect_remote": base_url = _prompt_remote_base_url() + _bootstrap_config_if_missing(base_url, workspace) + cli_context.reset_sdk_context() workspace = _resolve_setup_workspace(ctx, cli_context, workspace) _configure_remote_connection(cli_context, base_url, workspace) _ensure_platform_auth(cli_context) diff --git a/sdk/python/nemo-platform/src/nemo_platform/config/config.py b/sdk/python/nemo-platform/src/nemo_platform/config/config.py index f136ec9019..6b0ac52fe5 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/config/config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/config/config.py @@ -5,24 +5,24 @@ from __future__ import annotations -import logging import os import stat -from dataclasses import dataclass +import logging from pathlib import Path +from dataclasses import dataclass +from typing_extensions import Self import yaml -from pydantic import BaseModel, Field, HttpUrl, PrivateAttr, SecretStr -from typing_extensions import Self +from pydantic import Field, HttpUrl, BaseModel, SecretStr, PrivateAttr from .models import ( - DEFAULT_BASE_URL, DEFAULT_CONTEXT, + DEFAULT_BASE_URL, DEFAULT_WORKSPACE, - ConfigFile, - ConfigParams, Context, OAuthUser, + ConfigFile, + ConfigParams, OutputFormat, TimestampFormat, ) @@ -30,6 +30,14 @@ logger = logging.getLogger(__name__) +def _secure_chmod(path: Path, mode: int) -> None: + """Apply filesystem permissions, ignoring failures on dirs we do not own (e.g. /tmp).""" + try: + os.chmod(path, mode) + except PermissionError: + pass + + @dataclass(frozen=True) class _RuntimeAccessTokenSource: token: str @@ -266,7 +274,7 @@ def save(self, config_path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) path.parent.mkdir(parents=True, exist_ok=True) - os.chmod(path.parent, stat.S_IRWXU) # 700 + _secure_chmod(path.parent, stat.S_IRWXU) # 700 # Serialize with secrets revealed using context config_data = self._config_file.model_dump( @@ -279,7 +287,7 @@ def save(self, config_path: Path | None = None) -> None: yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) # Set secure file permissions (owner read/write only) - os.chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _secure_chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 600 # Update stored path if we saved to a new location self._config_path = path From 9761ea879116245bb372fd651f5ca3ad22b0ff82 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Wed, 22 Jul 2026 18:18:40 -0700 Subject: [PATCH 5/8] fix(cli): harden hosted setup probes and config file perms Address CodeRabbit review and lint-sdk-vendored drift: gate /status 404 on /cluster-info evidence, keep parent-dir chmod tolerant while enforcing 600 credential files, and re-vendor the SDK copy from source. Signed-off-by: Tyler Bray --- .../nemo_platform_ext/cli/commands/setup.py | 26 ++++- .../src/nemo_platform_ext/config/config.py | 32 ++++-- .../nemo_platform_ext/quickstart/config.py | 32 ++++-- .../tests/cli/commands/test_setup.py | 62 ++++++++-- .../tests/config/test_config.py | 25 +++- .../src/nemo_platform/cli/commands/setup.py | 96 +++++++++------- .../src/nemo_platform/config/config.py | 46 +++++--- .../src/nemo_platform/quickstart/config.py | 35 +++++- .../cli/commands/test_setup.py | 107 ++++++++++++++++-- .../nemo_platform_ext/config/test_config.py | 45 ++++++++ 10 files changed, 403 insertions(+), 103 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index 2fa3229dc5..2abb5b8701 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -394,22 +394,34 @@ def _platform_request_headers(cli_context: CLIContext) -> dict[str, str] | None: return {key: value for key, value in headers.items() if isinstance(key, str) and isinstance(value, str)} +def _hosted_platform_without_status(base_url: str, *, timeout: float, verify: str | bool) -> bool: + """Return True when ``/cluster-info`` confirms a hosted platform that omits ``/status``.""" + try: + resp = httpx.get(f"{base_url.rstrip('/')}/cluster-info", timeout=timeout, verify=verify) + except Exception: + return False + return resp.status_code == 200 + + def _check_controller_health(base_url: str, timeout: float = 5.0) -> tuple[bool, str]: """Query ``/status`` and assess controller health. Returns ``(True, "")`` when controllers are populated and all healthy. - Returns ``(True, detail)`` when ``/status`` is not published (hosted ingress). + Returns ``(True, detail)`` when ``/status`` is absent but ``/cluster-info`` confirms a hosted platform. Returns ``(False, detail)`` when unhealthy, unreachable, or empty after retry. If ``controllers.status`` is empty on the first call (startup timing race), waits ``_CONTROLLER_HEALTH_RETRY_DELAY`` seconds and retries once. """ verify = client_verify_from_env() + root = base_url.rstrip("/") for attempt in range(2): try: - resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout, verify=verify) + resp = httpx.get(f"{root}/status", timeout=timeout, verify=verify) if resp.status_code == 404: - return True, "Hosted deployment does not publish /status." + if _hosted_platform_without_status(root, timeout=timeout, verify=verify): + return True, "Hosted deployment does not publish /status." + return False, "Unexpected status 404 from /status endpoint." if resp.status_code != 200: return False, f"Unexpected status {resp.status_code} from /status endpoint." data = resp.json() @@ -446,8 +458,12 @@ def _verify_platform_health(base_url: str) -> bool: ok, detail = _check_controller_health(base_url) if ok: if detail: - console.print(f"\n{WARN} [yellow]{detail}[/yellow]") - console.print(" Setup may have succeeded, but controller health could not be verified.") + if "does not publish /status" in detail.lower(): + # Expected for hosted ingress that only exposes /cluster-info. + console.print(f"\n{CHECK} {detail}") + else: + console.print(f"\n{WARN} [yellow]{detail}[/yellow]") + console.print(" Setup may have succeeded, but controller health could not be verified.") return True if "no controllers" in detail.lower(): diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py index 0a94bdb418..086de28a1b 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py @@ -30,14 +30,32 @@ logger = logging.getLogger(__name__) -def _secure_chmod(path: Path, mode: int) -> None: - """Apply filesystem permissions, ignoring failures on dirs we do not own (e.g. /tmp).""" +def _try_secure_chmod_dir(path: Path) -> None: + """Best-effort owner-only permissions on a directory (e.g. ``/tmp`` may reject chmod).""" try: - os.chmod(path, mode) + os.chmod(path, stat.S_IRWXU) # 700 except PermissionError: pass +def _write_secure_yaml(path: Path, config_data: dict) -> None: + """Write YAML atomically with owner read/write permissions (600). + + File permission failures remain fatal so credentials are never left world-readable. + """ + mode = stat.S_IRUSR | stat.S_IWUSR # 600 + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) + try: + with os.fdopen(fd, "w") as f: + fd = -1 + yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) + finally: + if fd >= 0: + os.close(fd) + # Existing files keep prior mode under O_TRUNC; enforce 600 after write. + os.chmod(path, mode) + + @dataclass(frozen=True) class _RuntimeAccessTokenSource: token: str @@ -274,7 +292,7 @@ def save(self, config_path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) path.parent.mkdir(parents=True, exist_ok=True) - _secure_chmod(path.parent, stat.S_IRWXU) # 700 + _try_secure_chmod_dir(path.parent) # Serialize with secrets revealed using context config_data = self._config_file.model_dump( @@ -283,11 +301,7 @@ def save(self, config_path: Path | None = None) -> None: context={"include_secrets": True}, ) - with open(path, "w") as f: - yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) - - # Set secure file permissions (owner read/write only) - _secure_chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _write_secure_yaml(path, config_data) # Update stored path if we saved to a new location self._config_path = path diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py index 29cd4beeb5..c7907c08dd 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/quickstart/config.py @@ -19,14 +19,32 @@ from .gpu_config import parse_comma_separated_non_negative_integers -def _secure_chmod(path: Path, mode: int) -> None: - """Apply filesystem permissions, ignoring failures on dirs we do not own (e.g. /tmp).""" +def _try_secure_chmod_dir(path: Path) -> None: + """Best-effort owner-only permissions on a directory (e.g. ``/tmp`` may reject chmod).""" try: - os.chmod(path, mode) + os.chmod(path, stat.S_IRWXU) # 700 except PermissionError: pass +def _write_secure_yaml(path: Path, config_data: dict) -> None: + """Write YAML atomically with owner read/write permissions (600). + + File permission failures remain fatal so credentials are never left world-readable. + """ + mode = stat.S_IRUSR | stat.S_IWUSR # 600 + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) + try: + with os.fdopen(fd, "w") as f: + fd = -1 + yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) + finally: + if fd >= 0: + os.close(fd) + # Existing files keep prior mode under O_TRUNC; enforce 600 after write. + os.chmod(path, mode) + + InferenceProviderType = Literal["nvidia-build", "host-gpu"] # Registry and repo placeholder for SDK-stamped nightly/milestone tags. @@ -237,7 +255,7 @@ def save(self, path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) config_path.parent.mkdir(parents=True, exist_ok=True) - _secure_chmod(config_path.parent, stat.S_IRWXU) # 700 + _try_secure_chmod_dir(config_path.parent) # Serialize with secrets revealed config_data = self.model_dump( @@ -255,11 +273,7 @@ def save(self, path: Path | None = None) -> None: if "platform_config_path" in config_data: config_data["platform_config_path"] = str(config_data["platform_config_path"]) - with open(config_path, "w") as f: - yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) - - # Set secure file permissions (owner read/write only) - _secure_chmod(config_path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _write_secure_yaml(config_path, config_data) @classmethod def remove(cls) -> None: diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index d2ddae5895..94120177e9 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -167,12 +167,15 @@ class TestCheckPlatformReachable: def test_reachable_via_status(self): mock_resp = MagicMock() mock_resp.status_code = 200 - with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp) as mock_get: + with ( + patch(f"{SETUP_MOD}.client_verify_from_env", return_value="/tmp/custom-ca.pem"), + patch(f"{SETUP_MOD}.httpx.get", return_value=mock_resp) as mock_get, + ): assert _check_platform_reachable("http://localhost:8080") is True mock_get.assert_called_once_with( "http://localhost:8080/status", timeout=5.0, - verify=mock_get.call_args.kwargs["verify"], + verify="/tmp/custom-ca.pem", ) def test_reachable_via_cluster_info_when_status_missing(self): @@ -180,31 +183,40 @@ def test_reachable_via_cluster_info_when_status_missing(self): status_resp.status_code = 404 cluster_resp = MagicMock() cluster_resp.status_code = 200 + calls: list[tuple[str, object]] = [] def _get(url, **kwargs): + calls.append((url, kwargs.get("verify"))) if url.endswith("/status"): return status_resp if url.endswith("/cluster-info"): return cluster_resp raise AssertionError(f"unexpected url: {url}") - with patch("nemo_platform_ext.cli.commands.setup.httpx.get", side_effect=_get): + with ( + patch(f"{SETUP_MOD}.client_verify_from_env", return_value="/tmp/custom-ca.pem"), + patch(f"{SETUP_MOD}.httpx.get", side_effect=_get), + ): assert _check_platform_reachable("https://nemo-platform-freeplay.dev.aire.nvidia.com") is True + assert calls == [ + ("https://nemo-platform-freeplay.dev.aire.nvidia.com/status", "/tmp/custom-ca.pem"), + ("https://nemo-platform-freeplay.dev.aire.nvidia.com/cluster-info", "/tmp/custom-ca.pem"), + ] def test_unreachable_when_all_probes_fail(self): mock_resp = MagicMock() mock_resp.status_code = 404 - with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp): + with patch(f"{SETUP_MOD}.httpx.get", return_value=mock_resp): assert _check_platform_reachable("http://localhost:8080") is False def test_unreachable(self): - with patch("nemo_platform_ext.cli.commands.setup.httpx.get", side_effect=Exception("conn refused")): + with patch(f"{SETUP_MOD}.httpx.get", side_effect=Exception("conn refused")): assert _check_platform_reachable("http://localhost:8080") is False def test_non_200_status_without_cluster_info_fallback(self): mock_resp = MagicMock() mock_resp.status_code = 503 - with patch("nemo_platform_ext.cli.commands.setup.httpx.get", return_value=mock_resp): + with patch(f"{SETUP_MOD}.httpx.get", return_value=mock_resp): assert _check_platform_reachable("http://localhost:8080") is False @@ -2742,13 +2754,41 @@ def test_non_200_response(self): assert ok is False def test_status_not_published_on_hosted_deployment(self): - resp = MagicMock() - resp.status_code = 404 - with patch(f"{SETUP_MOD}.httpx.get", return_value=resp): + status_resp = MagicMock() + status_resp.status_code = 404 + cluster_resp = MagicMock() + cluster_resp.status_code = 200 + + def _get(url, **kwargs): + if url.endswith("/status"): + return status_resp + if url.endswith("/cluster-info"): + return cluster_resp + raise AssertionError(f"unexpected url: {url}") + + with patch(f"{SETUP_MOD}.httpx.get", side_effect=_get): ok, msg = _check_controller_health("https://nemo-platform-freeplay.dev.aire.nvidia.com") assert ok is True assert "does not publish /status" in msg + def test_status_404_without_cluster_info_is_unhealthy(self): + status_resp = MagicMock() + status_resp.status_code = 404 + cluster_resp = MagicMock() + cluster_resp.status_code = 404 + + def _get(url, **kwargs): + if url.endswith("/status"): + return status_resp + if url.endswith("/cluster-info"): + return cluster_resp + raise AssertionError(f"unexpected url: {url}") + + with patch(f"{SETUP_MOD}.httpx.get", side_effect=_get): + ok, msg = _check_controller_health("http://localhost:8080") + assert ok is False + assert "404" in msg + def test_invalid_json_response(self): resp = MagicMock() resp.status_code = 200 @@ -2786,7 +2826,7 @@ def test_healthy_returns_true(self): result = _verify_platform_health("http://localhost:8080") assert result is True - def test_missing_status_endpoint_returns_true_with_warning(self): + def test_missing_status_endpoint_returns_true_with_info(self): with ( patch( f"{SETUP_MOD}._check_controller_health", @@ -2798,6 +2838,8 @@ def test_missing_status_endpoint_returns_true_with_warning(self): assert result is True printed = " ".join(str(c) for c in mock_console.print.call_args_list) assert "does not publish /status" in printed + assert "could not be verified" not in printed + assert "yellow" not in printed.lower() def test_unhealthy_prints_red_error(self): with ( diff --git a/packages/nemo_platform_ext/tests/config/test_config.py b/packages/nemo_platform_ext/tests/config/test_config.py index 7d867a778b..49bc1ae069 100644 --- a/packages/nemo_platform_ext/tests/config/test_config.py +++ b/packages/nemo_platform_ext/tests/config/test_config.py @@ -868,7 +868,7 @@ def test_save_creates_file_with_600_permissions(self, tmp_path: Path): assert file_mode == 0o600, f"Expected 600, got {oct(file_mode)}" def test_save_tolerates_unowned_parent_directory(self, tmp_path: Path): - """Saving under a world-writable parent (e.g. /tmp) should not fail on chmod.""" + """Saving under a world-writable parent (e.g. /tmp) should not fail on dir chmod.""" config_path = tmp_path / "config.yaml" real_chmod = os.chmod @@ -886,6 +886,29 @@ def chmod_side_effect(path, mode): ) assert config_path.exists() + file_mode = config_path.stat().st_mode & 0o777 + assert file_mode == 0o600, f"Expected 600, got {oct(file_mode)}" + + def test_save_raises_when_file_chmod_fails(self, tmp_path: Path): + """Credential file permission failures must remain fatal.""" + config_path = tmp_path / "config.yaml" + + real_chmod = os.chmod + + def chmod_side_effect(path, mode): + if Path(path) == config_path: + raise PermissionError("Operation not permitted") + real_chmod(path, mode) + + with ( + patch("nemo_platform_ext.config.config.os.chmod", side_effect=chmod_side_effect), + pytest.raises(PermissionError), + ): + Config.write( + {"base_url": "http://test.example.com"}, + context_name="default", + config_path=config_path, + ) class TestUserTypeDiscriminator: diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 908cdb24d8..6a35f60530 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -10,60 +10,60 @@ from __future__ import annotations -import os -import time +import importlib.util import logging +import os import subprocess -import importlib.util -from typing import Literal, Annotated -from pathlib import Path +import time from dataclasses import dataclass -from urllib.parse import urlparse from importlib.resources import files from importlib.resources.abc import Traversable +from pathlib import Path +from typing import Annotated, Literal +from urllib.parse import urlparse -import yaml as _yaml import httpx import typer -from rich import box -from pydantic import SecretStr -from rich.panel import Panel -from rich.console import Console -from nmp.common.config import nmp_user_data_dir -from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig -from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest +import yaml as _yaml +from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.secrets.client import SecretsClient +from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest +from nmp.common.config import nmp_user_data_dir +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig +from pydantic import SecretStr +from rich import box +from rich.console import Console +from rich.panel import Panel -from nemo_platform import NeMoPlatform +from nemo_platform.cli.commands.skills import registry as skills_registry +from nemo_platform.cli.commands.skills.base import Scope, Skill +from nemo_platform.cli.commands.skills.registry import get_installer, load_skills +from nemo_platform.cli.core.context import CLIContext +from nemo_platform.cli.core.errors import handle_errors from nemo_platform.client.tls import client_verify_from_env +from nemo_platform.config.config import Config +from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform.local.process import ( + check_port_available_for_start, + compute_scope, + format_port_conflict, + log_path_for, + start_background, + stop_instance, +) from nemo_platform.ui.prompts import ( UserCancelled, - prompt_text, - prompt_choice, - prompt_select, is_interactive, + non_empty_validator, + prompt_choice, prompt_confirm, - prompt_password, prompt_multiselect, - non_empty_validator, + prompt_password, + prompt_select, + prompt_text, provider_name_validator, ) -from nemo_platform.config.config import Config -from nemo_platform.config.models import DEFAULT_BASE_URL, ConfigFile, ConfigParams, LocalServicesConfig -from nemo_platform.local.process import ( - log_path_for, - compute_scope, - stop_instance, - start_background, - format_port_conflict, - check_port_available_for_start, -) -from nemo_platform.cli.core.errors import handle_errors -from nemo_platform.cli.core.context import CLIContext -from nemo_platform.cli.commands.skills import registry as skills_registry -from nemo_platform.cli.commands.skills.base import Scope, Skill -from nemo_platform.cli.commands.skills.registry import load_skills, get_installer logger = logging.getLogger(__name__) console = Console(stderr=True) @@ -394,22 +394,34 @@ def _platform_request_headers(cli_context: CLIContext) -> dict[str, str] | None: return {key: value for key, value in headers.items() if isinstance(key, str) and isinstance(value, str)} +def _hosted_platform_without_status(base_url: str, *, timeout: float, verify: str | bool) -> bool: + """Return True when ``/cluster-info`` confirms a hosted platform that omits ``/status``.""" + try: + resp = httpx.get(f"{base_url.rstrip('/')}/cluster-info", timeout=timeout, verify=verify) + except Exception: + return False + return resp.status_code == 200 + + def _check_controller_health(base_url: str, timeout: float = 5.0) -> tuple[bool, str]: """Query ``/status`` and assess controller health. Returns ``(True, "")`` when controllers are populated and all healthy. - Returns ``(True, detail)`` when ``/status`` is not published (hosted ingress). + Returns ``(True, detail)`` when ``/status`` is absent but ``/cluster-info`` confirms a hosted platform. Returns ``(False, detail)`` when unhealthy, unreachable, or empty after retry. If ``controllers.status`` is empty on the first call (startup timing race), waits ``_CONTROLLER_HEALTH_RETRY_DELAY`` seconds and retries once. """ verify = client_verify_from_env() + root = base_url.rstrip("/") for attempt in range(2): try: - resp = httpx.get(f"{base_url.rstrip('/')}/status", timeout=timeout, verify=verify) + resp = httpx.get(f"{root}/status", timeout=timeout, verify=verify) if resp.status_code == 404: - return True, "Hosted deployment does not publish /status." + if _hosted_platform_without_status(root, timeout=timeout, verify=verify): + return True, "Hosted deployment does not publish /status." + return False, "Unexpected status 404 from /status endpoint." if resp.status_code != 200: return False, f"Unexpected status {resp.status_code} from /status endpoint." data = resp.json() @@ -446,8 +458,12 @@ def _verify_platform_health(base_url: str) -> bool: ok, detail = _check_controller_health(base_url) if ok: if detail: - console.print(f"\n{WARN} [yellow]{detail}[/yellow]") - console.print(" Setup may have succeeded, but controller health could not be verified.") + if "does not publish /status" in detail.lower(): + # Expected for hosted ingress that only exposes /cluster-info. + console.print(f"\n{CHECK} {detail}") + else: + console.print(f"\n{WARN} [yellow]{detail}[/yellow]") + console.print(" Setup may have succeeded, but controller health could not be verified.") return True if "no controllers" in detail.lower(): diff --git a/sdk/python/nemo-platform/src/nemo_platform/config/config.py b/sdk/python/nemo-platform/src/nemo_platform/config/config.py index 6b0ac52fe5..c4e34a40f8 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/config/config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/config/config.py @@ -5,24 +5,24 @@ from __future__ import annotations +import logging import os import stat -import logging -from pathlib import Path from dataclasses import dataclass -from typing_extensions import Self +from pathlib import Path import yaml -from pydantic import Field, HttpUrl, BaseModel, SecretStr, PrivateAttr +from pydantic import BaseModel, Field, HttpUrl, PrivateAttr, SecretStr +from typing_extensions import Self from .models import ( - DEFAULT_CONTEXT, DEFAULT_BASE_URL, + DEFAULT_CONTEXT, DEFAULT_WORKSPACE, - Context, - OAuthUser, ConfigFile, ConfigParams, + Context, + OAuthUser, OutputFormat, TimestampFormat, ) @@ -30,14 +30,32 @@ logger = logging.getLogger(__name__) -def _secure_chmod(path: Path, mode: int) -> None: - """Apply filesystem permissions, ignoring failures on dirs we do not own (e.g. /tmp).""" +def _try_secure_chmod_dir(path: Path) -> None: + """Best-effort owner-only permissions on a directory (e.g. ``/tmp`` may reject chmod).""" try: - os.chmod(path, mode) + os.chmod(path, stat.S_IRWXU) # 700 except PermissionError: pass +def _write_secure_yaml(path: Path, config_data: dict) -> None: + """Write YAML atomically with owner read/write permissions (600). + + File permission failures remain fatal so credentials are never left world-readable. + """ + mode = stat.S_IRUSR | stat.S_IWUSR # 600 + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) + try: + with os.fdopen(fd, "w") as f: + fd = -1 + yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) + finally: + if fd >= 0: + os.close(fd) + # Existing files keep prior mode under O_TRUNC; enforce 600 after write. + os.chmod(path, mode) + + @dataclass(frozen=True) class _RuntimeAccessTokenSource: token: str @@ -274,7 +292,7 @@ def save(self, config_path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) path.parent.mkdir(parents=True, exist_ok=True) - _secure_chmod(path.parent, stat.S_IRWXU) # 700 + _try_secure_chmod_dir(path.parent) # Serialize with secrets revealed using context config_data = self._config_file.model_dump( @@ -283,11 +301,7 @@ def save(self, config_path: Path | None = None) -> None: context={"include_secrets": True}, ) - with open(path, "w") as f: - yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) - - # Set secure file permissions (owner read/write only) - _secure_chmod(path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _write_secure_yaml(path, config_data) # Update stored path if we saved to a new location self._config_path = path diff --git a/sdk/python/nemo-platform/src/nemo_platform/quickstart/config.py b/sdk/python/nemo-platform/src/nemo_platform/quickstart/config.py index 6a21562158..c7907c08dd 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/quickstart/config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/quickstart/config.py @@ -18,6 +18,33 @@ from ._registry import image_registry_host from .gpu_config import parse_comma_separated_non_negative_integers + +def _try_secure_chmod_dir(path: Path) -> None: + """Best-effort owner-only permissions on a directory (e.g. ``/tmp`` may reject chmod).""" + try: + os.chmod(path, stat.S_IRWXU) # 700 + except PermissionError: + pass + + +def _write_secure_yaml(path: Path, config_data: dict) -> None: + """Write YAML atomically with owner read/write permissions (600). + + File permission failures remain fatal so credentials are never left world-readable. + """ + mode = stat.S_IRUSR | stat.S_IWUSR # 600 + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) + try: + with os.fdopen(fd, "w") as f: + fd = -1 + yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) + finally: + if fd >= 0: + os.close(fd) + # Existing files keep prior mode under O_TRUNC; enforce 600 after write. + os.chmod(path, mode) + + InferenceProviderType = Literal["nvidia-build", "host-gpu"] # Registry and repo placeholder for SDK-stamped nightly/milestone tags. @@ -228,7 +255,7 @@ def save(self, path: Path | None = None) -> None: # Ensure parent directory exists with secure permissions (owner-only access) config_path.parent.mkdir(parents=True, exist_ok=True) - os.chmod(config_path.parent, stat.S_IRWXU) # 700 + _try_secure_chmod_dir(config_path.parent) # Serialize with secrets revealed config_data = self.model_dump( @@ -246,11 +273,7 @@ def save(self, path: Path | None = None) -> None: if "platform_config_path" in config_data: config_data["platform_config_path"] = str(config_data["platform_config_path"]) - with open(config_path, "w") as f: - yaml.safe_dump(config_data, f, default_flow_style=False, sort_keys=False) - - # Set secure file permissions (owner read/write only) - os.chmod(config_path, stat.S_IRUSR | stat.S_IWUSR) # 600 + _write_secure_yaml(config_path, config_data) @classmethod def remove(cls) -> None: diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index ac74af7fcd..860e55c7ec 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -7,7 +7,7 @@ import logging import sys -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import httpx import nemo_platform.cli.commands.setup as setup_commands @@ -164,20 +164,59 @@ def test_each_known_provider_resolves(self): class TestCheckPlatformReachable: - def test_reachable(self): + def test_reachable_via_status(self): mock_resp = MagicMock() mock_resp.status_code = 200 - with patch("nemo_platform.cli.commands.setup.httpx.get", return_value=mock_resp): + with ( + patch(f"{SETUP_MOD}.client_verify_from_env", return_value="/tmp/custom-ca.pem"), + patch(f"{SETUP_MOD}.httpx.get", return_value=mock_resp) as mock_get, + ): assert _check_platform_reachable("http://localhost:8080") is True + mock_get.assert_called_once_with( + "http://localhost:8080/status", + timeout=5.0, + verify="/tmp/custom-ca.pem", + ) + + def test_reachable_via_cluster_info_when_status_missing(self): + status_resp = MagicMock() + status_resp.status_code = 404 + cluster_resp = MagicMock() + cluster_resp.status_code = 200 + calls: list[tuple[str, object]] = [] + + def _get(url, **kwargs): + calls.append((url, kwargs.get("verify"))) + if url.endswith("/status"): + return status_resp + if url.endswith("/cluster-info"): + return cluster_resp + raise AssertionError(f"unexpected url: {url}") + + with ( + patch(f"{SETUP_MOD}.client_verify_from_env", return_value="/tmp/custom-ca.pem"), + patch(f"{SETUP_MOD}.httpx.get", side_effect=_get), + ): + assert _check_platform_reachable("https://nemo-platform-freeplay.dev.aire.nvidia.com") is True + assert calls == [ + ("https://nemo-platform-freeplay.dev.aire.nvidia.com/status", "/tmp/custom-ca.pem"), + ("https://nemo-platform-freeplay.dev.aire.nvidia.com/cluster-info", "/tmp/custom-ca.pem"), + ] + + def test_unreachable_when_all_probes_fail(self): + mock_resp = MagicMock() + mock_resp.status_code = 404 + with patch(f"{SETUP_MOD}.httpx.get", return_value=mock_resp): + assert _check_platform_reachable("http://localhost:8080") is False def test_unreachable(self): - with patch("nemo_platform.cli.commands.setup.httpx.get", side_effect=Exception("conn refused")): + with patch(f"{SETUP_MOD}.httpx.get", side_effect=Exception("conn refused")): assert _check_platform_reachable("http://localhost:8080") is False - def test_non_200_status(self): + def test_non_200_status_without_cluster_info_fallback(self): mock_resp = MagicMock() mock_resp.status_code = 503 - with patch("nemo_platform.cli.commands.setup.httpx.get", return_value=mock_resp): + with patch(f"{SETUP_MOD}.httpx.get", return_value=mock_resp): assert _check_platform_reachable("http://localhost:8080") is False @@ -2606,7 +2645,10 @@ def test_remote_choice_preserves_active_workspace_when_flag_omitted(self): setup_command(ctx) mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") - mock_bootstrap.assert_called_once_with("https://remote.example.com", "team-a") + assert mock_bootstrap.call_args_list == [ + call("https://remote.example.com", "default"), + call("https://remote.example.com", "team-a"), + ] assert mock_run.call_args.args[2] == "team-a" def test_remote_choice_uses_explicit_workspace_flag(self): @@ -2711,6 +2753,42 @@ def test_non_200_response(self): ok, _ = _check_controller_health("http://localhost:8080") assert ok is False + def test_status_not_published_on_hosted_deployment(self): + status_resp = MagicMock() + status_resp.status_code = 404 + cluster_resp = MagicMock() + cluster_resp.status_code = 200 + + def _get(url, **kwargs): + if url.endswith("/status"): + return status_resp + if url.endswith("/cluster-info"): + return cluster_resp + raise AssertionError(f"unexpected url: {url}") + + with patch(f"{SETUP_MOD}.httpx.get", side_effect=_get): + ok, msg = _check_controller_health("https://nemo-platform-freeplay.dev.aire.nvidia.com") + assert ok is True + assert "does not publish /status" in msg + + def test_status_404_without_cluster_info_is_unhealthy(self): + status_resp = MagicMock() + status_resp.status_code = 404 + cluster_resp = MagicMock() + cluster_resp.status_code = 404 + + def _get(url, **kwargs): + if url.endswith("/status"): + return status_resp + if url.endswith("/cluster-info"): + return cluster_resp + raise AssertionError(f"unexpected url: {url}") + + with patch(f"{SETUP_MOD}.httpx.get", side_effect=_get): + ok, msg = _check_controller_health("http://localhost:8080") + assert ok is False + assert "404" in msg + def test_invalid_json_response(self): resp = MagicMock() resp.status_code = 200 @@ -2748,6 +2826,21 @@ def test_healthy_returns_true(self): result = _verify_platform_health("http://localhost:8080") assert result is True + def test_missing_status_endpoint_returns_true_with_info(self): + with ( + patch( + f"{SETUP_MOD}._check_controller_health", + return_value=(True, "Hosted deployment does not publish /status."), + ), + patch(f"{SETUP_MOD}.console") as mock_console, + ): + result = _verify_platform_health("https://nemo-platform-freeplay.dev.aire.nvidia.com") + assert result is True + printed = " ".join(str(c) for c in mock_console.print.call_args_list) + assert "does not publish /status" in printed + assert "could not be verified" not in printed + assert "yellow" not in printed.lower() + def test_unhealthy_prints_red_error(self): with ( patch(f"{SETUP_MOD}._check_controller_health", return_value=(False, "models_controller unhealthy")), diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py index 5af80018c9..47391148e1 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/config/test_config.py @@ -3,7 +3,9 @@ """Unit tests for the config module.""" +import os from pathlib import Path +from unittest.mock import patch import pytest import yaml @@ -865,6 +867,49 @@ def test_save_creates_file_with_600_permissions(self, tmp_path: Path): file_mode = config_path.stat().st_mode & 0o777 assert file_mode == 0o600, f"Expected 600, got {oct(file_mode)}" + def test_save_tolerates_unowned_parent_directory(self, tmp_path: Path): + """Saving under a world-writable parent (e.g. /tmp) should not fail on dir chmod.""" + config_path = tmp_path / "config.yaml" + + real_chmod = os.chmod + + def chmod_side_effect(path, mode): + if Path(path) == tmp_path: + raise PermissionError("Operation not permitted") + real_chmod(path, mode) + + with patch("nemo_platform.config.config.os.chmod", side_effect=chmod_side_effect): + Config.write( + {"base_url": "http://test.example.com"}, + context_name="default", + config_path=config_path, + ) + + assert config_path.exists() + file_mode = config_path.stat().st_mode & 0o777 + assert file_mode == 0o600, f"Expected 600, got {oct(file_mode)}" + + def test_save_raises_when_file_chmod_fails(self, tmp_path: Path): + """Credential file permission failures must remain fatal.""" + config_path = tmp_path / "config.yaml" + + real_chmod = os.chmod + + def chmod_side_effect(path, mode): + if Path(path) == config_path: + raise PermissionError("Operation not permitted") + real_chmod(path, mode) + + with ( + patch("nemo_platform.config.config.os.chmod", side_effect=chmod_side_effect), + pytest.raises(PermissionError), + ): + Config.write( + {"base_url": "http://test.example.com"}, + context_name="default", + config_path=config_path, + ) + class TestUserTypeDiscriminator: """Test User type discriminator and deserialization.""" From ef4d400d9d53fbb1fdc1ca645c3ef80d76366328 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Thu, 23 Jul 2026 08:56:16 -0700 Subject: [PATCH 6/8] feat(cli): re-prompt when remote platform is reachable during setup (#934) When setup finds a reachable remote URL, offer continue, start local, or connect to a different host instead of silently continuing. Document URL override escape hatches in help and prefill the remote URL prompt. Signed-off-by: Tyler Bray --- docs/cli/reference.mdx | 11 +++ .../cli/commands/manifest_registry.py | 13 ++- .../nemo_platform_ext/cli/commands/setup.py | 64 ++++++++++++-- .../tests/cli/commands/test_setup.py | 85 ++++++++++++++++++- .../cli/commands/manifest_registry.py | 13 ++- .../src/nemo_platform/cli/commands/setup.py | 64 ++++++++++++-- .../cli/commands/test_setup.py | 85 ++++++++++++++++++- 7 files changed, 321 insertions(+), 14 deletions(-) diff --git a/docs/cli/reference.mdx b/docs/cli/reference.mdx index 85a5fae77d..4a6b0ab135 100644 --- a/docs/cli/reference.mdx +++ b/docs/cli/reference.mdx @@ -49,6 +49,16 @@ CLI to an existing remote deployment. Then selects and registers an inference provider, picks a default model, installs coding agent skills, and optionally deploys a demo agent. +The active config context remembers the Platform URL. When a remote +deployment is already reachable, setup asks whether to continue with it, +start local services instead, or connect to a different remote URL. + +To override the URL for one run only: + nemo --base-url http://localhost:8080 setup + +To persist a different URL: + nemo config set --base-url http://localhost:8080 + Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -67,6 +77,7 @@ nemo setup --auto --start-services --ready-timeout 360 NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent +nemo --base-url http://localhost:8080 setup ``` **Usage:** diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py index 3cf4440931..13d5013840 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/manifest_registry.py @@ -41,6 +41,16 @@ inference provider, picks a default model, installs coding agent skills, and optionally deploys a demo agent. +The active config context remembers the Platform URL. When a remote +deployment is already reachable, setup asks whether to continue with it, +start local services instead, or connect to a different remote URL. + +To override the URL for one run only: + nemo --base-url http://localhost:8080 setup + +To persist a different URL: + nemo config set --base-url http://localhost:8080 + Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -56,7 +66,8 @@ nemo setup --auto --start-services --ready-timeout 360 NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace - nemo setup --no-install-skills --no-deploy-agent""", + nemo setup --no-install-skills --no-deploy-agent + nemo --base-url http://localhost:8080 setup""", name="setup", panel="Setup", kind="command", diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index 2abb5b8701..6c7b33ee3d 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -326,11 +326,12 @@ def _check_platform_reachable_with_retries( return False -def _prompt_remote_base_url() -> str: +def _prompt_remote_base_url(*, default_url: str = "") -> str: """Prompt until the user provides a reachable remote Platform URL.""" while True: base_url = prompt_text( "Enter the remote Platform base URL: ", + default=default_url, validator=non_empty_validator("Base URL"), ).strip() parsed = urlparse(base_url) @@ -736,6 +737,32 @@ def _is_local_base_url(base_url: str) -> bool: return parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"} +def _platform_host_label(base_url: str) -> str: + """Return a human-friendly host label for connection prompts.""" + parsed = urlparse(base_url) + return parsed.hostname or base_url.rstrip("/") + + +def _prompt_reachable_remote_connection(base_url: str) -> Literal["ready", "connect_remote", "start_local"]: + """Ask how to proceed when a configured remote Platform is already reachable.""" + hostname = _platform_host_label(base_url) + action = prompt_choice( + message=f"Platform reachable at {hostname} ({base_url}). What would you like to do?", + options=[ + ("continue", "Continue with this remote Platform"), + ("local", "Start local services instead"), + ("change", "Connect to a different remote URL"), + ], + default="continue", + ) + if action == "continue": + console.print(f"{CHECK} Platform already running at {base_url}\n") + return "ready" + if action == "change": + return "connect_remote" + return "start_local" + + def _start_services_background(base_url: str, data_dir: str | None = None) -> subprocess.Popen: """Launch ``nemo services run`` as a background process. @@ -819,7 +846,7 @@ def _maybe_start_services( auto: bool, start_services: bool | None, timeout: int = _SERVICE_STARTUP_TIMEOUT_SECONDS, -) -> Literal["ready", "connect_remote"]: +) -> Literal["ready", "connect_remote", "start_local"]: """Start services if requested, restarting if already running. In interactive mode (auto=False), prompts the user if start_services is None. @@ -839,8 +866,10 @@ def _maybe_start_services( already_running = _check_platform_reachable(base_url) if already_running and start_services is not True: - console.print(f"{CHECK} Platform already running at {base_url}\n") - return "ready" + if _is_local_base_url(base_url) or auto: + console.print(f"{CHECK} Platform already running at {base_url}\n") + return "ready" + return _prompt_reachable_remote_connection(base_url) should_start = start_services if should_start is None: @@ -1797,6 +1826,16 @@ def setup_command( inference provider, picks a default model, installs coding agent skills, and optionally deploys a demo agent. + The active config context remembers the Platform URL. When a remote + deployment is already reachable, setup asks whether to continue with it, + start local services instead, or connect to a different remote URL. + + To override the URL for one run only: + nemo --base-url http://localhost:8080 setup + + To persist a different URL: + nemo config set --base-url http://localhost:8080 + Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -1813,6 +1852,7 @@ def setup_command( NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent + nemo --base-url http://localhost:8080 setup """ cli_context: CLIContext = ctx.obj base_url = cli_context.get_base_url() or DEFAULT_BASE_URL @@ -1827,9 +1867,23 @@ def setup_command( if effective_timeout <= 0: raise typer.BadParameter("--ready-timeout must be greater than 0", param_hint="--ready-timeout") try: + configured_base_url = base_url service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) + if service_result == "start_local": + context_name = cli_context.get_sdk_context().context_name + _bootstrap_config_if_missing(DEFAULT_BASE_URL, workspace) + Config.write({"base_url": DEFAULT_BASE_URL}, context_name=context_name) + cli_context.overrides["base_url"] = DEFAULT_BASE_URL + cli_context.reset_sdk_context() + base_url = DEFAULT_BASE_URL + service_result = _maybe_start_services( + base_url, + auto, + start_services=True, + timeout=effective_timeout, + ) if service_result == "connect_remote": - base_url = _prompt_remote_base_url() + base_url = _prompt_remote_base_url(default_url=configured_base_url) _bootstrap_config_if_missing(base_url, workspace) cli_context.reset_sdk_context() workspace = _resolve_setup_workspace(ctx, cli_context, workspace) diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index 94120177e9..51b8eb4b3a 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -74,6 +74,7 @@ from nemo_platform_ext.cli.commands.skills.registry import UnsupportedAgentError from nemo_platform_ext.config.config import Config from nemo_platform_ext.config.models import ( + DEFAULT_BASE_URL, Cluster, ConfigFile, ConfigParams, @@ -628,6 +629,45 @@ def test_unreachable_remote_url_selects_remote_connection_without_local_prompt(s assert result == "connect_remote" mock_prompt.assert_not_called() + def test_reachable_remote_prompts_continue(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice", return_value="continue") as mock_prompt, + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "ready" + mock_prompt.assert_called_once() + assert "remote.example.com" in mock_prompt.call_args.kwargs["message"] + + def test_reachable_remote_start_local(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice", return_value="local"), + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "start_local" + + def test_reachable_remote_change_remote(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice", return_value="change"), + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "connect_remote" + + def test_reachable_remote_auto_skips_prompt(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice") as mock_prompt, + ): + result = _maybe_start_services("https://remote.example.com", auto=True, start_services=None) + + assert result == "ready" + mock_prompt.assert_not_called() + def test_rejects_start_services_for_remote_url(self): with ( patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), @@ -2604,7 +2644,7 @@ def test_remote_choice_connects_before_continuing_setup(self): with ( patch(f"{SETUP_MOD}.is_interactive", return_value=True), patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com") as mock_prompt, patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, patch(f"{SETUP_MOD}._ensure_platform_auth") as mock_auth, patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), @@ -2613,6 +2653,7 @@ def test_remote_choice_connects_before_continuing_setup(self): ): setup_command(ctx) + mock_prompt.assert_called_once_with(default_url="http://localhost:8080") mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "default") mock_auth.assert_called_once_with(cli_context) assert mock_run.call_args.args[3] == "https://remote.example.com" @@ -2681,6 +2722,48 @@ def test_remote_choice_uses_explicit_workspace_flag(self): mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") assert mock_run.call_args.args[2] == "shared-workspace" + def test_start_local_persists_default_url_and_starts_services(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT + cli_context = MagicMock() + cli_context.overrides = {} + cli_context.get_base_url.return_value = "https://remote.example.com" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="https://remote.example.com"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch( + f"{SETUP_MOD}._maybe_start_services", + side_effect=["start_local", "ready"], + ) as mock_start, + patch(f"{SETUP_MOD}.Config.write") as mock_write, + patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx) + + assert mock_start.call_count == 2 + assert mock_start.call_args_list[1] == call( + DEFAULT_BASE_URL, + False, + start_services=True, + timeout=_SERVICE_STARTUP_TIMEOUT_SECONDS, + ) + mock_bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") + mock_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") + assert cli_context.overrides["base_url"] == DEFAULT_BASE_URL + cli_context.reset_sdk_context.assert_called() + assert mock_run.call_args.args[3] == DEFAULT_BASE_URL + # --------------------------------------------------------------------------- # Controller health check diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py index cadad8e139..2e078d4fc2 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/manifest_registry.py @@ -41,6 +41,16 @@ inference provider, picks a default model, installs coding agent skills, and optionally deploys a demo agent. +The active config context remembers the Platform URL. When a remote +deployment is already reachable, setup asks whether to continue with it, +start local services instead, or connect to a different remote URL. + +To override the URL for one run only: + nemo --base-url http://localhost:8080 setup + +To persist a different URL: + nemo config set --base-url http://localhost:8080 + Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -56,7 +66,8 @@ nemo setup --auto --start-services --ready-timeout 360 NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace - nemo setup --no-install-skills --no-deploy-agent""", + nemo setup --no-install-skills --no-deploy-agent + nemo --base-url http://localhost:8080 setup""", name="setup", panel="Setup", kind="command", diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 6a35f60530..46805f8747 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -326,11 +326,12 @@ def _check_platform_reachable_with_retries( return False -def _prompt_remote_base_url() -> str: +def _prompt_remote_base_url(*, default_url: str = "") -> str: """Prompt until the user provides a reachable remote Platform URL.""" while True: base_url = prompt_text( "Enter the remote Platform base URL: ", + default=default_url, validator=non_empty_validator("Base URL"), ).strip() parsed = urlparse(base_url) @@ -736,6 +737,32 @@ def _is_local_base_url(base_url: str) -> bool: return parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"} +def _platform_host_label(base_url: str) -> str: + """Return a human-friendly host label for connection prompts.""" + parsed = urlparse(base_url) + return parsed.hostname or base_url.rstrip("/") + + +def _prompt_reachable_remote_connection(base_url: str) -> Literal["ready", "connect_remote", "start_local"]: + """Ask how to proceed when a configured remote Platform is already reachable.""" + hostname = _platform_host_label(base_url) + action = prompt_choice( + message=f"Platform reachable at {hostname} ({base_url}). What would you like to do?", + options=[ + ("continue", "Continue with this remote Platform"), + ("local", "Start local services instead"), + ("change", "Connect to a different remote URL"), + ], + default="continue", + ) + if action == "continue": + console.print(f"{CHECK} Platform already running at {base_url}\n") + return "ready" + if action == "change": + return "connect_remote" + return "start_local" + + def _start_services_background(base_url: str, data_dir: str | None = None) -> subprocess.Popen: """Launch ``nemo services run`` as a background process. @@ -819,7 +846,7 @@ def _maybe_start_services( auto: bool, start_services: bool | None, timeout: int = _SERVICE_STARTUP_TIMEOUT_SECONDS, -) -> Literal["ready", "connect_remote"]: +) -> Literal["ready", "connect_remote", "start_local"]: """Start services if requested, restarting if already running. In interactive mode (auto=False), prompts the user if start_services is None. @@ -839,8 +866,10 @@ def _maybe_start_services( already_running = _check_platform_reachable(base_url) if already_running and start_services is not True: - console.print(f"{CHECK} Platform already running at {base_url}\n") - return "ready" + if _is_local_base_url(base_url) or auto: + console.print(f"{CHECK} Platform already running at {base_url}\n") + return "ready" + return _prompt_reachable_remote_connection(base_url) should_start = start_services if should_start is None: @@ -1797,6 +1826,16 @@ def setup_command( inference provider, picks a default model, installs coding agent skills, and optionally deploys a demo agent. + The active config context remembers the Platform URL. When a remote + deployment is already reachable, setup asks whether to continue with it, + start local services instead, or connect to a different remote URL. + + To override the URL for one run only: + nemo --base-url http://localhost:8080 setup + + To persist a different URL: + nemo config set --base-url http://localhost:8080 + Requires an interactive terminal (TTY). In non-interactive contexts (CI, piped input), pass --auto to use environment variables instead. @@ -1813,6 +1852,7 @@ def setup_command( NMP_BASE_URL=https://nmp.example.com NMP_ACCESS_TOKEN=... nemo setup --auto --no-start-services nemo setup --workspace my-workspace nemo setup --no-install-skills --no-deploy-agent + nemo --base-url http://localhost:8080 setup """ cli_context: CLIContext = ctx.obj base_url = cli_context.get_base_url() or DEFAULT_BASE_URL @@ -1827,9 +1867,23 @@ def setup_command( if effective_timeout <= 0: raise typer.BadParameter("--ready-timeout must be greater than 0", param_hint="--ready-timeout") try: + configured_base_url = base_url service_result = _maybe_start_services(base_url, auto, start_services, timeout=effective_timeout) + if service_result == "start_local": + context_name = cli_context.get_sdk_context().context_name + _bootstrap_config_if_missing(DEFAULT_BASE_URL, workspace) + Config.write({"base_url": DEFAULT_BASE_URL}, context_name=context_name) + cli_context.overrides["base_url"] = DEFAULT_BASE_URL + cli_context.reset_sdk_context() + base_url = DEFAULT_BASE_URL + service_result = _maybe_start_services( + base_url, + auto, + start_services=True, + timeout=effective_timeout, + ) if service_result == "connect_remote": - base_url = _prompt_remote_base_url() + base_url = _prompt_remote_base_url(default_url=configured_base_url) _bootstrap_config_if_missing(base_url, workspace) cli_context.reset_sdk_context() workspace = _resolve_setup_workspace(ctx, cli_context, workspace) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index 860e55c7ec..f1261ed8a2 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -79,6 +79,7 @@ ConfigParams, Context, ContextDefinition, + DEFAULT_BASE_URL, NoAuthUser, OAuthUser, ) @@ -628,6 +629,45 @@ def test_unreachable_remote_url_selects_remote_connection_without_local_prompt(s assert result == "connect_remote" mock_prompt.assert_not_called() + def test_reachable_remote_prompts_continue(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice", return_value="continue") as mock_prompt, + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "ready" + mock_prompt.assert_called_once() + assert "remote.example.com" in mock_prompt.call_args.kwargs["message"] + + def test_reachable_remote_start_local(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice", return_value="local"), + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "start_local" + + def test_reachable_remote_change_remote(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice", return_value="change"), + ): + result = _maybe_start_services("https://remote.example.com", auto=False, start_services=None) + + assert result == "connect_remote" + + def test_reachable_remote_auto_skips_prompt(self): + with ( + patch(f"{SETUP_MOD}._check_platform_reachable", return_value=True), + patch(f"{SETUP_MOD}.prompt_choice") as mock_prompt, + ): + result = _maybe_start_services("https://remote.example.com", auto=True, start_services=None) + + assert result == "ready" + mock_prompt.assert_not_called() + def test_rejects_start_services_for_remote_url(self): with ( patch(f"{SETUP_MOD}._check_platform_reachable", return_value=False), @@ -2604,7 +2644,7 @@ def test_remote_choice_connects_before_continuing_setup(self): with ( patch(f"{SETUP_MOD}.is_interactive", return_value=True), patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com") as mock_prompt, patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, patch(f"{SETUP_MOD}._ensure_platform_auth") as mock_auth, patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), @@ -2613,6 +2653,7 @@ def test_remote_choice_connects_before_continuing_setup(self): ): setup_command(ctx) + mock_prompt.assert_called_once_with(default_url="http://localhost:8080") mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "default") mock_auth.assert_called_once_with(cli_context) assert mock_run.call_args.args[3] == "https://remote.example.com" @@ -2681,6 +2722,48 @@ def test_remote_choice_uses_explicit_workspace_flag(self): mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") assert mock_run.call_args.args[2] == "shared-workspace" + def test_start_local_persists_default_url_and_starts_services(self): + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = ParameterSource.DEFAULT + cli_context = MagicMock() + cli_context.overrides = {} + cli_context.get_base_url.return_value = "https://remote.example.com" + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url="https://remote.example.com"), + user=NoAuthUser(name="default-user"), + workspace="default", + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + + with ( + patch(f"{SETUP_MOD}.is_interactive", return_value=True), + patch( + f"{SETUP_MOD}._maybe_start_services", + side_effect=["start_local", "ready"], + ) as mock_start, + patch(f"{SETUP_MOD}.Config.write") as mock_write, + patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), + patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, + ): + setup_command(ctx) + + assert mock_start.call_count == 2 + assert mock_start.call_args_list[1] == call( + DEFAULT_BASE_URL, + False, + start_services=True, + timeout=_SERVICE_STARTUP_TIMEOUT_SECONDS, + ) + mock_bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") + mock_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") + assert cli_context.overrides["base_url"] == DEFAULT_BASE_URL + cli_context.reset_sdk_context.assert_called() + assert mock_run.call_args.args[3] == DEFAULT_BASE_URL + # --------------------------------------------------------------------------- # Controller health check From a80bc19d1f85737626c286b1194fbed8c1777210 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Thu, 23 Jul 2026 10:02:30 -0700 Subject: [PATCH 7/8] fix(cli): address setup review nits and CodeQL alert (#934) - Fix CodeQL "incomplete URL substring sanitization" by asserting the full prompt message instead of a hostname substring check. - Use a StrEnum for the reachable-remote connection choices. - Reword the secure-config-write docstring. - Re-vendor SDK copies so lint-sdk-vendored/lint-cli stay in sync. Signed-off-by: Tyler Bray --- .../nemo_platform_ext/cli/commands/setup.py | 21 +++++++++++++------ .../src/nemo_platform_ext/config/config.py | 3 ++- .../tests/cli/commands/test_setup.py | 5 ++++- .../src/nemo_platform/cli/commands/setup.py | 21 +++++++++++++------ .../src/nemo_platform/config/config.py | 3 ++- .../cli/commands/test_setup.py | 7 +++++-- 6 files changed, 43 insertions(+), 17 deletions(-) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index 6c7b33ee3d..5f06f6a4e4 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -16,6 +16,7 @@ import subprocess import time from dataclasses import dataclass +from enum import StrEnum from importlib.resources import files from importlib.resources.abc import Traversable from pathlib import Path @@ -743,22 +744,30 @@ def _platform_host_label(base_url: str) -> str: return parsed.hostname or base_url.rstrip("/") +class _RemoteConnectionChoice(StrEnum): + """Options offered when a configured remote Platform is already reachable.""" + + CONTINUE = "continue" + START_LOCAL = "local" + CHANGE_REMOTE = "change" + + def _prompt_reachable_remote_connection(base_url: str) -> Literal["ready", "connect_remote", "start_local"]: """Ask how to proceed when a configured remote Platform is already reachable.""" hostname = _platform_host_label(base_url) action = prompt_choice( message=f"Platform reachable at {hostname} ({base_url}). What would you like to do?", options=[ - ("continue", "Continue with this remote Platform"), - ("local", "Start local services instead"), - ("change", "Connect to a different remote URL"), + (_RemoteConnectionChoice.CONTINUE, "Continue with this remote Platform"), + (_RemoteConnectionChoice.START_LOCAL, "Start local services instead"), + (_RemoteConnectionChoice.CHANGE_REMOTE, "Connect to a different remote URL"), ], - default="continue", + default=_RemoteConnectionChoice.CONTINUE, ) - if action == "continue": + if action == _RemoteConnectionChoice.CONTINUE: console.print(f"{CHECK} Platform already running at {base_url}\n") return "ready" - if action == "change": + if action == _RemoteConnectionChoice.CHANGE_REMOTE: return "connect_remote" return "start_local" diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py index 086de28a1b..341d41924d 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/config/config.py @@ -41,7 +41,8 @@ def _try_secure_chmod_dir(path: Path) -> None: def _write_secure_yaml(path: Path, config_data: dict) -> None: """Write YAML atomically with owner read/write permissions (600). - File permission failures remain fatal so credentials are never left world-readable. + A failure to lock down the file's permissions raises instead of being + swallowed, so credentials are never left readable by other users. """ mode = stat.S_IRUSR | stat.S_IWUSR # 600 fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index 51b8eb4b3a..38ce54808b 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -638,7 +638,10 @@ def test_reachable_remote_prompts_continue(self): assert result == "ready" mock_prompt.assert_called_once() - assert "remote.example.com" in mock_prompt.call_args.kwargs["message"] + assert ( + mock_prompt.call_args.kwargs["message"] + == "Platform reachable at remote.example.com (https://remote.example.com). What would you like to do?" + ) def test_reachable_remote_start_local(self): with ( diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 46805f8747..d5d6b32ab9 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -16,6 +16,7 @@ import subprocess import time from dataclasses import dataclass +from enum import StrEnum from importlib.resources import files from importlib.resources.abc import Traversable from pathlib import Path @@ -743,22 +744,30 @@ def _platform_host_label(base_url: str) -> str: return parsed.hostname or base_url.rstrip("/") +class _RemoteConnectionChoice(StrEnum): + """Options offered when a configured remote Platform is already reachable.""" + + CONTINUE = "continue" + START_LOCAL = "local" + CHANGE_REMOTE = "change" + + def _prompt_reachable_remote_connection(base_url: str) -> Literal["ready", "connect_remote", "start_local"]: """Ask how to proceed when a configured remote Platform is already reachable.""" hostname = _platform_host_label(base_url) action = prompt_choice( message=f"Platform reachable at {hostname} ({base_url}). What would you like to do?", options=[ - ("continue", "Continue with this remote Platform"), - ("local", "Start local services instead"), - ("change", "Connect to a different remote URL"), + (_RemoteConnectionChoice.CONTINUE, "Continue with this remote Platform"), + (_RemoteConnectionChoice.START_LOCAL, "Start local services instead"), + (_RemoteConnectionChoice.CHANGE_REMOTE, "Connect to a different remote URL"), ], - default="continue", + default=_RemoteConnectionChoice.CONTINUE, ) - if action == "continue": + if action == _RemoteConnectionChoice.CONTINUE: console.print(f"{CHECK} Platform already running at {base_url}\n") return "ready" - if action == "change": + if action == _RemoteConnectionChoice.CHANGE_REMOTE: return "connect_remote" return "start_local" diff --git a/sdk/python/nemo-platform/src/nemo_platform/config/config.py b/sdk/python/nemo-platform/src/nemo_platform/config/config.py index c4e34a40f8..8c5f3e4b20 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/config/config.py +++ b/sdk/python/nemo-platform/src/nemo_platform/config/config.py @@ -41,7 +41,8 @@ def _try_secure_chmod_dir(path: Path) -> None: def _write_secure_yaml(path: Path, config_data: dict) -> None: """Write YAML atomically with owner read/write permissions (600). - File permission failures remain fatal so credentials are never left world-readable. + A failure to lock down the file's permissions raises instead of being + swallowed, so credentials are never left readable by other users. """ mode = stat.S_IRUSR | stat.S_IWUSR # 600 fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index f1261ed8a2..320db20cc2 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -74,12 +74,12 @@ from nemo_platform.cli.commands.skills.registry import UnsupportedAgentError from nemo_platform.config.config import Config from nemo_platform.config.models import ( + DEFAULT_BASE_URL, Cluster, ConfigFile, ConfigParams, Context, ContextDefinition, - DEFAULT_BASE_URL, NoAuthUser, OAuthUser, ) @@ -638,7 +638,10 @@ def test_reachable_remote_prompts_continue(self): assert result == "ready" mock_prompt.assert_called_once() - assert "remote.example.com" in mock_prompt.call_args.kwargs["message"] + assert ( + mock_prompt.call_args.kwargs["message"] + == "Platform reachable at remote.example.com (https://remote.example.com). What would you like to do?" + ) def test_reachable_remote_start_local(self): with ( From 77d28bd0c75bfd71491d9c60e9d2aca03c95e6a2 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Thu, 23 Jul 2026 11:18:47 -0700 Subject: [PATCH 8/8] refactor(cli): extract setup_command test helpers (#934) Collapse the repeated typer Context / CLIContext and with-patch stacks in TestNonTtyEarlyExit and TestSetupCommandRemoteFlow into shared helpers. Signed-off-by: Tyler Bray --- .../tests/cli/commands/test_setup.py | 262 ++++++++---------- .../cli/commands/test_setup.py | 262 ++++++++---------- 2 files changed, 246 insertions(+), 278 deletions(-) diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index 38ce54808b..fa128545f1 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -7,6 +7,9 @@ import logging import sys +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from types import SimpleNamespace from unittest.mock import MagicMock, call, patch import httpx @@ -2572,6 +2575,88 @@ def test_prompt_passes_hint(self): assert "lowercase" in hint.lower() +# --------------------------------------------------------------------------- +# setup_command entry-path helpers +# --------------------------------------------------------------------------- + + +def _make_setup_command_ctx( + *, + base_url: str = "http://localhost:8080", + workspace: str = "default", + workspace_source: ParameterSource = ParameterSource.DEFAULT, +) -> tuple[MagicMock, MagicMock]: + """Build a typer Context + CLIContext pair for invoking ``setup_command``.""" + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = workspace_source + cli_context = MagicMock() + cli_context.overrides = {} + cli_context.get_base_url.return_value = base_url + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url=base_url), + user=NoAuthUser(name="default-user"), + workspace=workspace, + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + return ctx, cli_context + + +@contextmanager +def _patch_setup_command( + *, + interactive: bool = True, + maybe_start_services: object | None = None, + remote_url: str | None = None, + include_config_write: bool = False, + auto_mode: bool = False, +) -> Iterator[SimpleNamespace]: + """Patch the common ``setup_command`` entry path and yield named mocks. + + When *remote_url* is set, also patches the connect-remote branch helpers and + defaults ``_maybe_start_services`` to return ``\"connect_remote\"``. + """ + maybe_start_kwargs: dict[str, object] = {} + if maybe_start_services is None: + if remote_url is not None: + maybe_start_kwargs["return_value"] = "connect_remote" + elif isinstance(maybe_start_services, list): + maybe_start_kwargs["side_effect"] = maybe_start_services + elif isinstance(maybe_start_services, MagicMock): + maybe_start_kwargs["new"] = maybe_start_services + else: + maybe_start_kwargs["return_value"] = maybe_start_services + + with ExitStack() as stack: + mocks = SimpleNamespace( + is_interactive=stack.enter_context(patch(f"{SETUP_MOD}.is_interactive", return_value=interactive)), + maybe_start_services=stack.enter_context(patch(f"{SETUP_MOD}._maybe_start_services", **maybe_start_kwargs)), + check_reachable=stack.enter_context( + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True) + ), + bootstrap=stack.enter_context(patch(f"{SETUP_MOD}._bootstrap_config_if_missing")), + run_interactive=stack.enter_context(patch(f"{SETUP_MOD}._run_interactive_mode")), + prompt_remote=None, + configure_remote=None, + ensure_auth=None, + config_write=None, + run_auto=None, + ) + if remote_url is not None: + mocks.prompt_remote = stack.enter_context( + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value=remote_url) + ) + mocks.configure_remote = stack.enter_context(patch(f"{SETUP_MOD}._configure_remote_connection")) + mocks.ensure_auth = stack.enter_context(patch(f"{SETUP_MOD}._ensure_platform_auth")) + if include_config_write: + mocks.config_write = stack.enter_context(patch(f"{SETUP_MOD}.Config.write")) + if auto_mode: + mocks.run_auto = stack.enter_context(patch(f"{SETUP_MOD}._run_auto_mode")) + yield mocks + + # --------------------------------------------------------------------------- # Non-TTY early exit guard # --------------------------------------------------------------------------- @@ -2580,49 +2665,32 @@ def test_prompt_passes_hint(self): class TestNonTtyEarlyExit: """setup_command must exit(1) when stdin is not a TTY and --auto is not passed.""" - def _invoke(self, *, auto: bool = False): - """Invoke setup_command with a minimal mock context.""" - ctx = MagicMock(spec=typer.Context) - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - ctx.obj = cli_context - setup_command(ctx, auto=auto) - def test_exits_when_non_tty_without_auto(self): + ctx, _cli_context = _make_setup_command_ctx() with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=False), + _patch_setup_command(interactive=False), pytest.raises(typer.Exit) as exc_info, ): - self._invoke(auto=False) + setup_command(ctx, auto=False) assert exc_info.value.exit_code == 1 def test_proceeds_when_non_tty_with_auto(self): - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=False), - patch(f"{SETUP_MOD}._maybe_start_services"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_auto_mode"), - ): - self._invoke(auto=True) + ctx, _cli_context = _make_setup_command_ctx() + with _patch_setup_command(interactive=False, auto_mode=True): + setup_command(ctx, auto=True) def test_proceeds_when_tty_without_auto(self): - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_interactive_mode"), - ): - self._invoke(auto=False) + ctx, _cli_context = _make_setup_command_ctx() + with _patch_setup_command(): + setup_command(ctx, auto=False) def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): + ctx, _cli_context = _make_setup_command_ctx() with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", side_effect=UserCancelled), + _patch_setup_command(maybe_start_services=MagicMock(side_effect=UserCancelled)), pytest.raises(typer.Exit) as exc_info, ): - self._invoke(auto=False) + setup_command(ctx, auto=False) assert exc_info.value.exit_code == 0 assert "Setup cancelled" in capsys.readouterr().err @@ -2630,142 +2698,58 @@ def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): class TestSetupCommandRemoteFlow: def test_remote_choice_connects_before_continuing_setup(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.DEFAULT - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), - user=NoAuthUser(name="default-user"), - workspace="default", - preferences={}, - ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com") as mock_prompt, - patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, - patch(f"{SETUP_MOD}._ensure_platform_auth") as mock_auth, - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + ctx, cli_context = _make_setup_command_ctx() + with _patch_setup_command(remote_url="https://remote.example.com") as mocks: setup_command(ctx) - mock_prompt.assert_called_once_with(default_url="http://localhost:8080") - mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "default") - mock_auth.assert_called_once_with(cli_context) - assert mock_run.call_args.args[3] == "https://remote.example.com" + mocks.prompt_remote.assert_called_once_with(default_url="http://localhost:8080") + mocks.configure_remote.assert_called_once_with(cli_context, "https://remote.example.com", "default") + mocks.ensure_auth.assert_called_once_with(cli_context) + assert mocks.run_interactive.call_args.args[3] == "https://remote.example.com" def test_remote_choice_preserves_active_workspace_when_flag_omitted(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.DEFAULT - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), - user=NoAuthUser(name="default-user"), - workspace="team-a", - preferences={}, - ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), - patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, - patch(f"{SETUP_MOD}._ensure_platform_auth"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + ctx, cli_context = _make_setup_command_ctx(workspace="team-a") + with _patch_setup_command(remote_url="https://remote.example.com") as mocks: setup_command(ctx) - mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") - assert mock_bootstrap.call_args_list == [ + mocks.configure_remote.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") + assert mocks.bootstrap.call_args_list == [ call("https://remote.example.com", "default"), call("https://remote.example.com", "team-a"), ] - assert mock_run.call_args.args[2] == "team-a" + assert mocks.run_interactive.call_args.args[2] == "team-a" def test_remote_choice_uses_explicit_workspace_flag(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.COMMANDLINE - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), - user=NoAuthUser(name="default-user"), + ctx, cli_context = _make_setup_command_ctx( workspace="team-a", - preferences={}, + workspace_source=ParameterSource.COMMANDLINE, ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), - patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, - patch(f"{SETUP_MOD}._ensure_platform_auth"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + with _patch_setup_command(remote_url="https://remote.example.com") as mocks: setup_command(ctx, workspace="shared-workspace") - mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") - assert mock_run.call_args.args[2] == "shared-workspace" + mocks.configure_remote.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") + assert mocks.run_interactive.call_args.args[2] == "shared-workspace" def test_start_local_persists_default_url_and_starts_services(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.DEFAULT - cli_context = MagicMock() - cli_context.overrides = {} - cli_context.get_base_url.return_value = "https://remote.example.com" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="https://remote.example.com"), - user=NoAuthUser(name="default-user"), - workspace="default", - preferences={}, - ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch( - f"{SETUP_MOD}._maybe_start_services", - side_effect=["start_local", "ready"], - ) as mock_start, - patch(f"{SETUP_MOD}.Config.write") as mock_write, - patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + ctx, cli_context = _make_setup_command_ctx(base_url="https://remote.example.com") + with _patch_setup_command( + maybe_start_services=["start_local", "ready"], + include_config_write=True, + ) as mocks: setup_command(ctx) - assert mock_start.call_count == 2 - assert mock_start.call_args_list[1] == call( + assert mocks.maybe_start_services.call_count == 2 + assert mocks.maybe_start_services.call_args_list[1] == call( DEFAULT_BASE_URL, False, start_services=True, timeout=_SERVICE_STARTUP_TIMEOUT_SECONDS, ) - mock_bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") - mock_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") + mocks.bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") + mocks.config_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") assert cli_context.overrides["base_url"] == DEFAULT_BASE_URL cli_context.reset_sdk_context.assert_called() - assert mock_run.call_args.args[3] == DEFAULT_BASE_URL + assert mocks.run_interactive.call_args.args[3] == DEFAULT_BASE_URL # --------------------------------------------------------------------------- diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index 320db20cc2..038bdba8fc 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -7,6 +7,9 @@ import logging import sys +from collections.abc import Iterator +from contextlib import ExitStack, contextmanager +from types import SimpleNamespace from unittest.mock import MagicMock, call, patch import httpx @@ -2572,6 +2575,88 @@ def test_prompt_passes_hint(self): assert "lowercase" in hint.lower() +# --------------------------------------------------------------------------- +# setup_command entry-path helpers +# --------------------------------------------------------------------------- + + +def _make_setup_command_ctx( + *, + base_url: str = "http://localhost:8080", + workspace: str = "default", + workspace_source: ParameterSource = ParameterSource.DEFAULT, +) -> tuple[MagicMock, MagicMock]: + """Build a typer Context + CLIContext pair for invoking ``setup_command``.""" + ctx = MagicMock(spec=typer.Context) + ctx.get_parameter_source.return_value = workspace_source + cli_context = MagicMock() + cli_context.overrides = {} + cli_context.get_base_url.return_value = base_url + cli_context.get_sdk_context.return_value = Context( + context_name="default", + cluster=Cluster(name="default-cluster", base_url=base_url), + user=NoAuthUser(name="default-user"), + workspace=workspace, + preferences={}, + ) + cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() + ctx.obj = cli_context + return ctx, cli_context + + +@contextmanager +def _patch_setup_command( + *, + interactive: bool = True, + maybe_start_services: object | None = None, + remote_url: str | None = None, + include_config_write: bool = False, + auto_mode: bool = False, +) -> Iterator[SimpleNamespace]: + """Patch the common ``setup_command`` entry path and yield named mocks. + + When *remote_url* is set, also patches the connect-remote branch helpers and + defaults ``_maybe_start_services`` to return ``\"connect_remote\"``. + """ + maybe_start_kwargs: dict[str, object] = {} + if maybe_start_services is None: + if remote_url is not None: + maybe_start_kwargs["return_value"] = "connect_remote" + elif isinstance(maybe_start_services, list): + maybe_start_kwargs["side_effect"] = maybe_start_services + elif isinstance(maybe_start_services, MagicMock): + maybe_start_kwargs["new"] = maybe_start_services + else: + maybe_start_kwargs["return_value"] = maybe_start_services + + with ExitStack() as stack: + mocks = SimpleNamespace( + is_interactive=stack.enter_context(patch(f"{SETUP_MOD}.is_interactive", return_value=interactive)), + maybe_start_services=stack.enter_context(patch(f"{SETUP_MOD}._maybe_start_services", **maybe_start_kwargs)), + check_reachable=stack.enter_context( + patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True) + ), + bootstrap=stack.enter_context(patch(f"{SETUP_MOD}._bootstrap_config_if_missing")), + run_interactive=stack.enter_context(patch(f"{SETUP_MOD}._run_interactive_mode")), + prompt_remote=None, + configure_remote=None, + ensure_auth=None, + config_write=None, + run_auto=None, + ) + if remote_url is not None: + mocks.prompt_remote = stack.enter_context( + patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value=remote_url) + ) + mocks.configure_remote = stack.enter_context(patch(f"{SETUP_MOD}._configure_remote_connection")) + mocks.ensure_auth = stack.enter_context(patch(f"{SETUP_MOD}._ensure_platform_auth")) + if include_config_write: + mocks.config_write = stack.enter_context(patch(f"{SETUP_MOD}.Config.write")) + if auto_mode: + mocks.run_auto = stack.enter_context(patch(f"{SETUP_MOD}._run_auto_mode")) + yield mocks + + # --------------------------------------------------------------------------- # Non-TTY early exit guard # --------------------------------------------------------------------------- @@ -2580,49 +2665,32 @@ def test_prompt_passes_hint(self): class TestNonTtyEarlyExit: """setup_command must exit(1) when stdin is not a TTY and --auto is not passed.""" - def _invoke(self, *, auto: bool = False): - """Invoke setup_command with a minimal mock context.""" - ctx = MagicMock(spec=typer.Context) - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - ctx.obj = cli_context - setup_command(ctx, auto=auto) - def test_exits_when_non_tty_without_auto(self): + ctx, _cli_context = _make_setup_command_ctx() with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=False), + _patch_setup_command(interactive=False), pytest.raises(typer.Exit) as exc_info, ): - self._invoke(auto=False) + setup_command(ctx, auto=False) assert exc_info.value.exit_code == 1 def test_proceeds_when_non_tty_with_auto(self): - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=False), - patch(f"{SETUP_MOD}._maybe_start_services"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_auto_mode"), - ): - self._invoke(auto=True) + ctx, _cli_context = _make_setup_command_ctx() + with _patch_setup_command(interactive=False, auto_mode=True): + setup_command(ctx, auto=True) def test_proceeds_when_tty_without_auto(self): - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_interactive_mode"), - ): - self._invoke(auto=False) + ctx, _cli_context = _make_setup_command_ctx() + with _patch_setup_command(): + setup_command(ctx, auto=False) def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): + ctx, _cli_context = _make_setup_command_ctx() with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", side_effect=UserCancelled), + _patch_setup_command(maybe_start_services=MagicMock(side_effect=UserCancelled)), pytest.raises(typer.Exit) as exc_info, ): - self._invoke(auto=False) + setup_command(ctx, auto=False) assert exc_info.value.exit_code == 0 assert "Setup cancelled" in capsys.readouterr().err @@ -2630,142 +2698,58 @@ def test_cancelling_initial_connection_prompt_exits_cleanly(self, capsys): class TestSetupCommandRemoteFlow: def test_remote_choice_connects_before_continuing_setup(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.DEFAULT - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), - user=NoAuthUser(name="default-user"), - workspace="default", - preferences={}, - ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com") as mock_prompt, - patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, - patch(f"{SETUP_MOD}._ensure_platform_auth") as mock_auth, - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + ctx, cli_context = _make_setup_command_ctx() + with _patch_setup_command(remote_url="https://remote.example.com") as mocks: setup_command(ctx) - mock_prompt.assert_called_once_with(default_url="http://localhost:8080") - mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "default") - mock_auth.assert_called_once_with(cli_context) - assert mock_run.call_args.args[3] == "https://remote.example.com" + mocks.prompt_remote.assert_called_once_with(default_url="http://localhost:8080") + mocks.configure_remote.assert_called_once_with(cli_context, "https://remote.example.com", "default") + mocks.ensure_auth.assert_called_once_with(cli_context) + assert mocks.run_interactive.call_args.args[3] == "https://remote.example.com" def test_remote_choice_preserves_active_workspace_when_flag_omitted(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.DEFAULT - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), - user=NoAuthUser(name="default-user"), - workspace="team-a", - preferences={}, - ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), - patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, - patch(f"{SETUP_MOD}._ensure_platform_auth"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + ctx, cli_context = _make_setup_command_ctx(workspace="team-a") + with _patch_setup_command(remote_url="https://remote.example.com") as mocks: setup_command(ctx) - mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") - assert mock_bootstrap.call_args_list == [ + mocks.configure_remote.assert_called_once_with(cli_context, "https://remote.example.com", "team-a") + assert mocks.bootstrap.call_args_list == [ call("https://remote.example.com", "default"), call("https://remote.example.com", "team-a"), ] - assert mock_run.call_args.args[2] == "team-a" + assert mocks.run_interactive.call_args.args[2] == "team-a" def test_remote_choice_uses_explicit_workspace_flag(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.COMMANDLINE - cli_context = MagicMock() - cli_context.get_base_url.return_value = "http://localhost:8080" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="http://localhost:8080"), - user=NoAuthUser(name="default-user"), + ctx, cli_context = _make_setup_command_ctx( workspace="team-a", - preferences={}, + workspace_source=ParameterSource.COMMANDLINE, ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch(f"{SETUP_MOD}._maybe_start_services", return_value="connect_remote"), - patch(f"{SETUP_MOD}._prompt_remote_base_url", return_value="https://remote.example.com"), - patch(f"{SETUP_MOD}._configure_remote_connection") as mock_configure, - patch(f"{SETUP_MOD}._ensure_platform_auth"), - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._bootstrap_config_if_missing"), - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + with _patch_setup_command(remote_url="https://remote.example.com") as mocks: setup_command(ctx, workspace="shared-workspace") - mock_configure.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") - assert mock_run.call_args.args[2] == "shared-workspace" + mocks.configure_remote.assert_called_once_with(cli_context, "https://remote.example.com", "shared-workspace") + assert mocks.run_interactive.call_args.args[2] == "shared-workspace" def test_start_local_persists_default_url_and_starts_services(self): - ctx = MagicMock(spec=typer.Context) - ctx.get_parameter_source.return_value = ParameterSource.DEFAULT - cli_context = MagicMock() - cli_context.overrides = {} - cli_context.get_base_url.return_value = "https://remote.example.com" - cli_context.get_sdk_context.return_value = Context( - context_name="default", - cluster=Cluster(name="default-cluster", base_url="https://remote.example.com"), - user=NoAuthUser(name="default-user"), - workspace="default", - preferences={}, - ) - cli_context.get_client.return_value.workspaces.retrieve.return_value = MagicMock() - ctx.obj = cli_context - - with ( - patch(f"{SETUP_MOD}.is_interactive", return_value=True), - patch( - f"{SETUP_MOD}._maybe_start_services", - side_effect=["start_local", "ready"], - ) as mock_start, - patch(f"{SETUP_MOD}.Config.write") as mock_write, - patch(f"{SETUP_MOD}._bootstrap_config_if_missing") as mock_bootstrap, - patch(f"{SETUP_MOD}._check_platform_reachable_with_retries", return_value=True), - patch(f"{SETUP_MOD}._run_interactive_mode") as mock_run, - ): + ctx, cli_context = _make_setup_command_ctx(base_url="https://remote.example.com") + with _patch_setup_command( + maybe_start_services=["start_local", "ready"], + include_config_write=True, + ) as mocks: setup_command(ctx) - assert mock_start.call_count == 2 - assert mock_start.call_args_list[1] == call( + assert mocks.maybe_start_services.call_count == 2 + assert mocks.maybe_start_services.call_args_list[1] == call( DEFAULT_BASE_URL, False, start_services=True, timeout=_SERVICE_STARTUP_TIMEOUT_SECONDS, ) - mock_bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") - mock_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") + mocks.bootstrap.assert_any_call(DEFAULT_BASE_URL, "default") + mocks.config_write.assert_called_once_with({"base_url": DEFAULT_BASE_URL}, context_name="default") assert cli_context.overrides["base_url"] == DEFAULT_BASE_URL cli_context.reset_sdk_context.assert_called() - assert mock_run.call_args.args[3] == DEFAULT_BASE_URL + assert mocks.run_interactive.call_args.args[3] == DEFAULT_BASE_URL # ---------------------------------------------------------------------------