fix(jobs): jobs-launcher authorization without token exchange - #920
Conversation
|
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:
📝 WalkthroughWalkthroughChangesWorkload identity authentication
Sequence Diagram(s)sequenceDiagram
participant Task as Workspace task
participant Factory as SDK factory
participant Exchange as Token exchange
participant Platform as Platform API
Task->>Factory: get_task_sdk(as_service="jobs")
Factory->>Exchange: Exchange token-file subject token
Exchange-->>Factory: Return access token
Factory->>Platform: Send workspace request with Bearer token
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)Error: can't load config: unsupported version of the configuration: "" See https://golangci-lint.run/docs/product/migration-guide for migration instructions Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py (1)
170-175: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClass docstring is now inaccurate.
Docstring says the provider only reads
NMP_BASE_URL/NMP_PRINCIPAL, butget_task_sdk/get_async_task_sdknow branch entirely onWORKLOAD_IDENTITY_TOKEN_FILE_ENVVARand drop principal delegation when set. Update the docstring to mention this third config source and the resulting behavior change.Also applies to: 178-183, 207-212
🤖 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/sdk_provider.py` around lines 170 - 175, Update the provider class docstring to document WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR as an additional configuration source: when set, get_task_sdk and get_async_task_sdk use workload identity instead of principal delegation. Retain the existing NMP_BASE_URL and NMP_PRINCIPAL behavior for the non-workload-identity path.
🤖 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/nmp_common/src/nmp/common/sdk_factory.py`:
- Around line 307-309: Update the workload-identity shortcut in get_task_sdk()
and get_async_task_sdk() so it does not bypass as_service handling for UDS
endpoints: either restrict the shortcut to endpoint.transport != "uds" or
forward as_service to get_platform_sdk()/get_async_platform_sdk(). Preserve the
existing shortcut behavior for non-UDS transports.
- Around line 277-283: Update the workload-identity bootstrap branches that
construct NeMoPlatform and AsyncNeMoPlatform to pass the resolved HTTP client
via http_client, preserving pooled-client reuse and async _test_http_client
injection; omit it only when the SDK must own the client for token-exchange
setup.
---
Nitpick comments:
In `@packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py`:
- Around line 170-175: Update the provider class docstring to document
WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR as an additional configuration source: when
set, get_task_sdk and get_async_task_sdk use workload identity instead of
principal delegation. Retain the existing NMP_BASE_URL and NMP_PRINCIPAL
behavior for the non-workload-identity path.
🪄 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: a3e8c36e-b055-48ec-af98-953a6512497a
📒 Files selected for processing (10)
packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/src/nmp/common/sdk_factory.pypackages/nmp_common/tests/auth/test_middleware.pypackages/nmp_common/tests/sdk_factory/test_sdk.pyservices/core/jobs/jobs-launcher/cmd/otel.goservices/core/jobs/jobs-launcher/cmd/workload_auth.goservices/core/jobs/jobs-launcher/cmd/workload_auth_test.goservices/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.pyservices/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py
|
442786e to
87c6ea7
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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
`@services/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py`:
- Around line 84-88: Update workspace_response to explicitly reject the obsolete
service Bearer token in the Authorization header, even when the
X-NMP-Principal-Id and X-NMP-Principal-On-Behalf-Of headers are valid; preserve
the existing 401 and 403 responses for invalid principal headers.
🪄 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: ce8d49b3-09d0-4af2-864c-adfe5e49074a
📒 Files selected for processing (10)
packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/src/nmp/common/sdk_factory.pypackages/nmp_common/tests/auth/test_middleware.pypackages/nmp_common/tests/sdk_factory/test_sdk.pyservices/core/jobs/jobs-launcher/cmd/otel.goservices/core/jobs/jobs-launcher/cmd/workload_auth.goservices/core/jobs/jobs-launcher/cmd/workload_auth_test.goservices/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.pyservices/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py
🚧 Files skipped from review as they are similar to previous changes (7)
- services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py
- services/core/jobs/jobs-launcher/cmd/workload_auth_test.go
- services/core/jobs/jobs-launcher/cmd/workload_auth.go
- packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py
- packages/nmp_common/src/nmp/common/sdk_factory.py
- packages/nmp_common/tests/auth/test_middleware.py
- services/core/jobs/jobs-launcher/cmd/otel.go
8661b25 to
9071930
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/nmp_common/tests/sdk_factory/test_sdk.py`:
- Around line 225-245: Update
test_get_async_platform_sdk_workload_identity_reuses_test_http_client to set
NMP_PRINCIPAL and assert the SDK client’s internal headers contain the expected
workload-identity credentials while omitting all X-NMP-Principal-* headers. Keep
the existing client-reuse and base URL assertions, and ensure the test
specifically exercises the token-file workload-identity branch rather than the
normal async factory path.
🪄 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: cd948b11-0b3e-400e-ad33-39405d939408
📒 Files selected for processing (10)
packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/src/nmp/common/sdk_factory.pypackages/nmp_common/tests/auth/test_middleware.pypackages/nmp_common/tests/sdk_factory/test_sdk.pyservices/core/jobs/jobs-launcher/cmd/otel.goservices/core/jobs/jobs-launcher/cmd/workload_auth.goservices/core/jobs/jobs-launcher/cmd/workload_auth_test.goservices/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.pyservices/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py
🚧 Files skipped from review as they are similar to previous changes (7)
- services/core/jobs/jobs-launcher/cmd/workload_auth.go
- services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py
- packages/nmp_common/src/nmp/common/sdk_factory.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py
- packages/nmp_common/tests/auth/test_middleware.py
- services/core/jobs/jobs-launcher/cmd/workload_auth_test.go
- services/core/jobs/jobs-launcher/cmd/otel.go
698e9e3 to
8c68822
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (1)
packages/nmp_common/tests/sdk_factory/test_sdk.py (1)
227-249: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winStill doesn't isolate the workload-identity branch.
Marked "Addressed" previously, but this test only asserts client reuse + base_url — both hold true on the normal (non-workload) async factory path too, since
_async_http_client_for_endpointis called in both branches. SetNMP_PRINCIPALand assert internal headers / absence ofX-NMP-Principal-*, as originally requested, so the test actually fails if the workload-identity branch is skipped.🤖 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/nmp_common/tests/sdk_factory/test_sdk.py` around lines 227 - 249, Update test_get_async_platform_sdk_workload_identity_reuses_test_http_client to set NMP_PRINCIPAL and assert the SDK’s internal headers reflect workload identity, including the expected token-derived headers and absence of X-NMP-Principal-* headers. Retain the existing client-reuse and base_url assertions so the test specifically distinguishes the workload-identity branch from the normal async factory path.
🤖 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.
Duplicate comments:
In `@packages/nmp_common/tests/sdk_factory/test_sdk.py`:
- Around line 227-249: Update
test_get_async_platform_sdk_workload_identity_reuses_test_http_client to set
NMP_PRINCIPAL and assert the SDK’s internal headers reflect workload identity,
including the expected token-derived headers and absence of X-NMP-Principal-*
headers. Retain the existing client-reuse and base_url assertions so the test
specifically distinguishes the workload-identity branch from the normal async
factory path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c6c24dda-d97a-4793-b882-3528633e0798
📒 Files selected for processing (13)
packages/nemo_platform_plugin/src/nemo_platform_plugin/client/constants.pypackages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.pypackages/nemo_platform_plugin/tests/test_sdk_provider.pypackages/nmp_common/src/nmp/common/platform_endpoint.pypackages/nmp_common/src/nmp/common/sdk_factory.pypackages/nmp_common/tests/auth/test_middleware.pypackages/nmp_common/tests/sdk_factory/test_sdk.pypackages/nmp_common/tests/test_platform_endpoint.pyservices/core/jobs/jobs-launcher/cmd/otel.goservices/core/jobs/jobs-launcher/cmd/workload_auth.goservices/core/jobs/jobs-launcher/cmd/workload_auth_test.goservices/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.pyservices/hello-world/tests/integration/tasks/test_workload_workspace_get_task.py
🚧 Files skipped from review as they are similar to previous changes (9)
- services/core/jobs/jobs-launcher/cmd/workload_auth.go
- packages/nmp_common/src/nmp/common/platform_endpoint.py
- services/hello-world/src/nmp/hello_world/tasks/workload_workspace_get/run.py
- packages/nmp_common/tests/test_platform_endpoint.py
- packages/nemo_platform_plugin/src/nemo_platform_plugin/sdk_provider.py
- services/core/jobs/jobs-launcher/cmd/workload_auth_test.go
- packages/nmp_common/src/nmp/common/sdk_factory.py
- services/core/jobs/jobs-launcher/cmd/otel.go
- packages/nmp_common/tests/auth/test_middleware.py
Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>
8c68822 to
adbc705
Compare
Fix Kubernetes job log upload auth without service bearer spoofing
Summary
This fixes Kubernetes job log upload and task-side SDK auth when platform auth is enabled but workload token exchange is disabled.
In that configuration, jobs do not receive
NMP_WORKLOAD_IDENTITY_TOKEN_FILE, so they cannot exchange a workload identity token for a real access token. The previous fallback sent:That is not valid auth for the Files OTLP log upload endpoint, so job logs failed with
401 Unauthorized:Problem
The failing Kubernetes setup has:
auth.enabled=truehttp://nemo-platform-api:8080/apis/files/v2/workspaces/{workspace}/filesets/{fileset}/otlp/v1/logsThe jobs launcher correctly detected that no workload token file was configured, but it fell back to a raw service bearer token. That raw bearer shortcut is only supported by the legacy HF-compatible Files path. It is intentionally not supported across normal API endpoints because an external caller that can reach the API directly could spoof a service principal with
Authorization: Bearer service:<name>.The same configuration also exposed a task-side SDK issue: individual tasks should not need to know whether workload token exchange is enabled. Task code should call a single SDK helper and let the SDK/factory layer choose the right auth mechanism.
Root Cause
The no-workload-token path conflated two different auth mechanisms:
Authorization: Bearer <access-token>from OIDC or workload token exchange.X-NMP-Principal-Id: service:<name>inside the trusted service boundary.When workload exchange is disabled,
Bearer service:jobsis not a real access token. The Files OTLP endpoint goes through normal middleware auth, so it returned 401.Solution
Use internal service-principal headers for the no-workload-token fallback, and keep real bearer auth only for real tokens.
Changes included:
Jobs launcher log upload:
NMP_WORKLOAD_IDENTITY_TOKEN_FILEis present, exchange the workload token and sendAuthorization: Bearer <access-token>.X-NMP-Principal-Id: service:jobs.Workload auth source:
Headers(ctx)method so token-backed auth and header-backed auth share one exporter path.Task SDK behavior:
nmp.common.sdk_factory.get_task_sdk(as_service=...)now chooses workload identity whenNMP_WORKLOAD_IDENTITY_TOKEN_FILEis set.get_task_sdk(as_service=...)usesX-NMP-Principal-Id: service:<name>plus on-behalf-of headers fromNMP_PRINCIPAL.Plugin SDK provider:
nemo_platform_plugin.sdk_provider.DefaultSDKProviderso plugin task authors also use one helper.Hello-world workload task:
get_task_sdk(as_service="jobs").Summary by CodeRabbit
X-NMP-Principal-Idwhen workload identity isn’t enabled.