feat(client): Add client auth support#462
Open
matthewgrossman wants to merge 11 commits into
Open
Conversation
Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…uth-support-to-nemoclient-replace Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds auth, config, client construction, provider resolution, and FastAPI dependency support for Nemo clients, with tests covering auth, config, and provider flows. ChangesNemo client auth and config plumbing
Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (1)
packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py (1)
16-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a regular import and concrete return annotation.
Move
AsyncNemoClientout ofTYPE_CHECKINGand annotate directly. As per coding guidelines,**/*.py: “Always prefer concrete type hints over string-based ones in Python code; do not import types under TYPE_CHECKING, instead import types as regular imports when possible.”Proposed fix
- from nemo_platform_plugin.client.client import AsyncNemoClient +from nemo_platform_plugin.client.client import AsyncNemoClient-def get_nemo_client() -> "AsyncNemoClient": +def get_nemo_client() -> AsyncNemoClient:Also applies to: 55-55
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py` at line 16, The dependency provider is still relying on a TYPE_CHECKING-only import and a string-based return annotation, which should be replaced with a regular import and a concrete type hint. Update the dependency module that exposes AsyncNemoClient so the symbol is imported normally at runtime and the provider function returns AsyncNemoClient directly, keeping the annotation concrete and removing the need for deferred typing.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.py`:
- Around line 141-151: The header builder in _build_headers() currently forwards
id, email, and groups from _read_principal_from_env() but drops the delegated
identity field. Update the principal-handling block in client_provider.py so
that when principal includes on_behalf_of, it is propagated into the headers
using the same X-NMP-Principal-On-Behalf-Of key used for the explicit
on_behalf_of parameter, and keep the existing explicit parameter behavior
intact.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/auth.py`:
- Around line 284-290: The TokenSet dataclass is exposing sensitive credentials
through its auto-generated repr, and OIDCTokenProvider may also print them via
its tokens field. Update TokenSet so access_token and refresh_token are excluded
from repr, and ensure any TokenSet stored or returned by OIDCTokenProvider does
not leak token values in debug output by using the repr-safe fields
consistently. Locate the fix in TokenSet and the OIDCTokenProvider token
handling logic.
- Around line 426-437: The token refresh flow in `AuthClient` updates
`self.tokens` and then calls `on_tokens_refreshed()`, but if persistence fails
after a rotated refresh token is received, the failure is only logged and the
process continues with an unsaved invalid token. Update the `refresh` logic
around `TokenSet.from_access_token`, `self.on_tokens_refreshed`, and the
`logger.warning` path so persistence failures are surfaced to the caller or
otherwise stop the refresh from being treated as successful when the refresh
token has rotated.
- Around line 246-263: refresh_token_grant currently sends refresh tokens to
whatever token_endpoint is provided, so add endpoint validation before the
httpx.post call. In refresh_token_grant, reject malformed URLs and non-HTTPS
schemes up front, but allow explicit HTTP loopback/local-dev endpoints such as
localhost or 127.0.0.1. Keep the check close to the existing token_endpoint
handling so the function only proceeds to build data and post when the endpoint
is safe and valid.
- Around line 426-430: The refresh flow in auth.py is rebuilding the TokenSet
with TokenSet.from_access_token(), which ignores expires_in from the refresh
response and can leave opaque access tokens without an expiration. Update the
token refresh path in the client auth logic to read token_data["expires_in"]
(falling back to the current JWT exp behavior when available), compute and store
expires_at accordingly, and preserve the rotated refresh token when constructing
self.tokens so auto-refresh continues to work after the first refresh.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py`:
- Around line 335-341: The from_config flow is ignoring the requested context,
so Config.load and resolve are using the active/default context instead of the
caller’s override. Update the client configuration path in from_config to pass
the provided context through to the config loading/resolution path, and ensure
the same context is used when computing the resolved config state via
Config.load, config.resolve, and any related current_context handling.
- Around line 345-353: The token passed from ctx.user in the OAuthUser branch is
still treated as non-explicit, which can cause resolve_oidc_provider() to cache
or persist a caller-owned access token. Update the auth resolution path in
client.py to distinguish runtime-provided tokens from stored credentials, and
pass an explicit_access_token flag through the call chain so
resolve_oidc_provider() can avoid persisting tokens sourced from
NMP_ACCESS_TOKEN. Use the ctx.user/OAuthUser handling near
resolve_oidc_provider() to locate the change.
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/config/config.py`:
- Around line 357-358: Preserve explicit false overrides in the config assembly
logic: the truncate setting is being skipped in the config path because the
`self.truncate` truthiness check drops `False`, so update the `Config`/config
resolution code that builds `params` to distinguish “unset” from an explicit
boolean value and always propagate `truncate=False` from `NMP_TRUNCATE=false` or
`{"truncate": False}`.
- Around line 139-144: The config lookup in the config path resolver should
honor XDG_CONFIG_HOME even when the nmp directory has not been created yet.
Update the logic in the config path selection method (the block in config.py
that checks xdg_config_home and returns the config.yaml path) so it returns the
XDG-based location unconditionally when XDG_CONFIG_HOME is set, instead of
requiring config_dir.exists(). Keep the existing fallback behavior only for when
XDG_CONFIG_HOME is unset.
- Around line 225-240: The write flow in Config.write resolves and creates the
target context before applying params["current_context"], which can persist a
new current_context without ensuring that context exists. Update Config.write to
read and apply params["current_context"] early, then use that effective context
when determining context_name, calling cls.load/cls.create, and passing through
config_file.ensure_context. Make sure current_context and the created/selected
context stay aligned so a write with only current_context still creates a valid
context before persisting.
- Around line 201-210: The config save flow in config.py writes secrets via
self._config_file.model_dump(include_secrets=True) and then opens the target
file with open(path, "w"), which can leave the file briefly created/truncated
with default umask permissions before os.chmod tightens it. Update the save
logic in the config-writing path to ensure the destination is set to 0600 before
any secret-containing content is written, ideally by preparing/chmodding the
file first or using a secure create/write sequence around the existing
yaml.safe_dump call.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.py`:
- Line 16: The dependency provider is still relying on a TYPE_CHECKING-only
import and a string-based return annotation, which should be replaced with a
regular import and a concrete type hint. Update the dependency module that
exposes AsyncNemoClient so the symbol is imported normally at runtime and the
provider function returns AsyncNemoClient directly, keeping the annotation
concrete and removing the need for deferred typing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 16701030-b694-4838-8d52-fe4b7285a39d
📒 Files selected for processing (7)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/auth.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/config/config.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/config/models.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client_provider.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/dependencies.pypackages/nemo_platform_plugin/tests/test_client_auth.py
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
Contributor
|
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…implify config_exists, add edge case tests - Remove unused _on_behalf_of_headers from client_provider.py - Remove duplicate _TOKEN_REFRESH_MARGIN_SECONDS, use DEFAULT_REFRESH_MARGIN_SECONDS - Simplify config_exists logic in _client_from_config - Add tests: invalid_grant recovery, nonexistent config path, Config.write round-trip Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…ules - auth.py: protocols (TokenProvider, AsyncTokenProvider), StaticToken, AuthError - oidc.py: OIDCTokenProvider, TokenSet, JWT helpers, OIDC discovery, TokenRefreshError - oidc_factory.py: provider cache, config persistence callbacks, cross-process locking, resolve_oidc_provider No re-exports — each consumer imports from the actual module. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py (1)
15-15: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove postponed annotations.
This stringifies annotations across the module. Convert the forward return on
TokenSet.from_access_token()to a concrete runtime-safe form, then drop the future import.As per coding guidelines, "
**/*.py: Always prefer concrete type hints over string-based ones in Python code; do not import types under TYPE_CHECKING, instead import types as regular imports when possible."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py` at line 15, The module is still relying on postponed annotations, which stringifies type hints across the file. Update the forward return annotation on TokenSet.from_access_token() to a concrete runtime-safe type so it no longer depends on future annotations, then remove the from __future__ import annotations line from the oidc module. Keep the existing TokenSet symbol and any referenced types as regular imports so the annotations remain concrete and runtime-safe.Source: Coding guidelines
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/auth.py (1)
13-13: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove postponed annotations.
from __future__ import annotationsmakes annotations string-backed module-wide; use concrete runtime annotations/imports instead.As per coding guidelines, "
**/*.py: Always prefer concrete type hints over string-based ones in Python code; do not import types under TYPE_CHECKING, instead import types as regular imports when possible."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/auth.py` at line 13, Remove the postponed-annotations import from auth.py and switch the module to concrete runtime type hints/imports instead of string-based annotations. Update any affected annotations in the relevant symbols in this file so they reference real imported types directly, and avoid relying on TYPE_CHECKING-only imports where regular imports are possible.Source: Coding guidelines
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py (1)
11-11: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove postponed annotations.
No forward references here require module-wide postponed annotations; keep type hints concrete at runtime.
As per coding guidelines, "
**/*.py: Always prefer concrete type hints over string-based ones in Python code; do not import types under TYPE_CHECKING, instead import types as regular imports when possible."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py` at line 11, Remove the module-wide postponed annotations import from oidc_factory since there are no forward references that need it. Update the type hints in the OIDC factory module to use concrete runtime types directly, following the existing style in symbols like the factory helpers/classes in this file, and keep any needed imports as regular imports rather than relying on postponed evaluation.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py`:
- Around line 160-165: The auto-refresh setup in oidc_factory should fail fast
when discovery does not provide refresh settings, instead of allowing empty
token_endpoint/client_id to reach the refresh flow. Update the logic in the OIDC
client factory around _discover_oidc_client_settings and the
TokenSet/auto-refresh provider construction to validate
oidc_config.token_endpoint and oidc_config.client_id before proceeding, and
raise a clear error or skip auto-refresh when either value is missing. Use the
existing symbols oidc_config, token_endpoint, client_id, and refresh_scope to
keep the check close to where the provider is built.
- Around line 154-170: `Config.resolve()` is treating `Config.access_token` like
a config-file token, so runtime env/CLI overrides can be cached or persisted
incorrectly. Update the `Config.resolve()` call path to pass
`explicit_access_token=True` whenever the token comes from a runtime override,
and keep `build_oidc_token_provider()` using that flag to set `share_provider`
only for tokens sourced from actual config files. Ensure the `OIDCTokenProvider`
creation path still distinguishes resolved config tokens from override tokens.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/auth.py`:
- Line 13: Remove the postponed-annotations import from auth.py and switch the
module to concrete runtime type hints/imports instead of string-based
annotations. Update any affected annotations in the relevant symbols in this
file so they reference real imported types directly, and avoid relying on
TYPE_CHECKING-only imports where regular imports are possible.
In
`@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.py`:
- Line 11: Remove the module-wide postponed annotations import from oidc_factory
since there are no forward references that need it. Update the type hints in the
OIDC factory module to use concrete runtime types directly, following the
existing style in symbols like the factory helpers/classes in this file, and
keep any needed imports as regular imports rather than relying on postponed
evaluation.
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py`:
- Line 15: The module is still relying on postponed annotations, which
stringifies type hints across the file. Update the forward return annotation on
TokenSet.from_access_token() to a concrete runtime-safe type so it no longer
depends on future annotations, then remove the from __future__ import
annotations line from the oidc module. Keep the existing TokenSet symbol and any
referenced types as regular imports so the annotations remain concrete and
runtime-safe.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 19f151bf-8d79-4041-bc94-7c1e1388aa13
📒 Files selected for processing (6)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/auth.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/config/config.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc_factory.pypackages/nemo_platform_plugin/tests/test_client_auth.py
🚧 Files skipped from review as they are similar to previous changes (3)
- packages/nemo_platform_plugin/src/nemo_platform_plugin/client/config/config.py
- packages/nemo_platform_plugin/tests/test_client_auth.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py
Bug fixes: - Fix context parameter silently ignored in from_config() - Propagate on_behalf_of from NMP_PRINCIPAL in client_provider headers Security: - Hide token values from TokenSet/OIDCTokenProvider repr - Validate HTTPS on token endpoint (allow HTTP loopback for local dev) - Fail fast when OIDC discovery can't provide token_endpoint/client_id Style: - Use concrete type hint for AsyncNemoClient in dependencies.py - Fix refresh_token_grant return type to dict[str, Any] Tests: - Add test for from_config context selection - Add tests for legacy api-key migration - Add tests for token repr hiding and endpoint validation Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…uth-support-to-nemoclient-replace Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
- Honor expires_in from refresh responses for opaque (non-JWT) tokens - Surface persistence failure when IdP rotates the refresh token (only swallow persistence errors when the refresh token wasn't rotated) - Pass explicit_access_token=True when token comes from NMP_ACCESS_TOKEN env var so the provider isn't cached/persisted against the config file - Honor XDG_CONFIG_HOME even before the directory exists - Set file permissions (0600) atomically via os.open before writing secrets (eliminates race window where file is world-readable) - Use params["current_context"] to determine write target in Config.write() - Preserve explicit truncate=False in default config resolution (use `is not None` checks instead of truthiness) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
…uth-support-to-nemoclient-replace Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
maxdubrinsky
approved these changes
Jun 26, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds first-class auth support to
NemoClient/AsyncNemoClientso they can authenticate independently without requiring aNeMoPlatforminstance. Also adds aget_nemo_client()global provider analogous toget_platform_sdk().Linear: AIRCORE-828
What changed
New
authparameter onNemoClient/AsyncNemoClient:auth="nvapi-xxx"— static token (wrapped inStaticToken)auth=OIDCTokenProvider(...)— OIDC with auto-refreshauth=None— no auth (default, backwards compatible)Auth injection in
send()—Authorization: Bearer <token>header set on every request. Asyncsend()handles three cases:get_access_token_async()— explicit async method (e.g.OIDCTokenProvider)async get_access_token()— pure async provider (coroutine function)get_access_token()— wrapped inasyncio.to_threadto avoid blocking the event loop during token refresh IONemoClient.from_config()/AsyncNemoClient.from_config()— creates a client from~/.config/nmp/config.yaml, wiring up OIDC token refresh, config persistence, and cross-process locking automatically.NemoClientProviderprotocol +get_nemo_client()— separate provider system (entry-point groupnemo.client_provider) analogous toSDKProvider/get_platform_sdk(). Services migrate incrementally as downstream services expose typed endpoints.Config + auth code moved from
nemo_platform_extintonemo_platform_plugin:client/auth.py—TokenProvider,AsyncTokenProvider,StaticToken,OIDCTokenProvider,TokenSet, OIDC discovery, provider cacheclient/config/models.py—OAuthUser,NoAuthUser,Cluster,Context,ConfigFile, etc.client/config/config.py—Configclass withload(),resolve(),save(),write()client_provider.py—NemoClientProvider,DefaultNemoClientProvider,get_nemo_client(),get_async_nemo_client()dependencies.py—get_nemo_client()FastAPI dependency stubFiles
nemo_platform_plugin/client/auth.pyStaticToken,OIDCTokenProvider, OIDC helpers, provider factorynemo_platform_plugin/client/config/models.pynemo_platform_plugin/client/config/config.pyload,resolve,save,write)nemo_platform_plugin/client/client.pyauthparam, auth injection insend(),from_config()classmethodsnemo_platform_plugin/client_provider.pyNemoClientProviderprotocol + default providernemo_platform_plugin/dependencies.pyget_nemo_client()FastAPI dependency stubtests/test_client_auth.pyMerge order
This PR should land before PR #429 (AIRCORE-839, fsspec migration). Once merged, #429 rebases on top so
FilesetFileSystemgets aNemoClientwith first-class auth from the start — no skipped tests needed.Test plan
NemoClient(auth="token")setsAuthorizationheader on requestsNemoClient(auth=CustomProvider())callsget_access_token()per requestAsyncNemoClientwith async provider usesget_access_token_async()AsyncNemoClientwith sync provider runs inasyncio.to_threadStaticTokensatisfiesTokenProviderprotocolOIDCTokenProviderrefreshes expired tokens and persists via callbackOIDCTokenProviderraises when no refresh token availableConfig.load()+resolve()with OAuth and NoAuth usersNemoClient.from_config()with OAuth, static token, and no-auth configsAsyncNemoClient.from_config()returns correct client typeDefaultNemoClientProviderreturnsNemoClient/AsyncNemoClientget_nemo_client()/get_async_nemo_client()delegate to provider🤖 Generated with Claude Code
Summary by CodeRabbit
Authorization: Bearer …header injection and shared sync/async token providers.