nvidia-nat-mcp 1.4.0a20251014__py3-none-any.whl → 1.5.0a20260115__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of nvidia-nat-mcp might be problematic. Click here for more details.

Files changed (40) hide show
  1. nat/meta/pypi.md +1 -1
  2. nat/plugins/mcp/__init__.py +1 -1
  3. nat/plugins/mcp/auth/__init__.py +1 -1
  4. nat/plugins/mcp/auth/auth_flow_handler.py +65 -1
  5. nat/plugins/mcp/auth/auth_provider.py +3 -2
  6. nat/plugins/mcp/auth/auth_provider_config.py +5 -2
  7. nat/plugins/mcp/auth/register.py +9 -1
  8. nat/plugins/mcp/auth/service_account/__init__.py +14 -0
  9. nat/plugins/mcp/auth/service_account/provider.py +136 -0
  10. nat/plugins/mcp/auth/service_account/provider_config.py +137 -0
  11. nat/plugins/mcp/auth/service_account/token_client.py +156 -0
  12. nat/plugins/mcp/auth/token_storage.py +2 -2
  13. nat/plugins/mcp/cli/__init__.py +15 -0
  14. nat/plugins/mcp/cli/commands.py +1094 -0
  15. nat/plugins/mcp/client/__init__.py +15 -0
  16. nat/plugins/mcp/{client_base.py → client/client_base.py} +18 -10
  17. nat/plugins/mcp/{client_config.py → client/client_config.py} +24 -9
  18. nat/plugins/mcp/{client_impl.py → client/client_impl.py} +253 -62
  19. nat/plugins/mcp/exception_handler.py +1 -1
  20. nat/plugins/mcp/exceptions.py +1 -1
  21. nat/plugins/mcp/register.py +5 -4
  22. nat/plugins/mcp/server/__init__.py +15 -0
  23. nat/plugins/mcp/server/front_end_config.py +109 -0
  24. nat/plugins/mcp/server/front_end_plugin.py +155 -0
  25. nat/plugins/mcp/server/front_end_plugin_worker.py +415 -0
  26. nat/plugins/mcp/server/introspection_token_verifier.py +72 -0
  27. nat/plugins/mcp/server/memory_profiler.py +320 -0
  28. nat/plugins/mcp/server/register_frontend.py +27 -0
  29. nat/plugins/mcp/server/tool_converter.py +290 -0
  30. nat/plugins/mcp/utils.py +153 -36
  31. {nvidia_nat_mcp-1.4.0a20251014.dist-info → nvidia_nat_mcp-1.5.0a20260115.dist-info}/METADATA +5 -5
  32. nvidia_nat_mcp-1.5.0a20260115.dist-info/RECORD +37 -0
  33. nvidia_nat_mcp-1.5.0a20260115.dist-info/entry_points.txt +9 -0
  34. nat/plugins/mcp/tool.py +0 -138
  35. nvidia_nat_mcp-1.4.0a20251014.dist-info/RECORD +0 -23
  36. nvidia_nat_mcp-1.4.0a20251014.dist-info/entry_points.txt +0 -3
  37. {nvidia_nat_mcp-1.4.0a20251014.dist-info → nvidia_nat_mcp-1.5.0a20260115.dist-info}/WHEEL +0 -0
  38. {nvidia_nat_mcp-1.4.0a20251014.dist-info → nvidia_nat_mcp-1.5.0a20260115.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  39. {nvidia_nat_mcp-1.4.0a20251014.dist-info → nvidia_nat_mcp-1.5.0a20260115.dist-info}/licenses/LICENSE.md +0 -0
  40. {nvidia_nat_mcp-1.4.0a20251014.dist-info → nvidia_nat_mcp-1.5.0a20260115.dist-info}/top_level.txt +0 -0
nat/meta/pypi.md CHANGED
@@ -1,5 +1,5 @@
1
1
  <!--
2
- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3
3
  SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -54,6 +54,9 @@ class MCPAuthenticationFlowHandler(ConsoleAuthenticationFlowHandler):
54
54
  self._redirect_app: FastAPI | None = None
55
55
  self._server_lock = asyncio.Lock()
56
56
  self._oauth_client: AsyncOAuth2Client | None = None
57
+ self._redirect_host: str = "localhost" # Default host, will be overridden from config
58
+ self._redirect_port: int = 8000 # Default port, will be overridden from config
59
+ self._server_task: asyncio.Task | None = None
57
60
 
58
61
  async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
59
62
  """
@@ -88,6 +91,34 @@ class MCPAuthenticationFlowHandler(ConsoleAuthenticationFlowHandler):
88
91
  async def _handle_oauth2_auth_code_flow(self, cfg: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext:
89
92
  logger.info("Starting MCP OAuth2 authorization code flow")
90
93
 
94
+ # Extract and validate host and port from redirect_uri for callback server
95
+ from urllib.parse import urlparse
96
+ parsed_uri = urlparse(str(cfg.redirect_uri))
97
+
98
+ # Validate scheme/host and choose a safe non-privileged bind port
99
+ scheme = (parsed_uri.scheme or "http").lower()
100
+ if scheme not in ("http", "https"):
101
+ raise ValueError(f"redirect_uri must use http or https scheme, got '{scheme}'")
102
+
103
+ host = parsed_uri.hostname
104
+ if not host:
105
+ raise ValueError("redirect_uri must include a hostname, for example http://localhost:8000/auth/redirect")
106
+
107
+ # Never auto-bind to 80/443; default to 8000 when port is not specified
108
+ port = parsed_uri.port or 8000
109
+ if not (1 <= port <= 65535):
110
+ raise ValueError(f"Invalid redirect port: {port}. Expected 1-65535.")
111
+
112
+ if scheme == "https" and parsed_uri.port is None:
113
+ logger.warning(
114
+ "redirect_uri uses https without an explicit port; binding to %d (plain HTTP). "
115
+ "Terminate TLS at a reverse proxy and forward to this port.",
116
+ port)
117
+
118
+ self._redirect_host = host
119
+ self._redirect_port = port
120
+ logger.info("MCP redirect server will use %s:%d", self._redirect_host, self._redirect_port)
121
+
91
122
  state = secrets.token_urlsafe(16)
92
123
  flow_state = _FlowState()
93
124
  client = self.construct_oauth_client(cfg)
@@ -142,3 +173,36 @@ class MCPAuthenticationFlowHandler(ConsoleAuthenticationFlowHandler):
142
173
  "raw_token": token,
143
174
  },
144
175
  )
176
+
177
+ async def _start_redirect_server(self) -> None:
178
+ """
179
+ Override to use the host and port from redirect_uri config instead of hardcoded localhost:8000.
180
+
181
+ This allows MCP authentication to work with custom redirect hosts and ports
182
+ specified in the configuration.
183
+ """
184
+ # If the server is already running, do nothing
185
+ if self._server_controller:
186
+ return
187
+ try:
188
+ if not self._redirect_app:
189
+ raise RuntimeError("Redirect app not built.")
190
+
191
+ self._server_controller = _FastApiFrontEndController(self._redirect_app)
192
+
193
+ self._server_task = asyncio.create_task(
194
+ self._server_controller.start_server(host=self._redirect_host, port=self._redirect_port))
195
+ logger.debug("MCP redirect server starting on %s:%d", self._redirect_host, self._redirect_port)
196
+
197
+ # Wait for the server to bind (max ~10s)
198
+ start = asyncio.get_running_loop().time()
199
+ while True:
200
+ server = getattr(self._server_controller, "_server", None)
201
+ if server and getattr(server, "started", False):
202
+ break
203
+ if asyncio.get_running_loop().time() - start > 10:
204
+ raise RuntimeError("Redirect server did not report ready within 10s")
205
+ await asyncio.sleep(0.1)
206
+ except Exception as exc:
207
+ raise RuntimeError(
208
+ f"Failed to start MCP redirect server on {self._redirect_host}:{self._redirect_port}: {exc}") from exc
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -34,6 +34,7 @@ from nat.authentication.interfaces import AuthFlowType
34
34
  from nat.authentication.interfaces import AuthProviderBase
35
35
  from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
36
36
  from nat.data_models.authentication import AuthResult
37
+ from nat.data_models.common import get_secret_value
37
38
  from nat.plugins.mcp.auth.auth_flow_handler import MCPAuthenticationFlowHandler
38
39
  from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
39
40
 
@@ -371,7 +372,7 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
371
372
  # Manual registration mode
372
373
  self._cached_credentials = OAuth2Credentials(
373
374
  client_id=self.config.client_id,
374
- client_secret=self.config.client_secret,
375
+ client_secret=get_secret_value(self.config.client_secret),
375
376
  )
376
377
  logger.info("Using manual client_id: %s", self._cached_credentials.client_id)
377
378
  else:
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,6 +18,7 @@ from pydantic import HttpUrl
18
18
  from pydantic import model_validator
19
19
 
20
20
  from nat.authentication.interfaces import AuthProviderBaseConfig
21
+ from nat.data_models.common import OptionalSecretStr
21
22
 
22
23
 
23
24
  class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
@@ -36,7 +37,8 @@ class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
36
37
 
37
38
  # Client registration (manual registration vs DCR)
38
39
  client_id: str | None = Field(default=None, description="OAuth2 client ID for pre-registered clients")
39
- client_secret: str | None = Field(default=None, description="OAuth2 client secret for pre-registered clients")
40
+ client_secret: OptionalSecretStr = Field(default=None,
41
+ description="OAuth2 client secret for pre-registered clients")
40
42
  enable_dynamic_registration: bool = Field(default=True,
41
43
  description="Enable OAuth2 Dynamic Client Registration (RFC 7591)")
42
44
  client_name: str = Field(default="NAT MCP Client", description="OAuth2 client name for dynamic registration")
@@ -50,6 +52,7 @@ class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
50
52
  # Advanced options
51
53
  use_pkce: bool = Field(default=True, description="Use PKCE for authorization code flow")
52
54
 
55
+ # These fields are only used for shared workflow (not per-user workflows)
53
56
  default_user_id: str | None = Field(default=None, description="Default user ID for authentication")
54
57
  allow_default_user_id_for_tool_calls: bool = Field(default=True, description="Allow default user ID for tool calls")
55
58
 
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -17,9 +17,17 @@ from nat.builder.builder import Builder
17
17
  from nat.cli.register_workflow import register_auth_provider
18
18
  from nat.plugins.mcp.auth.auth_provider import MCPOAuth2Provider
19
19
  from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
20
+ from nat.plugins.mcp.auth.service_account.provider import MCPServiceAccountProvider
21
+ from nat.plugins.mcp.auth.service_account.provider_config import MCPServiceAccountProviderConfig
20
22
 
21
23
 
22
24
  @register_auth_provider(config_type=MCPOAuth2ProviderConfig)
23
25
  async def mcp_oauth2_provider(authentication_provider: MCPOAuth2ProviderConfig, builder: Builder):
24
26
  """Register MCP OAuth2 authentication provider with NAT system."""
25
27
  yield MCPOAuth2Provider(authentication_provider, builder=builder)
28
+
29
+
30
+ @register_auth_provider(config_type=MCPServiceAccountProviderConfig)
31
+ async def mcp_service_account_provider(authentication_provider: MCPServiceAccountProviderConfig, builder: Builder):
32
+ """Register MCP Service Account authentication provider with NAT system."""
33
+ yield MCPServiceAccountProvider(authentication_provider, builder=builder)
@@ -0,0 +1,14 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
@@ -0,0 +1,136 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import asyncio
17
+ import importlib
18
+ import logging
19
+ import typing
20
+
21
+ from pydantic import SecretStr
22
+
23
+ from nat.authentication.interfaces import AuthProviderBase
24
+ from nat.data_models.authentication import AuthResult
25
+ from nat.data_models.authentication import Credential
26
+ from nat.data_models.authentication import HeaderCred
27
+ from nat.plugins.mcp.auth.service_account.provider_config import MCPServiceAccountProviderConfig
28
+ from nat.plugins.mcp.auth.service_account.token_client import ServiceAccountTokenClient
29
+
30
+ logger = logging.getLogger(__name__)
31
+
32
+
33
+ class MCPServiceAccountProvider(AuthProviderBase[MCPServiceAccountProviderConfig]):
34
+ """
35
+ MCP service account authentication provider using OAuth2 client credentials.
36
+
37
+ Provides headless authentication for MCP clients using service account credentials.
38
+ Supports two authentication patterns:
39
+
40
+ 1. Single authentication: OAuth2 service account token only
41
+ 2. Dual authentication: OAuth2 service account token + service-specific token
42
+
43
+ """
44
+
45
+ def __init__(self, config: MCPServiceAccountProviderConfig, builder=None):
46
+ super().__init__(config)
47
+
48
+ # Initialize token client
49
+ self._token_client = ServiceAccountTokenClient(
50
+ client_id=config.client_id,
51
+ client_secret=config.client_secret,
52
+ token_url=config.token_url,
53
+ scopes=" ".join(config.scopes), # Convert list to space-delimited string for OAuth2
54
+ token_cache_buffer_seconds=config.token_cache_buffer_seconds,
55
+ )
56
+
57
+ # Load dynamic service token function if configured
58
+ self._service_token_function = None
59
+ if config.service_token and config.service_token.function:
60
+ self._service_token_function = self._load_function(config.service_token.function)
61
+
62
+ logger.info("Initialized MCP service account auth provider: "
63
+ "token_url=%s, scopes=%s, has_service_token=%s",
64
+ config.token_url,
65
+ config.scopes,
66
+ config.service_token is not None)
67
+
68
+ def _load_function(self, function_path: str) -> typing.Callable:
69
+ """Load a Python function from a module path string (e.g., 'my_module.get_token')."""
70
+ try:
71
+ module_name, func_name = function_path.rsplit(".", 1)
72
+ module = importlib.import_module(module_name)
73
+ func = getattr(module, func_name)
74
+ logger.info("Loaded service token function: %s", function_path)
75
+ return func
76
+ except Exception as e:
77
+ raise ValueError(f"Failed to load service token function '{function_path}': {e}") from e
78
+
79
+ async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
80
+ """
81
+ Authenticate using OAuth2 client credentials flow.
82
+
83
+ Note: user_id is ignored for service accounts (non-session-specific).
84
+
85
+ Returns:
86
+ AuthResult with HeaderCred objects for service account authentication
87
+ """
88
+ # Get OAuth2 access token (cached if still valid)
89
+ access_token = await self._token_client.get_access_token()
90
+
91
+ # Build credentials list using HeaderCred
92
+ credentials: list[Credential] = [
93
+ HeaderCred(name="Authorization", value=SecretStr(f"Bearer {access_token.get_secret_value()}"))
94
+ ]
95
+
96
+ # Add service-specific token if configured
97
+ if self.config.service_token:
98
+ service_header = self.config.service_token.header
99
+ service_token_value = None
100
+
101
+ # Get service token from static config or dynamic function
102
+ if self.config.service_token.token:
103
+ # Static token from config
104
+ service_token_value = self.config.service_token.token.get_secret_value()
105
+
106
+ elif self._service_token_function:
107
+ # Dynamic token from function
108
+ try:
109
+ # Pass configured kwargs to the function
110
+ # Function can access runtime context via AIQContext.get() if needed
111
+ # Handle both sync and async functions
112
+ if asyncio.iscoroutinefunction(self._service_token_function):
113
+ result = await self._service_token_function(**self.config.service_token.kwargs)
114
+ else:
115
+ result = self._service_token_function(**self.config.service_token.kwargs)
116
+
117
+ # Handle function return type: str or tuple[str, str]
118
+ if isinstance(result, tuple):
119
+ service_header, service_token_value = result
120
+ else:
121
+ service_token_value = result
122
+
123
+ logger.debug("Retrieved service token via dynamic function")
124
+
125
+ except Exception as e:
126
+ raise RuntimeError(f"Failed to get service token from function: {e}") from e
127
+
128
+ if service_token_value:
129
+ credentials.append(HeaderCred(name=service_header, value=SecretStr(service_token_value)))
130
+
131
+ # Return AuthResult with HeaderCred objects
132
+ return AuthResult(
133
+ credentials=credentials,
134
+ token_expires_at=self._token_client.token_expires_at,
135
+ raw={},
136
+ )
@@ -0,0 +1,137 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import typing
17
+
18
+ from pydantic import BaseModel
19
+ from pydantic import Field
20
+ from pydantic import field_validator
21
+ from pydantic import model_validator
22
+
23
+ from nat.authentication.interfaces import AuthProviderBaseConfig
24
+ from nat.data_models.common import OptionalSecretStr
25
+ from nat.data_models.common import SerializableSecretStr
26
+
27
+
28
+ class ServiceTokenConfig(BaseModel):
29
+ """
30
+ Configuration for service-specific token in dual authentication patterns.
31
+
32
+ Supports two modes:
33
+
34
+ 1. Static token: Provide token and header directly
35
+ 2. Dynamic function: Provide function path and optional kwargs
36
+
37
+ The function will be called on every request and should have signature::
38
+
39
+ async def get_service_token(**kwargs) -> str | tuple[str, str]
40
+
41
+ If function returns ``tuple[str, str]``, it's interpreted as (header_name, token).
42
+ If function returns ``str``, it's the token and header field is used for header name.
43
+
44
+ The function can access runtime context via AIQContext.get() if needed.
45
+ """
46
+
47
+ # Static token approach
48
+ token: OptionalSecretStr = Field(
49
+ default=None,
50
+ description="Static service token value (mutually exclusive with function)",
51
+ )
52
+
53
+ header: str = Field(
54
+ default="X-Service-Account-Token",
55
+ description="HTTP header name for service token (default: 'X-Service-Account-Token')",
56
+ )
57
+
58
+ # Dynamic function approach
59
+ function: str | None = Field(
60
+ default=None,
61
+ description=("Python function path that returns service token dynamically (mutually exclusive with token). "
62
+ "Function signature: async def func(\\**kwargs) -> str | tuple[str, str]. "
63
+ "Access runtime context via AIQContext.get() if needed."),
64
+ )
65
+
66
+ kwargs: dict[str, typing.Any] = Field(
67
+ default_factory=dict,
68
+ description="Additional keyword arguments to pass to the custom function",
69
+ )
70
+
71
+ @model_validator(mode="after")
72
+ def validate_token_or_function(self):
73
+ """Ensure either token or function is provided, but not both."""
74
+ has_token = self.token is not None
75
+ has_function = self.function is not None
76
+
77
+ if not has_token and not has_function:
78
+ raise ValueError("Either 'token' or 'function' must be provided in service_token config")
79
+
80
+ if has_token and has_function:
81
+ raise ValueError("Cannot specify both 'token' and 'function' in service_token config. Choose one.")
82
+
83
+ return self
84
+
85
+
86
+ class MCPServiceAccountProviderConfig(AuthProviderBaseConfig, name="mcp_service_account"):
87
+ """
88
+ Configuration for MCP service account authentication using OAuth2 client credentials.
89
+
90
+ Generic implementation supporting any OAuth2 client credentials flow.
91
+
92
+ Supports two authentication patterns:
93
+ 1. Single authentication: OAuth2 service account token only
94
+ 2. Dual authentication: OAuth2 service account token + service-specific token
95
+
96
+ Common use cases:
97
+ - Headless/automated MCP workflows
98
+ - CI/CD pipelines
99
+ - Backend services without user interaction
100
+
101
+ All values must be provided via configuration. Use ${ENV_VAR} syntax in YAML
102
+ configs for environment variable substitution.
103
+ """
104
+
105
+ # Required: OAuth2 client credentials
106
+ client_id: str = Field(description="OAuth2 client identifier")
107
+
108
+ client_secret: SerializableSecretStr = Field(description="OAuth2 client secret")
109
+
110
+ # Required: Token endpoint URL
111
+ token_url: str = Field(description="OAuth2 token endpoint URL")
112
+
113
+ # Required: OAuth2 scopes
114
+ scopes: list[str] = Field(description="List of OAuth2 scopes (will be joined with spaces for OAuth2 request)")
115
+
116
+ # Optional: Service-specific token configuration for dual authentication patterns
117
+ service_token: ServiceTokenConfig | None = Field(
118
+ default=None,
119
+ description="Optional service token configuration for dual authentication patterns. "
120
+ "Provide either a static token or a dynamic function that returns the token at runtime.",
121
+ )
122
+
123
+ # Token caching configuration
124
+ token_cache_buffer_seconds: int = Field(default=300,
125
+ description="Seconds before token expiry to refresh (default: 300s/5min)")
126
+
127
+ @field_validator("scopes", mode="before")
128
+ @classmethod
129
+ def validate_scopes(cls, v):
130
+ """
131
+ Accept both list[str] and space-delimited string formats for scopes.
132
+ Converts string to list for consistency.
133
+ """
134
+ if isinstance(v, str):
135
+ # Split space-delimited string into list
136
+ return [scope.strip() for scope in v.split() if scope.strip()]
137
+ return v
@@ -0,0 +1,156 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import asyncio
17
+ import base64
18
+ import logging
19
+ from datetime import datetime
20
+ from datetime import timedelta
21
+
22
+ import httpx
23
+ from pydantic import SecretStr
24
+
25
+ logger = logging.getLogger(__name__)
26
+
27
+
28
+ class ServiceAccountTokenClient:
29
+ """
30
+ Generic OAuth2 client credentials token client for service accounts.
31
+
32
+ Implements standard OAuth2 client credentials flow with token caching.
33
+ """
34
+
35
+ def __init__(
36
+ self,
37
+ client_id: str,
38
+ client_secret: SecretStr,
39
+ token_url: str,
40
+ scopes: str,
41
+ token_cache_buffer_seconds: int = 300,
42
+ ):
43
+ """
44
+ Initialize service account token client.
45
+
46
+ Args:
47
+ client_id: OAuth2 client identifier
48
+ client_secret: OAuth2 client secret (SecretStr)
49
+ token_url: OAuth2 token endpoint URL
50
+ scopes: Space-separated list of scopes
51
+ token_cache_buffer_seconds: Seconds before expiry to refresh (default: 5 min)
52
+ """
53
+ self.client_id = client_id
54
+ self.client_secret = client_secret
55
+ self.token_url = token_url
56
+ self.scopes = scopes
57
+ self.token_cache_buffer_seconds = token_cache_buffer_seconds
58
+
59
+ # Token cache
60
+ self._cached_token: SecretStr | None = None
61
+ self._token_expires_at: datetime | None = None
62
+ self._lock = None # Will be initialized as asyncio.Lock when needed
63
+
64
+ @property
65
+ def token_expires_at(self) -> datetime | None:
66
+ return self._token_expires_at
67
+
68
+ async def _get_lock(self) -> asyncio.Lock:
69
+ """Lazy initialization of asyncio.Lock."""
70
+ if self._lock is None:
71
+ self._lock = asyncio.Lock()
72
+ return self._lock
73
+
74
+ def _is_token_valid(self) -> bool:
75
+ """Check if cached token is still valid (with buffer time)."""
76
+ if not self._cached_token or not self._token_expires_at:
77
+ return False
78
+ buffer = timedelta(seconds=self.token_cache_buffer_seconds)
79
+ return datetime.now() < (self._token_expires_at - buffer)
80
+
81
+ async def get_access_token(self) -> SecretStr:
82
+ """
83
+ Get OAuth2 access token, using cache if valid.
84
+
85
+ Returns:
86
+ Access token as SecretStr
87
+
88
+ Raises:
89
+ RuntimeError: If token acquisition fails
90
+ """
91
+ # Fast path: check cache without lock
92
+ if self._is_token_valid():
93
+ logger.debug("Using cached service account token")
94
+ assert self._cached_token is not None # _is_token_valid() ensures this
95
+ return self._cached_token
96
+
97
+ # Slow path: acquire lock and refresh token
98
+ lock = await self._get_lock()
99
+ async with lock:
100
+ # Double-check after acquiring lock
101
+ if self._is_token_valid():
102
+ logger.debug("Using cached service account token (acquired during lock wait)")
103
+ assert self._cached_token is not None # _is_token_valid() ensures this
104
+ return self._cached_token
105
+
106
+ logger.info("Fetching new service account token")
107
+ return await self._fetch_new_token()
108
+
109
+ async def _fetch_new_token(self) -> SecretStr:
110
+ """
111
+ Fetch a new token from the OAuth2 token endpoint.
112
+
113
+ Returns:
114
+ New access token as SecretStr
115
+
116
+ Raises:
117
+ RuntimeError: If token request fails
118
+ """
119
+ # Encode credentials for Basic authentication
120
+ credentials = f"{self.client_id}:{self.client_secret.get_secret_value()}"
121
+ encoded_credentials = base64.b64encode(credentials.encode()).decode()
122
+
123
+ headers = {"Authorization": f"Basic {encoded_credentials}", "Content-Type": "application/x-www-form-urlencoded"}
124
+
125
+ data = {"grant_type": "client_credentials", "scope": self.scopes}
126
+
127
+ try:
128
+ async with httpx.AsyncClient(timeout=30.0) as client:
129
+ response = await client.post(self.token_url, headers=headers, data=data)
130
+
131
+ if response.status_code == 200:
132
+ token_data = response.json()
133
+
134
+ # Cache the token
135
+ access_token = token_data.get("access_token")
136
+ if not access_token:
137
+ raise RuntimeError("Access token not found in token response")
138
+ self._cached_token = SecretStr(access_token)
139
+ expires_in = token_data.get("expires_in", 3600)
140
+ self._token_expires_at = datetime.now() + timedelta(seconds=expires_in)
141
+
142
+ logger.info("Service account token acquired (expires in %ss)", expires_in)
143
+ return self._cached_token
144
+
145
+ elif response.status_code == 401:
146
+ raise RuntimeError("Invalid service account credentials")
147
+ elif response.status_code == 429:
148
+ raise RuntimeError("Service account rate limit exceeded")
149
+ else:
150
+ raise RuntimeError(
151
+ f"Service account token request failed: {response.status_code} - {response.text}")
152
+
153
+ except httpx.TimeoutException as e:
154
+ raise RuntimeError(f"Service account token request timed out: {e}") from e
155
+ except httpx.RequestError as e:
156
+ raise RuntimeError(f"Service account token request failed: {e}") from e
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -206,7 +206,7 @@ class ObjectStoreTokenStorage(TokenStorageBase):
206
206
 
207
207
  class InMemoryTokenStorage(TokenStorageBase):
208
208
  """
209
- In-memory token storage using NeMo Agent toolkit's built-in object store.
209
+ In-memory token storage using the built-in object store provided by the NeMo Agent toolkit.
210
210
 
211
211
  This implementation uses the in-memory object store for token persistence,
212
212
  which provides a secure default option that doesn't require external storage
@@ -0,0 +1,15 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """MCP CLI commands."""