From 4ff6b1303a861b7d5307a75a2d215de878a1aae8 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Mon, 6 Jul 2026 15:32:09 -0700 Subject: [PATCH 1/2] fix(agents): Honor shared CLI context for base URL and auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `nemo agents` command group resolved its own base URL (only `--base-url`/`NEMO_BASE_URL`, else localhost) and attached no auth token, so it silently targeted the wrong platform and was rejected (401/403) on any secured cluster — breaking the local-to-remote deploy flow (AIRCORE-885). Add a shared `cli_context` module (resolve_base_url, resolve_context_headers) that reads the same CLIContext the rest of the CLI uses: - Base URL precedence: --base-url/NEMO_BASE_URL > `nemo config`/NMP_BASE_URL > localhost. The resolved target is echoed to stderr ("Targeting ") so mis-pointed commands are visible; stdout stays clean JSON. - Auth: attach the `nemo auth login` bearer token to the platform httpx calls (_api_request, gateway invoke) and to the usage-report SDK client. Applies to the platform commands and to `nemo agents usage show` (which builds its own SDK client). The module lives outside cli.py to avoid the cli.py <-> usage/cli.py import cycle. Tests: base-URL precedence + auth attachment (new test_cli_context_resolution.py), usage SDK client build with context base URL + auth, and stdout/stderr stream separation for JSON output. Signed-off-by: Tyler Bray --- .../nemo-agents/src/nemo_agents_plugin/cli.py | 160 +++++++++++-- .../src/nemo_agents_plugin/cli_context.py | 95 ++++++++ .../src/nemo_agents_plugin/usage/cli.py | 36 ++- .../tests/unit/test_cli_context_resolution.py | 222 ++++++++++++++++++ .../tests/unit/test_cli_list_output.py | 6 +- .../nemo-agents/tests/unit/usage/test_cli.py | 72 +++++- 6 files changed, 561 insertions(+), 30 deletions(-) create mode 100644 plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py create mode 100644 plugins/nemo-agents/tests/unit/test_cli_context_resolution.py diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py index f5a2cd5be6..eac620a265 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py @@ -45,6 +45,15 @@ import httpx import typer import yaml +from nemo_agents_plugin.cli_context import ( + DEFAULT_BASE_URL as _DEFAULT_BASE_URL, +) +from nemo_agents_plugin.cli_context import ( + resolve_base_url as _resolve_base_url, +) +from nemo_agents_plugin.cli_context import ( + resolve_context_headers as _resolve_context_headers, +) from nemo_agents_plugin.leaderboard.cli import register_leaderboard_commands from nemo_agents_plugin.usage.cli import register_usage_commands from nemo_platform.cli.core.formatters import Column, format_output @@ -54,7 +63,6 @@ logger = logging.getLogger(__name__) -_DEFAULT_BASE_URL = "http://localhost:8080" _DEFAULT_WORKSPACE = "default" _LIST_OUTPUT_FORMAT = Literal["table", "json", "yaml", "csv", "markdown", "raw"] _AGENT_LIST_COLUMNS = [ @@ -145,7 +153,16 @@ def invoke( help="Name of a specific deployment to invoke (platform required).", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), timeout: float = typer.Option( 300, "--timeout", @@ -160,6 +177,7 @@ def invoke( ), ) -> None: """Invoke an agent — locally (with --agent-config) or via the platform (with --agent or --agent-deployment).""" + base_url = _resolve_base_url(base_url) if agent_config: _local_invoke(agent_config, input, input_file, workspace=workspace, base_url=base_url) elif agent or agent_deployment: @@ -628,9 +646,19 @@ def create( ), description: str = typer.Option("", "--description"), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), ) -> None: """Register an agent on the platform.""" + base_url = _resolve_base_url(base_url) from nemo_agents_plugin.utils import inject_default_model config_dict = _load_yaml(agent_config) @@ -653,7 +681,16 @@ def create( def list_agents( ctx: typer.Context, workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), output_format: Optional[_LIST_OUTPUT_FORMAT] = typer.Option( None, "--format", @@ -671,6 +708,7 @@ def list_agents( ), ) -> None: """List agents on the platform.""" + base_url = _resolve_base_url(base_url) resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/agents") _print_list_response( ctx, @@ -684,9 +722,19 @@ def list_agents( def get( name: str = typer.Argument(..., help="Agent name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), ) -> None: """Get an agent by name.""" + base_url = _resolve_base_url(base_url) resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/agents/{name}") typer.echo(json.dumps(resp, indent=2)) @@ -694,10 +742,20 @@ def get( def delete( name: str = typer.Argument(..., help="Agent name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."), ) -> None: """Delete an agent from the platform.""" + base_url = _resolve_base_url(base_url) if not yes: typer.confirm(f"Delete agent '{name}'?", abort=True) _api_request("DELETE", base_url, f"/apis/agents/v2/workspaces/{workspace}/agents/{name}") @@ -725,7 +783,16 @@ def deploy( help="Maximum seconds to wait for a terminal status (only with --wait).", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), ) -> None: """Deploy an agent on the platform. @@ -736,6 +803,7 @@ def deploy( scripted pipelines that prefer to poll separately via ``nemo agents deployments wait``. """ + base_url = _resolve_base_url(base_url) payload: dict = {"agent": agent} if name: payload["name"] = name @@ -792,7 +860,16 @@ def logs( help="Print only the absolute log file path and exit (useful for scripting).", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), ) -> None: """Show logs for an agent deployment. @@ -816,6 +893,7 @@ def logs( raise typer.Exit(code=1) if agent and not name: + base_url = _resolve_base_url(base_url) candidates = [ d for d in _unwrap_list( @@ -859,10 +937,20 @@ def undeploy( None, "--agent", "--all", "-a", help="Remove all deployments for this agent." ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."), ) -> None: """Stop and remove a deployment (or all deployments for an agent).""" + base_url = _resolve_base_url(base_url) if name: if not yes: typer.confirm(f"Undeploy '{name}'?", abort=True) @@ -888,7 +976,16 @@ def undeploy( def deployments_list( ctx: typer.Context, workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), output_format: Optional[_LIST_OUTPUT_FORMAT] = typer.Option( None, "--format", @@ -906,6 +1003,7 @@ def deployments_list( ), ) -> None: """List deployments.""" + base_url = _resolve_base_url(base_url) resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/deployments") _print_list_response( ctx, @@ -919,9 +1017,19 @@ def deployments_list( def deployments_get( name: str = typer.Argument(..., help="Deployment name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), ) -> None: """Get a deployment by name.""" + base_url = _resolve_base_url(base_url) resp = _api_request("GET", base_url, f"/apis/agents/v2/workspaces/{workspace}/deployments/{name}") typer.echo(json.dumps(resp, indent=2)) @@ -929,10 +1037,20 @@ def deployments_get( def deployments_delete( name: str = typer.Argument(..., help="Deployment name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."), ) -> None: """Delete a deployment by name.""" + base_url = _resolve_base_url(base_url) if not yes: typer.confirm(f"Delete deployment '{name}'?", abort=True) _api_request("DELETE", base_url, f"/apis/agents/v2/workspaces/{workspace}/deployments/{name}") @@ -950,7 +1068,16 @@ def deployments_wait( timeout: int = typer.Option(300, "--timeout", "-t", help="Maximum seconds to wait."), interval: float = typer.Option(2.0, "--interval", help="Poll interval in seconds."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option(_DEFAULT_BASE_URL, "--base-url", envvar="NEMO_BASE_URL"), + base_url: Optional[str] = typer.Option( + None, + "--base-url", + envvar="NEMO_BASE_URL", + help=( + "Platform base URL. Defaults to the base URL from the shared CLI " + "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " + "http://localhost:8080." + ), + ), ) -> None: """Wait for a deployment to reach 'running' or 'failed' status. @@ -960,6 +1087,7 @@ def deployments_wait( Provide either a deployment name directly or --agent to resolve the latest active deployment for that agent automatically. """ + base_url = _resolve_base_url(base_url) if not name and not agent: typer.echo("Error: provide a deployment name or --agent.", err=True) raise typer.Exit(code=1) @@ -1200,13 +1328,14 @@ def _platform_invoke( path = f"/apis/agents/v2/workspaces/{workspace}/deployments/{deployment}/-/v1/chat/completions" url = base_url.rstrip("/") + path + headers = _resolve_context_headers() target_label = agent or deployment for query in queries: payload = {"messages": [{"role": "user", "content": query}], "stream": False} try: with request_progress(f"Waiting for agent '{target_label}'...", disabled=no_progress): with httpx.Client(timeout=timeout) as client: - resp = client.post(url, json=payload) + resp = client.post(url, json=payload, headers=headers or None) resp.raise_for_status() body = resp.json() typer.echo(json.dumps(body, indent=2)) @@ -1299,6 +1428,9 @@ def _api_request(method: str, base_url: str, path: str, *, json_body: dict[str, request_kwargs: dict[str, Any] = {} if json_body is not None: request_kwargs["json"] = json_body + headers = _resolve_context_headers() + if headers: + request_kwargs["headers"] = headers try: with httpx.Client(timeout=30) as client: resp = client.request(method, url, **request_kwargs) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py b/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py new file mode 100644 index 0000000000..800104ddcd --- /dev/null +++ b/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared CLI-context resolution for the ``nemo agents`` command group. + +The agents plugin makes platform calls from two places — the platform +commands in :mod:`nemo_agents_plugin.cli` (raw ``httpx``) and +``nemo agents usage show`` in :mod:`nemo_agents_plugin.usage.cli` (the +NeMoPlatform SDK client). Both must resolve the platform base URL and the +auth token the same way every other ``nemo`` command does: through the +shared CLI context object stored on ``typer.Context.obj``. + +These helpers read the *ambient* Click context so callers deep in a command's +call stack can resolve configuration without threading the context object +through every function signature. They live in their own module (rather than +in ``cli.py``) because ``cli.py`` imports the usage CLI at module load, so a +back-import would create a cycle. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import click + +logger = logging.getLogger(__name__) + +DEFAULT_BASE_URL = "http://localhost:8080" + + +def current_cli_state() -> Any: + """Return the shared CLI context object (``typer.Context.obj``) if present. + + Returns ``None`` when the plugin is exercised outside a Click invocation + (e.g. a direct unit test), so callers fall back to their own defaults. + """ + ctx = click.get_current_context(silent=True) + return ctx.obj if ctx is not None else None + + +def base_url_from_context() -> str | None: + """Return the base URL configured in the shared CLI context, if any.""" + state = current_cli_state() + if state is None or not hasattr(state, "get_base_url"): + return None + try: + return state.get_base_url(default=None) + except Exception: + logger.debug("Failed to resolve base URL from CLI context", exc_info=True) + return None + + +def resolve_base_url(base_url: str | None) -> str: + """Resolve the platform base URL and announce the target on stderr. + + Precedence: + 1. Explicit ``--base-url`` / ``NEMO_BASE_URL`` on the command. + 2. The shared CLI context — ``nemo config set --base-url`` and the + ``NMP_BASE_URL`` env var — so ``nemo agents`` targets the same + platform as every other ``nemo`` command. + 3. The built-in localhost default. + + The resolved target is echoed to stderr (never stdout, so piped/JSON + output stays clean) so a mis-pointed command is visible instead of + silently hitting the wrong platform. + """ + resolved = base_url or base_url_from_context() or DEFAULT_BASE_URL + click.echo(f"Targeting {resolved}", err=True) + return resolved + + +def resolve_context_headers() -> dict[str, str]: + """Return auth (and other) default headers from the shared CLI context. + + Mirrors ``nemo_platform_plugin.commands._resolve_submit_auth_headers``: + reads the SDK client config off the shared context so ``nemo agents`` + attaches the same ``Authorization: Bearer`` token as the rest of the CLI + (i.e. the token established by ``nemo auth login``). Returns an empty + mapping when no context or token is available — leaving requests + unauthenticated exactly as before, so local unauthenticated dev keeps + working. + """ + state = current_cli_state() + if state is None or not hasattr(state, "get_sdk_context"): + return {} + try: + client_config = state.get_sdk_context().user.get_client_config() + except Exception: + logger.debug("Failed to resolve auth headers from CLI context", exc_info=True) + return {} + headers = client_config.get("default_headers") if isinstance(client_config, dict) else None + if isinstance(headers, dict): + return {str(key): str(value) for key, value in headers.items()} + return {} diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py b/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py index e33920a2d4..eebccf123c 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py @@ -19,6 +19,7 @@ from typing import Optional import typer +from nemo_agents_plugin.cli_context import resolve_base_url, resolve_context_headers from nemo_agents_plugin.usage import compute, render from nemo_agents_plugin.usage import parser as parser_module from nemo_agents_plugin.usage.models import ( @@ -33,7 +34,6 @@ logger = logging.getLogger(__name__) -_DEFAULT_BASE_URL = "http://localhost:8080" _DEFAULT_WORKSPACE = "default" @@ -73,10 +73,15 @@ def show_cmd( "public number — leave unset and compute_units stays null.", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: str = typer.Option( - _DEFAULT_BASE_URL, + base_url: Optional[str] = typer.Option( + None, "--base-url", envvar="NEMO_BASE_URL", + help=( + "Platform base URL for fileset downloads. Defaults to the base URL " + "from the shared CLI config (`nemo config set --base-url`) / " + "NMP_BASE_URL env var, then http://localhost:8080." + ), ), ) -> None: """Show a usage report for *ref*.""" @@ -98,7 +103,7 @@ def _show( *, total_params: float | None, workspace: str, - base_url: str, + base_url: str | None, ) -> None: try: report = _resolve_and_score( @@ -118,7 +123,7 @@ def _resolve_and_score( *, total_params: float | None, workspace: str, - base_url: str, + base_url: str | None, ) -> UsageReport | BatchUsageReport: """End-to-end pipeline: source → parse → (rewrite if fileset) → score. @@ -146,7 +151,9 @@ def _resolve_and_score( # Path-shaped but missing — clearer error than a fileset 404. raise UsageSourceError(f"local path does not exist: {candidate}") - sdk = _build_sdk(base_url=base_url) + # Only fileset refs contact the platform, so resolve/announce the target + # (and attach auth) here rather than for purely-local reads above. + sdk = _build_sdk(base_url=resolve_base_url(base_url)) with fileset_path(FilesetRef(ref), sdk=sdk, workspace=workspace) as path: report = parser_module.parse_path(path) report = _rewrite_source_dirs(report, original_ref=ref, staged_root=path) @@ -193,13 +200,20 @@ def _rewrite_task_source( def _build_sdk(*, base_url: str) -> NeMoPlatform: - """Construct a NeMoPlatform SDK client. + """Construct a NeMoPlatform SDK client for fileset downloads. + + *base_url* is the value already resolved by ``resolve_base_url`` (flag / + ``NEMO_BASE_URL`` > shared CLI config / ``NMP_BASE_URL`` > localhost), so + don't re-read the env here — that would invert precedence. - *base_url* is the already-resolved value from the Typer option, which - handles env-var fallback via ``envvar="NEMO_BASE_URL"``. Don't re-read - the env here — that would invert precedence and let a stale env - variable beat an explicit ``--base-url`` flag. + Attaches the CLI auth token from the shared context (the same + ``Authorization: Bearer`` header the rest of the CLI sends) so fileset + downloads succeed against a secured cluster. Falls back to an + unauthenticated client when no token is configured. """ + headers = resolve_context_headers() + if headers: + return NeMoPlatform(base_url=base_url, default_headers=headers) return NeMoPlatform(base_url=base_url) diff --git a/plugins/nemo-agents/tests/unit/test_cli_context_resolution.py b/plugins/nemo-agents/tests/unit/test_cli_context_resolution.py new file mode 100644 index 0000000000..4b026ef120 --- /dev/null +++ b/plugins/nemo-agents/tests/unit/test_cli_context_resolution.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI tests for shared-context base-URL resolution and auth-token attachment. + +These pin the two behaviours that make ``nemo agents`` usable against a +remote, secured platform: + +- **Base URL** resolves through the shared CLI context the rest of the CLI + uses (``nemo config set --base-url`` / ``NMP_BASE_URL``), with an explicit + ``--base-url`` / ``NEMO_BASE_URL`` still taking precedence, and the resolved + target echoed to stderr so a mis-pointed command is visible instead of + silently hitting localhost. +- **Auth** headers from the shared context (the ``Authorization: Bearer`` + token behind ``nemo auth login``) are attached to every platform HTTP call, + so agents commands are not rejected 401/403 on a secured cluster. +""" + +from __future__ import annotations + +from contextlib import AbstractContextManager +from typing import Any +from unittest.mock import patch + +import httpx +from nemo_agents_plugin.cli import AgentsCLI +from typer.testing import CliRunner + + +def _install_mock_transport(handler) -> AbstractContextManager[Any]: + """Patch ``httpx.Client`` in the CLI module to use a ``MockTransport``.""" + transport = httpx.MockTransport(handler) + real_client = httpx.Client + + def _factory(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + return patch("nemo_agents_plugin.cli.httpx.Client", _factory) + + +def _capturing(captured: list[httpx.Request], *, json_body: Any = None): + """Return a handler that records every request and replies 200.""" + + def handler(req: httpx.Request) -> httpx.Response: + captured.append(req) + return httpx.Response(200, json=json_body if json_body is not None else {"data": []}) + + return handler + + +class _FakeUser: + def __init__(self, token: str | None) -> None: + self._token = token + + def get_client_config(self) -> dict[str, object]: + if self._token is None: + return {} + return {"default_headers": {"Authorization": f"Bearer {self._token}"}} + + +class _FakeCluster: + def __init__(self, base_url: str) -> None: + self.base_url = base_url + + +class _FakeSDKContext: + def __init__(self, base_url: str, token: str | None) -> None: + self.user = _FakeUser(token) + self.cluster = _FakeCluster(base_url) + + +class _FakeCLIContext: + """Minimal stand-in for ``CLIContext`` (typer.Context.obj).""" + + def __init__(self, base_url: str = "http://config-host:9999", token: str | None = "cfg-token") -> None: + self._sdk = _FakeSDKContext(base_url, token) + + def get_sdk_context(self) -> _FakeSDKContext: + return self._sdk + + def get_base_url(self, default: str | None = None) -> str | None: + return str(self._sdk.cluster.base_url) + + +# --------------------------------------------------------------------------- +# Base URL resolution +# --------------------------------------------------------------------------- + + +def test_base_url_flag_overrides_configured_context() -> None: + """An explicit ``--base-url`` wins over the configured context base URL.""" + captured: list[httpx.Request] = [] + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing(captured)): + result = CliRunner().invoke( + app, + ["list", "--base-url", "http://flag-host:1111"], + obj=_FakeCLIContext(base_url="http://config-host:9999"), + ) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert captured, "expected a request to be issued" + assert captured[0].url.host == "flag-host" + assert captured[0].url.port == 1111 + + +def test_base_url_falls_back_to_configured_context() -> None: + """With no flag/env, agents commands target the configured context base URL. + + This is the P0 regression: previously agents ignored the shared config + and silently hit localhost:8080. + """ + captured: list[httpx.Request] = [] + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing(captured)): + result = CliRunner().invoke( + app, + ["list"], + obj=_FakeCLIContext(base_url="http://config-host:9999"), + ) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert captured[0].url.host == "config-host" + assert captured[0].url.port == 9999 + + +def test_base_url_env_overrides_configured_context() -> None: + """``NEMO_BASE_URL`` (command-level env) still takes precedence over config.""" + captured: list[httpx.Request] = [] + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing(captured)): + result = CliRunner().invoke( + app, + ["list"], + obj=_FakeCLIContext(base_url="http://config-host:9999"), + env={"NEMO_BASE_URL": "http://env-host:2222"}, + ) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert captured[0].url.host == "env-host" + assert captured[0].url.port == 2222 + + +def test_base_url_defaults_to_localhost_without_context() -> None: + """Backwards compatibility: no context and no flag -> localhost:8080.""" + captured: list[httpx.Request] = [] + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing(captured)): + result = CliRunner().invoke(app, ["list"]) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert captured[0].url.host == "localhost" + assert captured[0].url.port == 8080 + + +def test_resolved_target_is_echoed_to_stderr_only() -> None: + """The resolved target is announced on stderr, keeping stdout clean for pipes.""" + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing([])): + result = CliRunner().invoke( + app, + ["list", "--base-url", "http://flag-host:1234", "-o", "json"], + obj=_FakeCLIContext(), + ) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert "Targeting http://flag-host:1234" in (result.stderr or "") + assert "Targeting" not in result.stdout + + +# --------------------------------------------------------------------------- +# Auth token attachment +# --------------------------------------------------------------------------- + + +def test_auth_header_attached_from_context() -> None: + """The bearer token from the shared context is attached to platform calls.""" + captured: list[httpx.Request] = [] + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing(captured)): + result = CliRunner().invoke( + app, + ["list", "--base-url", "http://h:1"], + obj=_FakeCLIContext(token="secret-token"), + ) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert captured[0].headers.get("authorization") == "Bearer secret-token" + + +def test_no_auth_header_without_context() -> None: + """No context -> no auth header (unauthenticated local dev keeps working).""" + captured: list[httpx.Request] = [] + app = AgentsCLI().get_cli() + with _install_mock_transport(_capturing(captured)): + result = CliRunner().invoke(app, ["list", "--base-url", "http://h:1"]) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert "authorization" not in captured[0].headers + + +def test_platform_invoke_attaches_auth_and_targets_context_base_url() -> None: + """``invoke --agent`` routes through the gateway with the context token+URL.""" + captured: list[httpx.Request] = [] + + def handler(req: httpx.Request) -> httpx.Response: + captured.append(req) + return httpx.Response(200, json={"choices": [{"message": {"content": "96"}}]}) + + app = AgentsCLI().get_cli() + with _install_mock_transport(handler): + result = CliRunner().invoke( + app, + ["invoke", "--agent", "calc", "--input", "12*8", "--no-progress"], + obj=_FakeCLIContext(base_url="http://config-host:9999", token="tkn"), + ) + + assert result.exit_code == 0, result.stdout + (result.stderr or "") + assert captured[0].url.host == "config-host" + assert captured[0].url.port == 9999 + assert captured[0].headers.get("authorization") == "Bearer tkn" diff --git a/plugins/nemo-agents/tests/unit/test_cli_list_output.py b/plugins/nemo-agents/tests/unit/test_cli_list_output.py index ec4db3977c..e863fb8e11 100644 --- a/plugins/nemo-agents/tests/unit/test_cli_list_output.py +++ b/plugins/nemo-agents/tests/unit/test_cli_list_output.py @@ -78,7 +78,8 @@ def test_agents_list_supports_json_output(self, app, flag: str) -> None: result = runner.invoke(app, ["list", flag, "json"]) assert result.exit_code == 0, result.output - assert json.loads(result.output) == response + # The resolved-target banner goes to stderr; stdout stays clean JSON. + assert json.loads(result.stdout) == response class TestDeploymentsListOutput: @@ -103,4 +104,5 @@ def test_deployments_list_supports_json_output(self, app, flag: str) -> None: result = runner.invoke(app, ["deployments", "list", flag, "json"]) assert result.exit_code == 0, result.output - assert json.loads(result.output) == response + # The resolved-target banner goes to stderr; stdout stays clean JSON. + assert json.loads(result.stdout) == response diff --git a/plugins/nemo-agents/tests/unit/usage/test_cli.py b/plugins/nemo-agents/tests/unit/usage/test_cli.py index 46421ac421..e805018216 100644 --- a/plugins/nemo-agents/tests/unit/usage/test_cli.py +++ b/plugins/nemo-agents/tests/unit/usage/test_cli.py @@ -6,6 +6,7 @@ from __future__ import annotations import json +from contextlib import contextmanager from pathlib import Path from unittest.mock import patch @@ -16,6 +17,35 @@ runner = CliRunner() +class _FakeSDKUser: + def __init__(self, token: str | None) -> None: + self._token = token + + def get_client_config(self) -> dict[str, object]: + if self._token is None: + return {} + return {"default_headers": {"Authorization": f"Bearer {self._token}"}} + + +class _FakeSDKContext: + def __init__(self, base_url: str, token: str | None) -> None: + self.user = _FakeSDKUser(token) + self.cluster = type("_Cluster", (), {"base_url": base_url})() + + +class _FakeCLIContext: + """Minimal stand-in for ``CLIContext`` (typer.Context.obj).""" + + def __init__(self, base_url: str = "http://config-host:9999", token: str | None = "cfg-token") -> None: + self._sdk = _FakeSDKContext(base_url, token) + + def get_sdk_context(self) -> _FakeSDKContext: + return self._sdk + + def get_base_url(self, default: str | None = None) -> str | None: + return str(self._sdk.cluster.base_url) + + @pytest.fixture def app(): """Build the actual ``nemo agents`` Typer app (with ``usage`` registered).""" @@ -92,7 +122,8 @@ def test_usage_show_with_fileset_ref_uses_sdk(app, tmp_natjobs_dir: Path, fake_s ) assert result.exit_code == 0, result.output - payload = json.loads(result.output) + # The resolved-target banner goes to stderr; stdout stays clean JSON. + payload = json.loads(result.stdout) assert len(payload["runs"]) == 4 assert len(fake.files.calls) == 1 call = fake.files.calls[0] @@ -262,7 +293,8 @@ def test_usage_show_fileset_rewrites_source_dirs(app, tmp_natjobs_dir: Path, fak ) assert result.exit_code == 0, result.output - payload = json.loads(result.output) + # The resolved-target banner goes to stderr; stdout stays clean JSON. + payload = json.loads(result.stdout) # source_dir should be the synthetic / form, not the dead tempdir for run in payload["runs"]: assert run["source_dir"].startswith("my-fileset/") @@ -282,7 +314,8 @@ def test_usage_show_fileset_single_run_rel_dot(app, tmp_path: Path, fake_sdk_fac result = runner.invoke(app, ["usage", "show", "single-fileset"]) assert result.exit_code == 0, result.output - payload = json.loads(result.output) + # The resolved-target banner goes to stderr; stdout stays clean JSON. + payload = json.loads(result.stdout) # rel == "." branch: no slash + rel suffix; just the bare ref. assert payload["task"]["source_dir"] == "single-fileset" @@ -309,6 +342,39 @@ def test_usage_show_rejects_empty_or_dot_relative_fileset_name(app, tmp_path: Pa assert "must be a real fileset name" in result.output, (ref, result.output) +def test_usage_show_fileset_builds_sdk_with_context_base_url_and_auth(app, tmp_natjobs_dir: Path) -> None: + """A fileset ref builds the SDK client with the shared context's base URL + auth token. + + Pins P0 parity for ``usage show``: it must honor ``nemo config`` / + ``NMP_BASE_URL`` and attach the ``Authorization`` bearer token, instead + of defaulting to localhost with no auth. + """ + captured: dict[str, object] = {} + + def fake_platform(**kwargs: object) -> object: + captured.update(kwargs) + return object() + + @contextmanager + def fake_fileset_path(_ref, *, sdk, workspace): + yield tmp_natjobs_dir + + with ( + patch("nemo_agents_plugin.usage.cli.NeMoPlatform", fake_platform), + patch("nemo_agents_plugin.usage.cli.fileset_path", fake_fileset_path), + ): + result = runner.invoke( + app, + ["usage", "show", "my-fileset"], + obj=_FakeCLIContext(base_url="http://config-host:9999", token="tkn"), + ) + + assert result.exit_code == 0, result.output + assert captured["base_url"] == "http://config-host:9999" + assert captured["default_headers"] == {"Authorization": "Bearer tkn"} + assert "Targeting http://config-host:9999" in (result.stderr or "") + + def test_usage_show_with_workspace_qualified_fileset_ref(app, tmp_natjobs_dir: Path, fake_sdk_factory) -> None: """A ``ws/name`` ref overrides the default workspace.""" fake = fake_sdk_factory(tmp_natjobs_dir) From 0156e8cc0bd0aec5e5cf9e7a4b1a39464af58008 Mon Sep 17 00:00:00 2001 From: Tyler Bray Date: Tue, 7 Jul 2026 16:02:09 -0700 Subject: [PATCH 2/2] refactor(agents): Extract shared --base-url typer option Address review feedback: the --base-url option (and its help text) was repeated inline across all 12 platform commands plus `usage show`. Define it once in cli_context as a reusable Annotated `BaseUrlOption` (with the help text in a `BASE_URL_HELP` constant) and reuse it everywhere. Also reformat the help as a numbered resolution-order list. No behavior change. Signed-off-by: Tyler Bray --- .../nemo-agents/src/nemo_agents_plugin/cli.py | 135 ++---------------- .../src/nemo_agents_plugin/cli_context.py | 18 ++- .../src/nemo_agents_plugin/usage/cli.py | 13 +- 3 files changed, 34 insertions(+), 132 deletions(-) diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py index eac620a265..14a0a57619 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/cli.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/cli.py @@ -48,6 +48,9 @@ from nemo_agents_plugin.cli_context import ( DEFAULT_BASE_URL as _DEFAULT_BASE_URL, ) +from nemo_agents_plugin.cli_context import ( + BaseUrlOption, +) from nemo_agents_plugin.cli_context import ( resolve_base_url as _resolve_base_url, ) @@ -153,16 +156,7 @@ def invoke( help="Name of a specific deployment to invoke (platform required).", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, timeout: float = typer.Option( 300, "--timeout", @@ -646,16 +640,7 @@ def create( ), description: str = typer.Option("", "--description"), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Register an agent on the platform.""" base_url = _resolve_base_url(base_url) @@ -681,16 +666,7 @@ def create( def list_agents( ctx: typer.Context, workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, output_format: Optional[_LIST_OUTPUT_FORMAT] = typer.Option( None, "--format", @@ -722,16 +698,7 @@ def list_agents( def get( name: str = typer.Argument(..., help="Agent name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Get an agent by name.""" base_url = _resolve_base_url(base_url) @@ -742,16 +709,7 @@ def get( def delete( name: str = typer.Argument(..., help="Agent name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."), ) -> None: """Delete an agent from the platform.""" @@ -783,16 +741,7 @@ def deploy( help="Maximum seconds to wait for a terminal status (only with --wait).", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Deploy an agent on the platform. @@ -860,16 +809,7 @@ def logs( help="Print only the absolute log file path and exit (useful for scripting).", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Show logs for an agent deployment. @@ -937,16 +877,7 @@ def undeploy( None, "--agent", "--all", "-a", help="Remove all deployments for this agent." ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."), ) -> None: """Stop and remove a deployment (or all deployments for an agent).""" @@ -976,16 +907,7 @@ def undeploy( def deployments_list( ctx: typer.Context, workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, output_format: Optional[_LIST_OUTPUT_FORMAT] = typer.Option( None, "--format", @@ -1017,16 +939,7 @@ def deployments_list( def deployments_get( name: str = typer.Argument(..., help="Deployment name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Get a deployment by name.""" base_url = _resolve_base_url(base_url) @@ -1037,16 +950,7 @@ def deployments_get( def deployments_delete( name: str = typer.Argument(..., help="Deployment name."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt."), ) -> None: """Delete a deployment by name.""" @@ -1068,16 +972,7 @@ def deployments_wait( timeout: int = typer.Option(300, "--timeout", "-t", help="Maximum seconds to wait."), interval: float = typer.Option(2.0, "--interval", help="Poll interval in seconds."), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL. Defaults to the base URL from the shared CLI " - "config (`nemo config set --base-url`) / NMP_BASE_URL env var, then " - "http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Wait for a deployment to reach 'running' or 'failed' status. diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py b/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py index 800104ddcd..0a9fdc9156 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/cli_context.py @@ -20,14 +20,30 @@ from __future__ import annotations import logging -from typing import Any +from typing import Annotated, Any, Optional import click +import typer logger = logging.getLogger(__name__) DEFAULT_BASE_URL = "http://localhost:8080" +BASE_URL_HELP = ( + "Platform base URL. Resolution order: " + "(1) this --base-url flag or NEMO_BASE_URL; " + "(2) shared CLI config (`nemo config set --base-url`) or NMP_BASE_URL; " + f"(3) {DEFAULT_BASE_URL} (default)." +) + +# Reusable ``--base-url`` option shared across every ``nemo agents`` command +# so the option and its help text are defined once. ``None`` means "unset" so +# ``resolve_base_url`` can fall back to the shared CLI context / config. +BaseUrlOption = Annotated[ + Optional[str], + typer.Option("--base-url", envvar="NEMO_BASE_URL", help=BASE_URL_HELP), +] + def current_cli_state() -> Any: """Return the shared CLI context object (``typer.Context.obj``) if present. diff --git a/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py b/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py index eebccf123c..ca199d5b18 100644 --- a/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py +++ b/plugins/nemo-agents/src/nemo_agents_plugin/usage/cli.py @@ -19,7 +19,7 @@ from typing import Optional import typer -from nemo_agents_plugin.cli_context import resolve_base_url, resolve_context_headers +from nemo_agents_plugin.cli_context import BaseUrlOption, resolve_base_url, resolve_context_headers from nemo_agents_plugin.usage import compute, render from nemo_agents_plugin.usage import parser as parser_module from nemo_agents_plugin.usage.models import ( @@ -73,16 +73,7 @@ def show_cmd( "public number — leave unset and compute_units stays null.", ), workspace: str = typer.Option(_DEFAULT_WORKSPACE, "--workspace", "-w"), - base_url: Optional[str] = typer.Option( - None, - "--base-url", - envvar="NEMO_BASE_URL", - help=( - "Platform base URL for fileset downloads. Defaults to the base URL " - "from the shared CLI config (`nemo config set --base-url`) / " - "NMP_BASE_URL env var, then http://localhost:8080." - ), - ), + base_url: BaseUrlOption = None, ) -> None: """Show a usage report for *ref*.""" _show(