update

Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>

Author: ironcommit

openapi/ga/individual/platform.openapi.yaml

@@ -38,6 +38,7 @@
     StaticToken,
     TokenProvider,
 )
+from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR
 from nemo_platform_plugin.client.errors import (
     NemoResponseValidationError,
     NemoTransportError,
@@ -72,6 +73,59 @@
 logger = logging.getLogger(__name__)
 
 DEFAULT_TIMEOUT = 60.0
+_AUTHORIZATION_HEADER = "Authorization"
+_PRINCIPAL_ID_HEADER = "X-NMP-Principal-Id"
+
+
+def _has_header(headers: Mapping[str, str] | None, name: str) -> bool:
+    if not headers:
+        return False
+    normalized = name.lower()
+    return any(header.lower() == normalized for header in headers)
+
+
+@overload
+def _resolve_implicit_workload_auth(
+    *,
+    base_url: str,
+    auth: TokenProvider | str | None,
+    default_headers: Mapping[str, str] | None,
+    allow_env_bootstrap: bool,
+) -> TokenProvider | str | None: ...
+
+
+@overload
+def _resolve_implicit_workload_auth(
+    *,
+    base_url: str,
+    auth: TokenProvider | AsyncTokenProvider | str | None,
+    default_headers: Mapping[str, str] | None,
+    allow_env_bootstrap: bool,
+) -> TokenProvider | AsyncTokenProvider | str | None: ...
+
+
+def _resolve_implicit_workload_auth(
+    *,
+    base_url: str,
+    auth: TokenProvider | AsyncTokenProvider | str | None,
+    default_headers: Mapping[str, str] | None,
+    allow_env_bootstrap: bool,
+) -> TokenProvider | AsyncTokenProvider | str | None:
+    if auth is not None or not allow_env_bootstrap:
+        return auth
+    if _has_header(default_headers, _AUTHORIZATION_HEADER) or _has_header(default_headers, _PRINCIPAL_ID_HEADER):
+        return None
+
+    subject_token_file = os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR)
+    if not subject_token_file:
+        return None
+
+    from nemo_platform_plugin.client.oidc_factory import resolve_workload_exchange_provider
+
+    return resolve_workload_exchange_provider(
+        base_url=base_url,
+        subject_token_file=Path(subject_token_file),
+    )
 
 
 @cache
@@ -420,6 +474,12 @@ def __init__(
         defers to the transport's timeout, giving one we build ourselves
         :data:`DEFAULT_TIMEOUT`; ``httpx.Timeout(None)`` waits indefinitely.
         """
+        auth = _resolve_implicit_workload_auth(
+            base_url=base_url,
+            auth=auth,
+            default_headers=default_headers,
+            allow_env_bootstrap=http_client is None,
+        )
         super().__init__(
             base_url=base_url,
             workspace=workspace,
@@ -671,6 +731,12 @@ def __init__(
         url_resolver: Callable[[str], str | httpx.URL] | None = None,
     ) -> None:
         """Create a client. See :meth:`NemoClient.__init__` for *timeout*."""
+        auth = _resolve_implicit_workload_auth(
+            base_url=base_url,
+            auth=auth,
+            default_headers=default_headers,
+            allow_env_bootstrap=http_client is None,
+        )
         super().__init__(
             base_url=base_url,
             workspace=workspace,
@@ -911,8 +977,7 @@ def _client_from_config(
     """Shared implementation for NemoClient.from_config / AsyncNemoClient.from_config."""
     from nemo_platform_plugin.client.config.config import Config
     from nemo_platform_plugin.client.config.models import ConfigParams, OAuthUser
-    from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR
-    from nemo_platform_plugin.client.oidc_factory import resolve_oidc_provider, resolve_workload_exchange_provider
+    from nemo_platform_plugin.client.oidc_factory import resolve_oidc_provider
 
     resolved_path = Path(config_path) if isinstance(config_path, str) else config_path
     overrides: ConfigParams | None = None
@@ -928,13 +993,9 @@ def _client_from_config(
 
     auth: TokenProvider | str | None = None
     workload_identity_token_file = os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR)
+    use_implicit_workload_auth = bool(workload_identity_token_file and not explicit_access_token)
 
-    if workload_identity_token_file and not explicit_access_token:
-        auth = resolve_workload_exchange_provider(
-            base_url=str(ctx.cluster.base_url),
-            subject_token_file=Path(workload_identity_token_file),
-        )
-    elif isinstance(ctx.user, OAuthUser):
+    if not use_implicit_workload_auth and isinstance(ctx.user, OAuthUser):
         auth = resolve_oidc_provider(
             base_url=str(ctx.cluster.base_url),
             context_name=ctx.context_name,
@@ -944,7 +1005,7 @@ def _client_from_config(
             config_path=actual_config_path,
             explicit_access_token=explicit_access_token,
         )
-    elif ctx.user:
+    elif not use_implicit_workload_auth and ctx.user:
         client_config = ctx.user.get_client_config()
         raw_headers = client_config.get("default_headers")
         if isinstance(raw_headers, dict):

openapi/ga/openapi.yaml

@@ -16,13 +16,9 @@
 import json
 import logging
 import os
-from pathlib import Path
 from typing import Any
 
 from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient
-from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR
-from nemo_platform_plugin.client.oidc import WorkloadTokenExchangeProvider
-from nemo_platform_plugin.client.oidc_factory import resolve_workload_exchange_provider
 
 logger = logging.getLogger(__name__)
 
@@ -82,23 +78,6 @@ def _base_url() -> str:
     return os.environ.get("NMP_BASE_URL", "http://localhost:8080")
 
 
-def _workload_exchange_auth_from_env(
-    base_url: str,
-    headers: dict[str, str],
-) -> WorkloadTokenExchangeProvider | None:
-    if headers.get("X-NMP-Principal-Id"):
-        return None
-
-    subject_token_file = os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR)
-    if not subject_token_file:
-        return None
-
-    return resolve_workload_exchange_provider(
-        base_url=base_url,
-        subject_token_file=Path(subject_token_file),
-    )
-
-
 def get_nemo_client(
     *,
     as_service: str | None = None,
@@ -112,8 +91,7 @@ def get_nemo_client(
     """
     headers = _build_headers(as_service=as_service, internal=internal, on_behalf_of=on_behalf_of)
     base_url = _base_url()
-    auth = _workload_exchange_auth_from_env(base_url, headers)
-    return NemoClient(base_url=base_url, auth=auth, default_headers=headers or None)
+    return NemoClient(base_url=base_url, default_headers=headers or None)
 
 
 def get_async_nemo_client(
@@ -129,5 +107,4 @@ def get_async_nemo_client(
     """
     headers = _build_headers(as_service=as_service, internal=internal, on_behalf_of=on_behalf_of)
     base_url = _base_url()
-    auth = _workload_exchange_auth_from_env(base_url, headers)
-    return AsyncNemoClient(base_url=base_url, auth=auth, default_headers=headers or None)
+    return AsyncNemoClient(base_url=base_url, default_headers=headers or None)

openapi/openapi.yaml

@@ -130,6 +130,52 @@ def test_no_auth_no_header(self):
         assert route.called
         assert "Authorization" not in route.calls[0].request.headers
 
+    def test_constructor_uses_workload_exchange_provider_from_env(self, monkeypatch, tmp_path):
+        subject_token_file = tmp_path / "workload-token"
+        subject_token_file.write_text("subject-token\n", encoding="utf-8")
+        provider = object()
+        monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file))
+
+        with patch(
+            "nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider",
+            return_value=provider,
+        ) as resolve_provider:
+            client = NemoClient(base_url="https://nemo.example.com")
+
+        assert client._auth is provider
+        resolve_provider.assert_called_once_with(
+            base_url="https://nemo.example.com",
+            subject_token_file=subject_token_file,
+        )
+
+    def test_constructor_workload_exchange_does_not_override_authorization_header(self, monkeypatch, tmp_path):
+        subject_token_file = tmp_path / "workload-token"
+        subject_token_file.write_text("subject-token\n", encoding="utf-8")
+        monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file))
+
+        with patch("nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider") as resolve_provider:
+            client = NemoClient(
+                base_url="https://nemo.example.com",
+                default_headers={"Authorization": "Bearer explicit-token"},
+            )
+
+        assert client._auth is None
+        resolve_provider.assert_not_called()
+
+    def test_constructor_workload_exchange_does_not_override_principal_header(self, monkeypatch, tmp_path):
+        subject_token_file = tmp_path / "workload-token"
+        subject_token_file.write_text("subject-token\n", encoding="utf-8")
+        monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file))
+
+        with patch("nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider") as resolve_provider:
+            client = NemoClient(
+                base_url="https://nemo.example.com",
+                default_headers={"X-NMP-Principal-Id": "service:jobs"},
+            )
+
+        assert client._auth is None
+        resolve_provider.assert_not_called()
+
 
 # ---------------------------------------------------------------------------
 # AsyncNemoClient auth parameter
@@ -176,6 +222,24 @@ def test_sync_provider_works_in_async_client(self):
         assert route.called
         assert route.calls[0].request.headers["Authorization"] == "Bearer sync-token"
 
+    def test_constructor_uses_workload_exchange_provider_from_env(self, monkeypatch, tmp_path):
+        subject_token_file = tmp_path / "workload-token"
+        subject_token_file.write_text("subject-token\n", encoding="utf-8")
+        provider = object()
+        monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file))
+
+        with patch(
+            "nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider",
+            return_value=provider,
+        ) as resolve_provider:
+            client = AsyncNemoClient(base_url="https://nemo.example.com")
+
+        assert client._auth is provider
+        resolve_provider.assert_called_once_with(
+            base_url="https://nemo.example.com",
+            subject_token_file=subject_token_file,
+        )
+
 
 # ---------------------------------------------------------------------------
 # OIDCTokenProvider
@@ -802,7 +866,7 @@ def test_sync_client_uses_workload_exchange_provider_from_env(self, monkeypatch,
         monkeypatch.delenv("NMP_PRINCIPAL", raising=False)
 
         with patch(
-            "nemo_platform_plugin.client_provider.resolve_workload_exchange_provider",
+            "nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider",
             return_value=provider,
         ) as resolve_provider:
             client = get_nemo_client()
@@ -822,7 +886,7 @@ def test_async_client_uses_workload_exchange_provider_from_env(self, monkeypatch
         monkeypatch.delenv("NMP_PRINCIPAL", raising=False)
 
         with patch(
-            "nemo_platform_plugin.client_provider.resolve_workload_exchange_provider",
+            "nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider",
             return_value=provider,
         ) as resolve_provider:
             client = get_async_nemo_client()
@@ -839,7 +903,7 @@ def test_workload_exchange_provider_does_not_override_explicit_service(self, mon
         monkeypatch.setenv(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, str(subject_token_file))
         monkeypatch.delenv("NMP_PRINCIPAL", raising=False)
 
-        with patch("nemo_platform_plugin.client_provider.resolve_workload_exchange_provider") as resolve_provider:
+        with patch("nemo_platform_plugin.client.oidc_factory.resolve_workload_exchange_provider") as resolve_provider:
             client = get_nemo_client(as_service="jobs")
 
         assert client._auth is None

packages/nemo_platform_plugin/src/nemo_platform_plugin/client/client.py

@@ -99,14 +99,17 @@
                         "description": "Requested audience for the issued access token.",
                     },
                     "resource": {
+                        "type": "string",
                         "description": "Unsupported for the jobs workload token exchange profile.",
                         "not": {},
                     },
                     "actor_token": {
+                        "type": "string",
                         "description": "Unsupported for the jobs workload token exchange profile.",
                         "not": {},
                     },
                     "actor_token_type": {
+                        "type": "string",
                         "description": "Unsupported for the jobs workload token exchange profile.",
                         "not": {},
                     },