From 4a0426125e047473ce0f652fddd3c17fb47df15c Mon Sep 17 00:00:00 2001 From: David Gardner Date: Mon, 31 Aug 2026 16:33:17 -0700 Subject: [PATCH 1/6] Improved user identity resolution * Using a header is now an opt-in, disabled by default, allowing the user to configure the header name. * Document situations where using a header is and isn't secure * Remove fallback to default_user in auto memory wrapper * Borrow the JWT improvements from #2196 to apply to all endpoints not just websockets Signed-off-by: David Gardner --- docs/source/build-workflows/memory.md | 24 +++++-- .../auto-memory-wrapper.md | 48 ++++++++++--- docs/source/reference/rest-api/websockets.md | 16 +++++ examples/agents/auto_memory_wrapper/README.md | 35 ++++++--- .../http_basic_auth_provider.py | 12 ++-- .../oauth2/oauth2_auth_code_flow_provider.py | 12 ++-- .../src/nat/data_models/user_info.py | 7 ++ .../fastapi/fastapi_front_end_config.py | 25 +++++++ .../fastapi_front_end_plugin_worker.py | 6 ++ .../nat/front_ends/fastapi/message_handler.py | 18 ++++- .../front_ends/fastapi/routes/websocket.py | 1 + .../src/nat/runtime/session.py | 14 ++-- .../src/nat/runtime/user_manager.py | 39 ++++++++++ .../test_http_basic_auth_exchanger.py | 20 +++++- .../authentication/test_oauth_exchanger.py | 16 +++++ .../fastapi/test_fastapi_front_end_config.py | 24 +++++++ .../fastapi/test_message_handler.py | 37 ++++++++++ .../tests/nat/runtime/test_user_manager.py | 71 ++++++++++++++++++- .../agent/auto_memory_wrapper/agent.py | 57 ++++----------- .../agent/auto_memory_wrapper/register.py | 6 +- .../agent/auto_memory_wrapper/state.py | 1 + .../tests/agent/test_auto_memory_wrapper.py | 47 +++++++----- .../nat/plugins/mcp/client/fastapi_routes.py | 6 +- .../tests/server/test_mcp_client_endpoint.py | 3 +- 24 files changed, 428 insertions(+), 117 deletions(-) diff --git a/docs/source/build-workflows/memory.md b/docs/source/build-workflows/memory.md index 575d25e5cb..546ffaa55a 100644 --- a/docs/source/build-workflows/memory.md +++ b/docs/source/build-workflows/memory.md @@ -141,15 +141,29 @@ The automatic memory wrapper agent supports several configuration parameters: ### Multi-Tenant Memory Isolation -User ID is automatically extracted at runtime for memory isolation via: -1. `SessionManager.session(user_id=...)` - For production with custom auth middleware (recommended) -2. `X-User-ID` HTTP header - Illustrative/testing only; assumes a trusted upstream proxy authenticates the request and injects this header -3. Console front end `user_id` - Defaults to `"nat_run_user_id"` for `nat run` +The automatic memory wrapper reads only the identity resolved by the runtime session. Resolve identity through +authenticated front-end credentials, `SessionManager.session(user_id=...)`, or the console front end `user_id` (which +defaults to `"nat_run_user_id"` for `nat run`). Memory operations fail closed when no identity is available. + +For local testing or an isolated deployment behind an authenticating reverse proxy, you can explicitly opt in to a +trusted upstream identity header: + +```yaml +general: + front_end: + _type: fastapi + identity_header: X-User-ID +``` + +Do not enable this setting unless clients cannot reach `nat serve` directly, the proxy authenticates every request and +overwrites the header, the toolkit port is not published outside the trusted backend network, and every container on +that network is trusted. A client-supplied identity header is not authentication. Conversation-aware memory backends can also use `conversation_id` to isolate separate conversations for the same user. For `nat run`, pass `--conversation_id` when testing independent memory conversations from the CLI. -> Never treat a client-supplied `X-User-ID` header as authentication. +When configured, the header is authoritative for HTTP and WebSocket requests. Missing, empty, and repeated values are +rejected, and other credentials cannot override it. For detailed configuration and usage examples, refer to the `examples/agents/auto_memory_wrapper/README.md` guide. diff --git a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md index 120933939e..0a255e587a 100644 --- a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md +++ b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md @@ -163,11 +163,14 @@ The wrapper automatically: The automatic memory wrapper agent provides multi-tenant support through runtime user ID extraction. Configure user IDs through the front end or session runtime, not the `auto_memory_agent` workflow block. -### User ID Extraction Priority +### User ID Resolution -1. **`SessionManager.session(user_id=...)`** - For production with custom auth middleware (recommended) -2. **`X-User-ID` HTTP header** - Illustrative/testing only; assumes a hypothetical trusted upstream proxy authenticates the request and injects the header -3. **Console front end `user_id`** - Defaults to `"nat_run_user_id"` for `nat run` +The wrapper reads only the user ID resolved by the runtime session. It does not read request headers directly and does +not use a shared fallback identity. If a memory operation is enabled and the runtime has not resolved an identity, the +request fails instead of sharing memory across unauthenticated users. + +Resolve identity with authenticated front-end credentials, `SessionManager.session(user_id=...)`, or the console front +end `user_id` (which defaults to `"nat_run_user_id"` for `nat run`). Conversation-aware memory backends can also use `conversation_id` to isolate separate conversations for the same user. For Zep Cloud, if no conversation ID is supplied, the integration uses a deterministic per-user default thread. @@ -196,9 +199,32 @@ async def handle_request(request): return result ``` -### Testing: X-User-ID Header +### Local Testing or a Trusted Upstream Identity Header + +The FastAPI front end can explicitly trust one upstream identity header: + +```yaml +general: + front_end: + _type: fastapi + identity_header: X-User-ID +``` + +Use this setting only for local testing or behind an authenticating reverse proxy. It is secure only when all of the +following are true: + +- Clients cannot connect to `nat serve` directly; only the trusted proxy can reach its listening port. +- The proxy authenticates every request and overwrites `X-User-ID` with the authenticated principal. It must not + preserve, append, or forward a client-supplied value. +- The backend Docker network and every container attached to it are trusted. Do not publish the `nat serve` port to the + host; publish only the proxy port. +- External traffic is protected as appropriate, typically with TLS at the proxy. + +When configured, the header is authoritative for HTTP and WebSocket requests. A missing, empty, or repeated header is +rejected. Other connection credentials and WebSocket auth messages cannot replace it. The setting cannot be combined +with `accepted_identity_credentials` or `identity_authentication`. -For quick testing without custom middleware: +After configuring the trusted boundary, a request forwarded by the proxy can contain: ```bash curl -X POST http://localhost:8000/chat \ @@ -327,14 +353,14 @@ workflow: ## Important Notes -1. **User ID is runtime/front-end scoped** - Set via `SessionManager.session(user_id=...)`, `X-User-ID`, or - `nat run --user_id`. When no runtime identity or header is available, the wrapper falls back to `"default_user"`. - Use this fallback for development and testing only. Production deployments must require an authenticated runtime - identity. +1. **User ID is runtime/front-end scoped** - Set through authenticated front-end identity resolution, + `SessionManager.session(user_id=...)`, an explicitly configured trusted `identity_header`, or `nat run --user_id`. + Memory operations fail when no identity is available. 2. **Memory backends are interchangeable** - Works with any implementation of `MemoryEditor` interface 3. **No memory tools needed** - The wrapped agent does not need explicit memory tools configured 4. **Transparent to inner agent** - The wrapped agent is unaware of memory operations -5. **X-User-ID** - We used this header for illustrative purposes only. Do not rely on it for authentication in production. +5. **Trusted identity headers are opt-in** - Never enable `identity_header` on a server that untrusted clients or + untrusted containers can reach directly. --- diff --git a/docs/source/reference/rest-api/websockets.md b/docs/source/reference/rest-api/websockets.md index 7c0b8516a5..fb1ffeaa70 100644 --- a/docs/source/reference/rest-api/websockets.md +++ b/docs/source/reference/rest-api/websockets.md @@ -112,6 +112,22 @@ general: Each JWT provider requires `issuer_url`, `jwks_uri`, and `audience`. It can also require `scopes` and configure `timeout` and `leeway`. Add another named provider and select it in `identity_authentication` to accept JWT tokens from another issuer. The issuer claim selects the matching provider; an unknown issuer or a failed signature, time, audience, or scope check returns an authentication error and does not restore workflow state. If `identity_authentication` is omitted, JWT tokens retain the existing decode-only behavior. Configuring `identity_authentication` while excluding `jwt` from `accepted_identity_credentials` is invalid. +For local testing or a deployment where an authenticating reverse proxy is the only service that can reach +`nat serve`, the FastAPI front end can instead trust an upstream identity header: + +```yaml +general: + front_end: + _type: fastapi + identity_header: X-User-ID +``` + +This setting is secure only if the proxy authenticates every connection, overwrites any client-supplied value, and the +toolkit port is not exposed outside a trusted backend network. Every container that can reach that network must also +be trusted. When enabled, the header is authoritative: the connection is rejected if it is missing, empty, or repeated, +and an `auth_message` cannot replace the resolved identity. `identity_header` cannot be combined with +`accepted_identity_credentials` or `identity_authentication`. + ## Auth Message This message allows clients to authenticate over a WebSocket connection when header-based or cookie-based authentication is not feasible (e.g., browser WebSocket APIs that do not support custom headers). diff --git a/examples/agents/auto_memory_wrapper/README.md b/examples/agents/auto_memory_wrapper/README.md index aece1e59b0..4481731b77 100644 --- a/examples/agents/auto_memory_wrapper/README.md +++ b/examples/agents/auto_memory_wrapper/README.md @@ -130,11 +130,11 @@ See `config_zep.yml` for comprehensive parameter examples. User ID is extracted at runtime for memory isolation. Configure it through the front end or session runtime, not the `auto_memory_agent` workflow block. -### User ID Extraction Priority +### User ID Resolution -1. **`SessionManager.session(user_id=...)`** - For production with custom auth middleware (recommended) -2. **`X-User-ID` HTTP header** - Illustrative/testing only; assumes a hypothetical trusted upstream proxy authenticates the request and injects the header -3. **Console front end `user_id`** - Defaults to `"nat_run_user_id"` for `nat run` +The wrapper reads only the identity resolved by the runtime session. Use authenticated front-end credentials, +`SessionManager.session(user_id=...)`, or the console front end `user_id` (which defaults to `"nat_run_user_id"` for +`nat run`). Memory operations fail closed when no identity is available. Conversation-aware memory backends can also use `conversation_id` to isolate separate conversations for the same user. For Zep Cloud, if no conversation ID is supplied, the integration uses a deterministic per-user default thread. @@ -163,9 +163,23 @@ async def handle_request(request): return result ``` -### Testing: X-User-ID Header +### Local Testing or a Trusted Upstream Identity Header -For quick testing without custom middleware: +To explicitly trust an identity header, configure the FastAPI front end: + +```yaml +general: + front_end: + _type: fastapi + identity_header: X-User-ID +``` + +Use this only for local testing or when `nat serve` is isolated behind an authenticating reverse proxy. Clients must +not be able to reach the toolkit service directly. The proxy must overwrite the header after authentication, the +toolkit port must not be published outside the trusted backend network, and every container on that network must be +trusted. + +After configuring that trust boundary: ```bash curl -X POST http://localhost:8000/chat \ @@ -175,7 +189,8 @@ curl -X POST http://localhost:8000/chat \ -d '{"messages": [{"role": "user", "content": "Hello!"}]}' ``` -The example usage of the `X-User-ID` header is for illustrative purposes only; do not accept this header directly from untrusted clients. +Never accept this header directly from untrusted clients. When `identity_header` is configured, it is authoritative; +missing, empty, or repeated values are rejected and other credentials cannot override it. ### Local Development: Console User and Conversation IDs @@ -216,9 +231,11 @@ workflow: ## Important Notes -1. **User ID is runtime/front-end scoped** - Set via `SessionManager.session(user_id=...)`, `X-User-ID`, or `nat run --user_id` +1. **User ID is runtime/front-end scoped** - Set through authenticated front-end identity resolution, + `SessionManager.session(user_id=...)`, an explicitly configured trusted `identity_header`, or `nat run --user_id` 2. **Memory backends are interchangeable** - Works with any implementation of `MemoryEditor` interface -3. `X-User-ID` HTTP header - Illustrative/testing only; assumes a trusted upstream proxy authenticates the request and injects this header +3. **Trusted headers require network enforcement** - Never expose a trusted identity-header deployment directly to + untrusted clients or containers ## Examples diff --git a/packages/nvidia_nat_core/src/nat/authentication/http_basic_auth/http_basic_auth_provider.py b/packages/nvidia_nat_core/src/nat/authentication/http_basic_auth/http_basic_auth_provider.py index dc37784590..df450c3e8f 100644 --- a/packages/nvidia_nat_core/src/nat/authentication/http_basic_auth/http_basic_auth_provider.py +++ b/packages/nvidia_nat_core/src/nat/authentication/http_basic_auth/http_basic_auth_provider.py @@ -23,7 +23,6 @@ from nat.data_models.authentication import AuthResult from nat.data_models.authentication import BasicAuthCred from nat.data_models.authentication import BearerTokenCred -from nat.runtime.session import SESSION_COOKIE_NAME class HTTPBasicAuthProvider(AuthProviderBase): @@ -46,13 +45,10 @@ async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult context = Context.get() - if user_id is None and hasattr(context, "metadata") and hasattr( - context.metadata, "cookies") and context.metadata.cookies is not None: - session_id = context.metadata.cookies.get(SESSION_COOKIE_NAME, None) - if not session_id: - raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.") - - user_id = session_id + if user_id is None: + user_id = context.user_id + if not user_id: + raise RuntimeError("Authentication failed. No resolved user identity is available.") if user_id and user_id in self._authenticated_tokens: return self._authenticated_tokens[user_id] diff --git a/packages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider.py b/packages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider.py index 2979046e0f..0de9cfbd5f 100644 --- a/packages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +++ b/packages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider.py @@ -31,7 +31,6 @@ from nat.data_models.authentication import AuthFlowType from nat.data_models.authentication import AuthResult from nat.data_models.authentication import BearerTokenCred -from nat.runtime.session import SESSION_COOKIE_NAME logger = logging.getLogger(__name__) @@ -91,13 +90,10 @@ def _set_custom_auth_callback(self, async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult: context = Context.get() - if user_id is None and hasattr(context, "metadata") and hasattr( - context.metadata, "cookies") and context.metadata.cookies is not None: - session_id = context.metadata.cookies.get(SESSION_COOKIE_NAME, None) - if not session_id: - raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.") - - user_id = session_id + if user_id is None: + user_id = context.user_id + if not user_id: + raise RuntimeError("Authentication failed. No resolved user identity is available.") if user_id: # Try to retrieve from token storage diff --git a/packages/nvidia_nat_core/src/nat/data_models/user_info.py b/packages/nvidia_nat_core/src/nat/data_models/user_info.py index 52313d2ea9..b996a70e2c 100644 --- a/packages/nvidia_nat_core/src/nat/data_models/user_info.py +++ b/packages/nvidia_nat_core/src/nat/data_models/user_info.py @@ -173,6 +173,13 @@ def _from_session_cookie(cls, cookie: str) -> "UserInfo": def _from_api_key(cls, api_key: str) -> "UserInfo": return cls(api_key=SecretStr(api_key)) + @classmethod + def _from_identity_header(cls, header_name: str, header_value: str) -> "UserInfo": + """Create a user from an identity asserted by a trusted upstream proxy.""" + instance: UserInfo = cls() + instance._set_user_id(f"trusted-header:{header_name.lower()}\x1f{header_value}") + return instance + @classmethod def _from_jwt(cls, jwt_info: JwtUserInfo, *, issuer_scoped: bool = False) -> "UserInfo": identity: str | None = jwt_info.identity_claim diff --git a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_config.py b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_config.py index ffcd4f387b..3920003f2d 100644 --- a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_config.py +++ b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_config.py @@ -15,6 +15,7 @@ import logging import os +import re import sys import typing from datetime import datetime @@ -295,8 +296,32 @@ class CrossOriginResourceSharing(BaseModel): "If omitted, JWT claims retain the existing decode-only behavior."), ) + identity_header: str | None = Field( + default=None, + description=("Name of an HTTP header containing an identity asserted by a trusted upstream proxy. " + "This is disabled by default and must only be enabled when untrusted clients cannot reach " + "the server directly and the proxy removes any client-supplied value before setting it."), + ) + + @field_validator("identity_header") + @classmethod + def validate_identity_header(cls, identity_header: str | None) -> str | None: + if identity_header is None: + return None + identity_header = identity_header.strip() + if not identity_header: + raise ValueError("identity_header must not be empty") + if re.fullmatch(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+", identity_header) is None: + raise ValueError("identity_header must be a valid HTTP header name") + return identity_header + @model_validator(mode="after") def validate_jwt_identity_policy(self) -> typing.Self: + if self.identity_header is not None: + if self.accepted_identity_credentials is not None: + raise ValueError("identity_header cannot be combined with accepted_identity_credentials") + if self.identity_authentication: + raise ValueError("identity_header cannot be combined with identity_authentication") if (self.identity_authentication and self.accepted_identity_credentials is not None and "jwt" not in self.accepted_identity_credentials): raise ValueError( diff --git a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py index 7ee9b9eab3..9fc03d9879 100644 --- a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +++ b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py @@ -27,12 +27,14 @@ from fastapi import Response from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel +from starlette.responses import JSONResponse from nat.builder.evaluator import EvaluatorInfo from nat.builder.workflow_builder import WorkflowBuilder from nat.builder.workflow_builder import WorkflowEvalBuilderBase from nat.data_models.config import Config from nat.runtime.session import SessionManager +from nat.runtime.user_manager import IdentityHeaderError from nat.utils.log_utils import setup_logging from .auth_flow_handlers.http_flow_handler import HTTPAuthenticationFlowHandler @@ -126,6 +128,10 @@ async def lifespan(starting_app: FastAPI): nat_app = FastAPI(lifespan=lifespan) + @nat_app.exception_handler(IdentityHeaderError) + async def identity_header_error_handler(_request: Request, exc: IdentityHeaderError) -> JSONResponse: + return JSONResponse(status_code=401, content={"detail": str(exc)}) + # Configure app CORS. self.set_cors_config(nat_app) diff --git a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py index 7a1d42fb34..07488c180a 100644 --- a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py +++ b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/message_handler.py @@ -66,6 +66,7 @@ from nat.front_ends.fastapi.step_adaptor import StepAdaptor from nat.runtime.session import SessionManager from nat.runtime.user_manager import IdentityCredentialNotAcceptedError +from nat.runtime.user_manager import IdentityHeaderError from nat.runtime.user_manager import JwtVerificationError from nat.runtime.user_manager import UserManager @@ -96,6 +97,7 @@ def __init__( worker: "FastApiFrontEndPluginWorker", accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, jwt_validators: typing.Mapping[str, BearerTokenValidator] | None = None, + identity_header: str | None = None, ): self._socket: WebSocket = socket self._session_manager: SessionManager = session_manager @@ -103,6 +105,7 @@ def __init__( self._worker: FastApiFrontEndPluginWorker = worker self._accepted_identity_credentials = accepted_identity_credentials self._jwt_validators = jwt_validators + self._identity_header = identity_header self._message_validator: MessageValidator = MessageValidator() self._running_workflow_task: asyncio.Task | None = None @@ -192,8 +195,9 @@ async def __aenter__(self) -> "WebSocketMessageHandler": self._socket, accepted_identity_credentials=self._accepted_identity_credentials, jwt_validators=self._jwt_validators, + identity_header=self._identity_header, ) - except (IdentityCredentialNotAcceptedError, JwtVerificationError) as exc: + except (IdentityCredentialNotAcceptedError, IdentityHeaderError, JwtVerificationError) as exc: self._connection_rejected = True response = WebSocketAuthResponseMessage( status=AuthMessageStatus.ERROR, @@ -313,6 +317,18 @@ async def _process_auth_message(self, message: WebSocketAuthMessage) -> None: self._flow_handler.set_oauth_mode(message.payload.mode) return + if self._identity_header is not None: + response = WebSocketAuthResponseMessage( + status=AuthMessageStatus.ERROR, + payload=Error( + code=ErrorTypes.USER_AUTH_ERROR, + message="Authentication failed", + details="WebSocket auth messages cannot replace an identity asserted by a trusted header", + ), + ) + await self._socket.send_json(response.model_dump()) + return + identity_resolved = False try: user_info: UserInfo = await UserManager.from_auth_payload_with_verification( diff --git a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.py b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.py index 8e5828950e..86f733c92c 100644 --- a/packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.py +++ b/packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.py @@ -114,6 +114,7 @@ async def _websocket_endpoint(websocket: WebSocket): worker, accepted_identity_credentials=worker.front_end_config.accepted_identity_credentials, jwt_validators=jwt_validators, + identity_header=worker.front_end_config.identity_header, ) as handler: origin = websocket.headers.get("origin") allowed_origins = worker.front_end_config.cors.allow_origins or [] diff --git a/packages/nvidia_nat_core/src/nat/runtime/session.py b/packages/nvidia_nat_core/src/nat/runtime/session.py index 8ca5a3ca13..fc8aa2c380 100644 --- a/packages/nvidia_nat_core/src/nat/runtime/session.py +++ b/packages/nvidia_nat_core/src/nat/runtime/session.py @@ -210,6 +210,9 @@ def __init__(self, from nat.cli.type_registry import GlobalTypeRegistry self._config = config + front_end = getattr(config.general, "front_end", None) + identity_header = getattr(front_end, "identity_header", None) + self._identity_header = identity_header if isinstance(identity_header, str) else None self._max_concurrency = max_concurrency self._entry_function = entry_function @@ -472,6 +475,7 @@ async def session(self, builder_info: PerUserBuilderInfo | None = None request_start_time: float | None = None request_success = True + identity_header = getattr(self, "_identity_header", None) try: if user_input_callback is not None: @@ -487,15 +491,17 @@ async def session(self, token_user_message_id = self._context_state.user_message_id.set(user_message_id) if isinstance(http_connection, WebSocket): - if user_id is None: - user_info: UserInfo | None = UserManager.extract_user_from_connection(http_connection) + if identity_header is not None or user_id is None: + user_info: UserInfo | None = UserManager.extract_user_from_connection( + http_connection, identity_header=identity_header) if user_info is not None: user_id = user_info.get_user_id() self.set_metadata_from_websocket(http_connection, user_message_id, conversation_id) if isinstance(http_connection, Request): - if user_id is None: - user_info = UserManager.extract_user_from_connection(http_connection) + if identity_header is not None or user_id is None: + user_info = UserManager.extract_user_from_connection( + http_connection, identity_header=identity_header) if user_info is not None: user_id = user_info.get_user_id() token_workflow_parent_id, token_workflow_parent_name = \ diff --git a/packages/nvidia_nat_core/src/nat/runtime/user_manager.py b/packages/nvidia_nat_core/src/nat/runtime/user_manager.py index 2b1d0a4d47..15928c2f6c 100644 --- a/packages/nvidia_nat_core/src/nat/runtime/user_manager.py +++ b/packages/nvidia_nat_core/src/nat/runtime/user_manager.py @@ -43,6 +43,10 @@ class IdentityCredentialNotAcceptedError(ValueError): """Raised when a supplied identity credential method is disabled by policy.""" +class IdentityHeaderError(ValueError): + """Raised when a configured trusted identity header is missing or ambiguous.""" + + class JwtVerificationError(ValueError): """Raised when an enabled JWT verification policy rejects a token.""" @@ -55,6 +59,7 @@ def extract_user_from_connection( cls, connection: Request | WebSocket, accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + identity_header: str | None = None, ) -> UserInfo | None: """Resolve an HTTP/WebSocket connection into a ``UserInfo``. @@ -69,6 +74,10 @@ def extract_user_from_connection( ValueError: If a credential is found but cannot be resolved to a valid user identity. """ + if identity_header is not None: + identity = cls._get_identity_header(connection, identity_header) + return UserInfo._from_identity_header(identity_header, identity) + cookie: str | None = cls._get_session_cookie(connection) if cookie: cls._ensure_identity_credential_accepted("session_cookie", accepted_identity_credentials) @@ -93,7 +102,12 @@ async def extract_user_from_connection_with_verification( connection: Request | WebSocket, accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, jwt_validators: typing.Mapping[str, BearerTokenValidator] | None = None, + identity_header: str | None = None, ) -> UserInfo | None: + if identity_header is not None: + identity = cls._get_identity_header(connection, identity_header) + return UserInfo._from_identity_header(identity_header, identity) + cookie = cls._get_session_cookie(connection) if cookie: cls._ensure_identity_credential_accepted("session_cookie", accepted_identity_credentials) @@ -271,6 +285,31 @@ def _ensure_identity_credential_accepted( if accepted_identity_credentials is not None and credential_type not in accepted_identity_credentials: raise IdentityCredentialNotAcceptedError(f"Identity credential type '{credential_type}' is not accepted") + @staticmethod + def _get_identity_header(connection: Request | WebSocket, header_name: str) -> str: + """Read one non-empty value for a trusted identity header, failing closed otherwise.""" + values: list[str] = [] + if isinstance(connection, Request): + values = list(connection.headers.getlist(header_name)) + elif isinstance(connection, WebSocket) and hasattr(connection, "scope"): + target = header_name.lower() + for name, value in connection.scope.get("headers", []): + try: + if name.decode("latin-1").lower() == target: + values.append(value.decode("latin-1")) + except (AttributeError, UnicodeDecodeError): + raise IdentityHeaderError(f"Configured identity header '{header_name}' is malformed") from None + + if not values: + raise IdentityHeaderError(f"Configured identity header '{header_name}' is missing") + if len(values) != 1: + raise IdentityHeaderError(f"Configured identity header '{header_name}' must occur exactly once") + + identity = values[0].strip() + if not identity: + raise IdentityHeaderError(f"Configured identity header '{header_name}' must not be empty") + return identity + @staticmethod def _get_session_cookie(connection: Request | WebSocket) -> str | None: """Extract the ``nat-session`` cookie value from a Request or WebSocket.""" diff --git a/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py b/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py index 9343aedab5..3c97b1fd30 100644 --- a/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py +++ b/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py @@ -28,13 +28,14 @@ # --------------------------------------------------------------------------- # -def _patch_context(monkeypatch: pytest.MonkeyPatch, callback): +def _patch_context(monkeypatch: pytest.MonkeyPatch, callback, user_id: str | None = None): """Replace Context.get() so the exchanger sees *our* callback.""" class _DummyCtx: def __init__(self, cb): self.user_auth_callback = cb + self.user_id = user_id monkeypatch.setattr(Context, "get", staticmethod(lambda: _DummyCtx(callback)), raising=True) @@ -94,6 +95,23 @@ async def cb(cfg, flow): # noqa: D401 assert hits["n"] == 1 +async def test_uses_resolved_context_identity(monkeypatch): + """Authentication caching uses the authoritative runtime identity when no ID is passed.""" + + async def cb(cfg, flow): + return AuthenticatedContext( + headers={"Authorization": "Basic YQ=="}, + metadata={"username": "a", "password": "b"}, + ) + + _patch_context(monkeypatch, cb, user_id="resolved-user") + exchanger = HTTPBasicAuthProvider(HTTPBasicAuthProviderConfig()) + + await exchanger.authenticate() + + assert "resolved-user" in exchanger._authenticated_tokens + + async def test_missing_authorization_header(monkeypatch): """Callback returns no `Authorization` header → RuntimeError.""" diff --git a/packages/nvidia_nat_core/tests/nat/authentication/test_oauth_exchanger.py b/packages/nvidia_nat_core/tests/nat/authentication/test_oauth_exchanger.py index 4f34a4ad64..52b0fad5fc 100644 --- a/packages/nvidia_nat_core/tests/nat/authentication/test_oauth_exchanger.py +++ b/packages/nvidia_nat_core/tests/nat/authentication/test_oauth_exchanger.py @@ -37,12 +37,14 @@ def _patch_context( monkeypatch: pytest.MonkeyPatch, callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType], Awaitable[AuthenticatedContext]], + user_id: str | None = None, ) -> None: class _DummyCtx: def __init__(self, cb): self.user_auth_callback = cb + self.user_id = user_id monkeypatch.setattr(Context, "get", staticmethod(lambda: _DummyCtx(callback)), raising=True) @@ -138,6 +140,20 @@ async def cb(conf, flow): assert calls["n"] == 1 +async def test_authenticate_uses_resolved_context_identity(monkeypatch, cfg): + """Token storage uses the authoritative runtime identity when no ID is passed.""" + + async def cb(conf, flow): + return _bearer_ctx(token="tok", expires_at=datetime.now(UTC) + timedelta(minutes=10)) + + _patch_context(monkeypatch, cb, user_id="resolved-user") + client = OAuth2AuthCodeFlowProvider(cfg) + + await client.authenticate() + + assert await client._token_storage.retrieve("resolved-user") is not None + + # --------------------------------------------------------------------------- # # 4. Token refresh succeeds # --------------------------------------------------------------------------- # diff --git a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_fastapi_front_end_config.py b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_fastapi_front_end_config.py index 95eeea626b..4406f94008 100644 --- a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_fastapi_front_end_config.py +++ b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_fastapi_front_end_config.py @@ -195,3 +195,27 @@ def test_identity_authentication_rejects_policy_that_disables_jwt(): accepted_identity_credentials=["session_cookie"], identity_authentication=["corporate_jwt"], ) + + +def test_identity_header_is_disabled_by_default(): + assert FastApiFrontEndConfig().identity_header is None + + +def test_identity_header_accepts_and_trims_valid_header_name(): + assert FastApiFrontEndConfig(identity_header=" X-User-ID ").identity_header == "X-User-ID" + + +@pytest.mark.parametrize("header_name", ["", " ", "X User ID", "X-User-ID\r\nInjected"]) +def test_identity_header_rejects_invalid_header_name(header_name): + with pytest.raises(ValueError, match="identity_header"): + FastApiFrontEndConfig(identity_header=header_name) + + +def test_identity_header_cannot_be_combined_with_credential_policy(): + with pytest.raises(ValueError, match="identity_header cannot be combined"): + FastApiFrontEndConfig(identity_header="X-User-ID", accepted_identity_credentials=[]) + + +def test_identity_header_cannot_be_combined_with_jwt_authentication(): + with pytest.raises(ValueError, match="identity_header cannot be combined"): + FastApiFrontEndConfig(identity_header="X-User-ID", identity_authentication=["corporate_jwt"]) diff --git a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py index 9709a88e21..fe97e79ff4 100644 --- a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py +++ b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py @@ -45,6 +45,7 @@ def _make_jwt(claims: dict) -> str: def _make_message_handler( accepted_identity_credentials=None, jwt_validators=None, + identity_header=None, ) -> tuple[WebSocketMessageHandler, MagicMock, WebSocketAuthenticationFlowHandler]: """Build a WebSocketMessageHandler with a mockable socket and a real flow handler.""" socket = MagicMock(spec=WebSocket) @@ -58,6 +59,7 @@ def _make_message_handler( worker=MagicMock(), accepted_identity_credentials=accepted_identity_credentials, jwt_validators=jwt_validators, + identity_header=identity_header, ) flow_handler = WebSocketAuthenticationFlowHandler( add_flow_cb=AsyncMock(), @@ -87,6 +89,41 @@ async def test_context_manager_resolves_connection_identity_before_restoration() restore.assert_awaited_once() +async def test_context_manager_uses_configured_identity_header(): + """The configured identity header is passed through the verified connection resolver.""" + handler, socket, _ = _make_message_handler(identity_header="X-User-ID") + user_info = MagicMock() + user_info.get_user_id.return_value = "resolved-user" + + with patch( + "nat.front_ends.fastapi.message_handler.UserManager.extract_user_from_connection_with_verification", + return_value=user_info, + ) as resolver: + await handler.__aenter__() + + assert resolver.await_args.kwargs["identity_header"] == "X-User-ID" + assert handler._user_id == "resolved-user" + + +async def test_auth_message_cannot_replace_trusted_header_identity(): + """A WebSocket auth message cannot override an upstream-asserted identity.""" + handler, socket, _ = _make_message_handler(identity_header="X-User-ID") + handler._user_id = "resolved-user" + msg = WebSocketAuthMessage( + type=WebSocketMessageType.AUTH_MESSAGE, + payload=ApiKeyAuthPayload(method="api_key", token="replacement-key"), + ) + + with patch( + "nat.front_ends.fastapi.message_handler.UserManager.from_auth_payload_with_verification", + ) as resolver: + await handler._process_auth_message(msg) + + resolver.assert_not_called() + assert handler._user_id == "resolved-user" + assert socket.send_json.await_args.args[0]["status"] == AuthMessageStatus.ERROR + + async def test_context_manager_rejects_disabled_connection_credential_without_restoration(): """A disabled upgrade credential closes the socket and never attempts state restoration.""" handler, socket, _ = _make_message_handler(accepted_identity_credentials=["jwt"]) diff --git a/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py b/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py index 17c22e3c56..226bbd9487 100644 --- a/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py +++ b/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py @@ -34,6 +34,7 @@ from nat.data_models.user_info import UserInfo from nat.runtime.session import SESSION_COOKIE_NAME from nat.runtime.user_manager import IdentityCredentialNotAcceptedError +from nat.runtime.user_manager import IdentityHeaderError from nat.runtime.user_manager import JwtVerificationError from nat.runtime.user_manager import UserManager @@ -50,7 +51,9 @@ def _mock_request(cookies: dict[str, str] | None = None, headers: dict[str, str] mock = MagicMock(spec=Request) mock.cookies = cookies or {} mock.headers = MagicMock() - mock.headers.get = (headers or {}).get + header_values = headers or {} + mock.headers.get = header_values.get + mock.headers.getlist = lambda name: [value for key, value in header_values.items() if key.lower() == name.lower()] return mock @@ -58,6 +61,7 @@ def _mock_websocket( cookie_header: str | None = None, auth_header: str | None = None, api_key_header: str | None = None, + extra_headers: list[tuple[bytes, bytes]] | None = None, ) -> MagicMock: """Create a MagicMock that passes ``isinstance(obj, WebSocket)``.""" raw_headers: list[tuple[bytes, bytes]] = [] @@ -67,12 +71,62 @@ def _mock_websocket( raw_headers.append((b"authorization", auth_header.encode())) if api_key_header: raw_headers.append((b"x-api-key", api_key_header.encode())) + raw_headers.extend(extra_headers or []) mock = MagicMock(spec=WebSocket) mock.scope = {"headers": raw_headers} return mock +class TestTrustedIdentityHeader: + """A configured upstream identity header is authoritative and fails closed.""" + + def test_request_identity_header_returns_deterministic_user(self): + request = _mock_request(headers={"X-User-ID": "alice"}) + + first = UserManager.extract_user_from_connection(request, identity_header="X-User-ID") + second = UserManager.extract_user_from_connection(request, identity_header="x-user-id") + + assert first is not None + assert second is not None + assert first.get_user_id() == second.get_user_id() + assert first.get_user_details() is None + + def test_websocket_identity_header_is_supported(self): + websocket = _mock_websocket(extra_headers=[(b"x-user-id", b"alice")]) + + info = UserManager.extract_user_from_connection(websocket, identity_header="X-User-ID") + + assert info is not None + assert info.get_user_id() + + @pytest.mark.parametrize("value", [None, "", " "]) + def test_missing_or_empty_identity_header_is_rejected(self, value): + headers = {} if value is None else {"X-User-ID": value} + request = _mock_request(headers=headers) + + with pytest.raises(IdentityHeaderError): + UserManager.extract_user_from_connection(request, identity_header="X-User-ID") + + def test_duplicate_identity_header_is_rejected(self): + websocket = _mock_websocket(extra_headers=[(b"x-user-id", b"alice"), (b"X-User-ID", b"mallory")]) + + with pytest.raises(IdentityHeaderError, match="exactly once"): + UserManager.extract_user_from_connection(websocket, identity_header="X-User-ID") + + def test_identity_header_takes_precedence_over_other_credentials(self): + request = _mock_request( + cookies={SESSION_COOKIE_NAME: "cookie-user"}, + headers={"X-User-ID": "header-user", "authorization": "Bearer api-key"}, + ) + + header_info = UserManager.extract_user_from_connection(request, identity_header="X-User-ID") + expected = UserInfo._from_identity_header("X-User-ID", "header-user") + + assert header_info is not None + assert header_info.get_user_id() == expected.get_user_id() + + class TestFromConnectionRequestCookie: """extract_user_from_connection resolves a UserInfo from a session cookie on an HTTP Request.""" @@ -896,6 +950,21 @@ async def test_user_id_provided_skips_extraction(self): assert session._user_id == "explicit-id" mock_extract.assert_not_called() + async def test_identity_header_overrides_explicit_user_id_for_connection(self): + """A caller cannot bypass configured header identity with an explicit user_id.""" + from unittest.mock import patch + + sm = self._make_session_manager() + sm._identity_header = "X-User-ID" + ws = _mock_websocket(extra_headers=[(b"x-user-id", b"proxy-user")]) + header_info = UserInfo._from_identity_header("X-User-ID", "proxy-user") + + with patch.object(UserManager, "extract_user_from_connection", return_value=header_info) as resolver: + async with sm.session(user_id="caller-controlled", http_connection=ws) as session: + assert session._user_id == header_info.get_user_id() + + resolver.assert_called_once_with(ws, identity_header="X-User-ID") + async def test_websocket_cookie_sets_user_id_in_context(self): """Input: WebSocket with session cookie. Asserts session user_id matches cookie-derived UUID.""" from unittest.mock import patch diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py index 62562b59a5..f853e61a88 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py @@ -68,53 +68,20 @@ def __init__( self._context = Context.get() def _get_user_id_from_context(self) -> str: - """ - Extract user_id from runtime context. - - Priority order: - - 1. Context.user_id - For authenticated sessions (set via SessionManager.session()) - 2. user_manager.get_id() - Legacy/custom context compatibility - 3. X-User-ID HTTP header - Illustrative/testing only; assumes a trusted upstream proxy - has authenticated the request and injected the header - 4. "default_user" - Fallback for development/testing without authentication - - Returns: - str: The user ID for memory operations - """ - # Priority 1: Get user_id from the runtime context. + """Return the identity resolved by the runtime session.""" try: user_id = self._context.user_id if user_id: - logger.debug(f"Using user_id from context: {user_id}") - return user_id - except Exception as e: - logger.debug(f"Failed to get user_id from context: {e}") - - # Priority 2: Get user_id from user_manager (legacy/custom context compatibility) - user_manager = getattr(self._context, "user_manager", None) - if user_manager and hasattr(user_manager, 'get_id'): - try: - user_id = user_manager.get_id() - if user_id: - logger.debug(f"Using user_id from user_manager: {user_id}") - return user_id - except Exception as e: - logger.debug(f"Failed to get user_id from user_manager: {e}") - - # Priority 3: Extract an identity header injected by a trusted upstream proxy. This is illustrative/testing - # support only; an application must not accept a client-supplied X-User-ID header as authentication. - metadata = getattr(self._context, "metadata", None) - headers = getattr(metadata, "headers", None) if metadata else None - if headers: - user_id = headers.get("x-user-id") or headers.get("X-User-ID") - if user_id: - logger.debug(f"Using user_id from X-User-ID header: {user_id}") return user_id + except Exception: + logger.debug("Failed to get user_id from context", exc_info=True) + raise RuntimeError("No resolved user identity is available for automatic memory operations") - # Fallback: default for development/testing - logger.debug("Using default user_id: default_user") - return "default_user" + def _get_user_id(self, state: AutoMemoryWrapperState) -> str: + """Resolve the identity once and retain it for the full graph invocation.""" + if state.user_id is None: + state.user_id = self._get_user_id_from_context() + return state.user_id def get_wrapper_node_count(self) -> int: """ @@ -164,7 +131,7 @@ async def capture_user_message_node(self, state: AutoMemoryWrapperState) -> Auto user_message = state.messages[-1] if isinstance(user_message, HumanMessage): # Get user_id from runtime context - user_id = self._get_user_id_from_context() + user_id = self._get_user_id(state) # Add to memory await self.memory_editor.add_items( @@ -183,7 +150,7 @@ async def memory_retrieve_node(self, state: AutoMemoryWrapperState) -> AutoMemor user_message = state.messages[-1] # Get user_id from runtime context - user_id = self._get_user_id_from_context() + user_id = self._get_user_id(state) # Retrieve memory from memory provider memory_items = await self.memory_editor.search( @@ -249,7 +216,7 @@ async def capture_ai_response_node(self, state: AutoMemoryWrapperState) -> AutoM ai_message = state.messages[-1] if isinstance(ai_message, AIMessage): # Get user_id from runtime context - user_id = self._get_user_id_from_context() + user_id = self._get_user_id(state) # Add to memory await self.memory_editor.add_items( diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/register.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/register.py index 0201c98396..d6c23b86dd 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/register.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/register.py @@ -69,9 +69,9 @@ class AutoMemoryAgentConfig(AgentBaseConfig, name="auto_memory_agent"): **Multi-tenant User Isolation:** - User ID is automatically extracted from runtime context (user_manager.get_id()) for proper - multi-tenant memory isolation. Set user_manager via SessionManager.session() in production. - Defaults to "default_user" for testing/development. See README.md for deployment examples. + User ID is read only from the identity resolved by the runtime session. Memory operations fail + closed when no identity is available. Configure authentication on the front end or pass an + explicit user ID to ``SessionManager.session()``. See README.md for deployment examples. """ # Memory configuration diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/state.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/state.py index 29427dee99..cc7a21c4a7 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/state.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/state.py @@ -28,3 +28,4 @@ class AutoMemoryWrapperState(BaseModel): """ messages: list[BaseMessage] = Field(default_factory=list, description="Conversation messages with context injection") + user_id: str | None = Field(default=None, description="Resolved runtime identity used for memory operations") diff --git a/packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py b/packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py index 0c438948a7..0fb1e3213d 100644 --- a/packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py +++ b/packages/nvidia_nat_langchain/tests/agent/test_auto_memory_wrapper.py @@ -69,7 +69,7 @@ async def _ainvoke(chat_request: ChatRequest): def fixture_mock_context() -> Mock: """Create a mock Context for testing.""" context = Mock(spec=Context) - context.user_id = None + context.user_id = "test-user" context.metadata = None return context @@ -140,10 +140,11 @@ def test_get_wrapper_node_count_minimal(self, mock_inner_agent, mock_memory_edit count = wrapper.get_wrapper_node_count() assert count == 1 # only inner_agent - def test_get_user_id_default(self, wrapper_graph): - """Test user ID extraction defaults to 'default_user'.""" - user_id = wrapper_graph._get_user_id_from_context() - assert user_id == "default_user" + def test_get_user_id_missing_raises(self, wrapper_graph, mock_context): + """Test memory access fails closed when the runtime did not resolve an identity.""" + mock_context.user_id = None + with pytest.raises(RuntimeError, match="No resolved user identity"): + wrapper_graph._get_user_id_from_context() def test_get_user_id_from_context(self, wrapper_graph, mock_context): """Test user ID extraction from Context.user_id.""" @@ -153,24 +154,22 @@ def test_get_user_id_from_context(self, wrapper_graph, mock_context): user_id = wrapper_graph._get_user_id_from_context() assert user_id == "user-from-context" - def test_get_user_id_from_header(self, wrapper_graph, mock_context): - """Test user ID extraction from X-User-ID header.""" + def test_get_user_id_does_not_read_raw_header(self, wrapper_graph, mock_context): + """Test the wrapper does not trust request metadata directly.""" + mock_context.user_id = None mock_context.metadata = Mock() mock_context.metadata.headers = {"x-user-id": "test-user-123"} - with patch('nat.plugins.langchain.agent.auto_memory_wrapper.agent.Context.get', return_value=mock_context): - wrapper_graph._context = mock_context - user_id = wrapper_graph._get_user_id_from_context() - assert user_id == "test-user-123" + with pytest.raises(RuntimeError, match="No resolved user identity"): + wrapper_graph._get_user_id_from_context() - def test_get_user_id_from_user_manager(self, wrapper_graph, mock_context): - """Test user ID extraction from user_manager.""" + def test_get_user_id_does_not_use_legacy_user_manager(self, wrapper_graph, mock_context): + """Test the wrapper consumes only the authoritative runtime identity.""" + mock_context.user_id = None mock_user_manager = Mock() mock_user_manager.get_id.return_value = "user-from-manager" mock_context.user_manager = mock_user_manager - with patch('nat.plugins.langchain.agent.auto_memory_wrapper.agent.Context.get', return_value=mock_context): - wrapper_graph._context = mock_context - user_id = wrapper_graph._get_user_id_from_context() - assert user_id == "user-from-manager" + with pytest.raises(RuntimeError, match="No resolved user identity"): + wrapper_graph._get_user_id_from_context() def test_langchain_message_to_nat_message_human(self): """Test conversion of HumanMessage to NAT Message.""" @@ -204,7 +203,19 @@ async def test_capture_user_message_node(self, wrapper_graph, mock_memory_editor items = call_args[0][0] assert len(items) == 1 assert items[0].conversation == [{"role": "user", "content": "Test message"}] - assert items[0].user_id == "default_user" + assert items[0].user_id == "test-user" + + async def test_user_id_is_resolved_once_per_graph_state(self, wrapper_graph, mock_context, mock_memory_editor): + state = AutoMemoryWrapperState(messages=[HumanMessage(content="Remember this")]) + + await wrapper_graph.memory_retrieve_node(state) + mock_context.user_id = "different-user" + await wrapper_graph.capture_user_message_node(state) + + assert state.user_id == "test-user" + assert mock_memory_editor.search.call_args.kwargs["user_id"] == "test-user" + saved_item = mock_memory_editor.add_items.call_args.args[0][0] + assert saved_item.user_id == "test-user" async def test_capture_user_message_node_disabled(self, mock_inner_agent, mock_memory_editor, mock_context): """Test capture_user_message_node when disabled.""" diff --git a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py index a979708192..fe880dec61 100644 --- a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py +++ b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py @@ -25,6 +25,7 @@ from nat.builder.function import FunctionGroup from nat.builder.workflow_builder import WorkflowBuilder from nat.runtime.session import SessionManager +from nat.runtime.user_manager import IdentityHeaderError logger = logging.getLogger(__name__) @@ -197,7 +198,6 @@ async def get_mcp_client_tool_list() -> MCPClientToolListResponse: async def get_per_user_mcp_client_tool_list( request: Request, - user_id: str | None = None, ) -> MCPClientToolListResponse: """Get the list of MCP tools for a specific user in per-user workflows. @@ -209,9 +209,11 @@ async def get_per_user_mcp_client_tool_list( raise HTTPException(status_code=400, detail="No per-user workflow is configured.") try: - async with per_user_manager.session(user_id=user_id, http_connection=request) as session: + async with per_user_manager.session(http_connection=request) as session: mcp_clients_info = await _collect_mcp_client_tool_list(session.workflow.function_groups) return MCPClientToolListResponse(mcp_clients=mcp_clients_info) + except IdentityHeaderError: + raise except Exception as e: logger.exception("Error in per-user MCP client tool list endpoint: %s", e) raise HTTPException(status_code=500, diff --git a/packages/nvidia_nat_mcp/tests/server/test_mcp_client_endpoint.py b/packages/nvidia_nat_mcp/tests/server/test_mcp_client_endpoint.py index cbd119ab24..c9ca3b2bfe 100644 --- a/packages/nvidia_nat_mcp/tests/server/test_mcp_client_endpoint.py +++ b/packages/nvidia_nat_mcp/tests/server/test_mcp_client_endpoint.py @@ -231,7 +231,8 @@ async def test_mcp_client_tool_list_per_user_success(app_worker): assert group["total_tools"] == 1 assert group["available_tools"] == 1 assert group["tools"][0]["name"] == "alias_tool" - assert per_user_manager._user_ids == ["alice"] + # The caller-controlled query parameter is ignored; identity comes from the request session. + assert per_user_manager._user_ids == [None] async def test_mcp_client_tool_list_per_user_missing_config(app_worker): From 4ec210603496354e5166dcd6c359fc663fe70b4b Mon Sep 17 00:00:00 2001 From: David Gardner Date: Mon, 31 Aug 2026 16:49:11 -0700 Subject: [PATCH 2/6] Formatting Signed-off-by: David Gardner --- packages/nvidia_nat_core/src/nat/runtime/session.py | 4 ++-- .../nat/authentication/test_http_basic_auth_exchanger.py | 4 +++- .../tests/nat/front_ends/fastapi/test_message_handler.py | 4 +--- .../nvidia_nat_core/tests/nat/runtime/test_user_manager.py | 4 +++- .../src/nat/plugins/mcp/client/fastapi_routes.py | 4 +--- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/packages/nvidia_nat_core/src/nat/runtime/session.py b/packages/nvidia_nat_core/src/nat/runtime/session.py index fc8aa2c380..669ebb7371 100644 --- a/packages/nvidia_nat_core/src/nat/runtime/session.py +++ b/packages/nvidia_nat_core/src/nat/runtime/session.py @@ -500,8 +500,8 @@ async def session(self, if isinstance(http_connection, Request): if identity_header is not None or user_id is None: - user_info = UserManager.extract_user_from_connection( - http_connection, identity_header=identity_header) + user_info = UserManager.extract_user_from_connection(http_connection, + identity_header=identity_header) if user_info is not None: user_id = user_info.get_user_id() token_workflow_parent_id, token_workflow_parent_name = \ diff --git a/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py b/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py index 3c97b1fd30..5885a5f695 100644 --- a/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py +++ b/packages/nvidia_nat_core/tests/nat/authentication/test_http_basic_auth_exchanger.py @@ -101,7 +101,9 @@ async def test_uses_resolved_context_identity(monkeypatch): async def cb(cfg, flow): return AuthenticatedContext( headers={"Authorization": "Basic YQ=="}, - metadata={"username": "a", "password": "b"}, + metadata={ + "username": "a", "password": "b" + }, ) _patch_context(monkeypatch, cb, user_id="resolved-user") diff --git a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py index fe97e79ff4..d88146701c 100644 --- a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py +++ b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_message_handler.py @@ -114,9 +114,7 @@ async def test_auth_message_cannot_replace_trusted_header_identity(): payload=ApiKeyAuthPayload(method="api_key", token="replacement-key"), ) - with patch( - "nat.front_ends.fastapi.message_handler.UserManager.from_auth_payload_with_verification", - ) as resolver: + with patch("nat.front_ends.fastapi.message_handler.UserManager.from_auth_payload_with_verification", ) as resolver: await handler._process_auth_message(msg) resolver.assert_not_called() diff --git a/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py b/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py index 226bbd9487..5ded64b8c4 100644 --- a/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py +++ b/packages/nvidia_nat_core/tests/nat/runtime/test_user_manager.py @@ -117,7 +117,9 @@ def test_duplicate_identity_header_is_rejected(self): def test_identity_header_takes_precedence_over_other_credentials(self): request = _mock_request( cookies={SESSION_COOKIE_NAME: "cookie-user"}, - headers={"X-User-ID": "header-user", "authorization": "Bearer api-key"}, + headers={ + "X-User-ID": "header-user", "authorization": "Bearer api-key" + }, ) header_info = UserManager.extract_user_from_connection(request, identity_header="X-User-ID") diff --git a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py index fe880dec61..e4ce1168e2 100644 --- a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py +++ b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py @@ -196,9 +196,7 @@ async def get_mcp_client_tool_list() -> MCPClientToolListResponse: logger.error(f"Error in MCP client tool list endpoint: {e}") raise HTTPException(status_code=500, detail=f"Failed to retrieve MCP client information: {str(e)}") from e - async def get_per_user_mcp_client_tool_list( - request: Request, - ) -> MCPClientToolListResponse: + async def get_per_user_mcp_client_tool_list(request: Request, ) -> MCPClientToolListResponse: """Get the list of MCP tools for a specific user in per-user workflows. Uses the per-user workflow builder to resolve function groups and From 22881b6bc702c92422910d896773cf139d0edf07 Mon Sep 17 00:00:00 2001 From: David Gardner Date: Mon, 31 Aug 2026 16:56:49 -0700 Subject: [PATCH 3/6] Update phrasing Signed-off-by: David Gardner --- .../agents/auto-memory-wrapper/auto-memory-wrapper.md | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md index 0a255e587a..936b84619f 100644 --- a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md +++ b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md @@ -210,15 +210,11 @@ general: identity_header: X-User-ID ``` -Use this setting only for local testing or behind an authenticating reverse proxy. It is secure only when all of the -following are true: +Use this setting only for local testing or behind an authenticating reverse proxy. It is secure only when all of the following are true: - Clients cannot connect to `nat serve` directly; only the trusted proxy can reach its listening port. -- The proxy authenticates every request and overwrites `X-User-ID` with the authenticated principal. It must not - preserve, append, or forward a client-supplied value. -- The backend Docker network and every container attached to it are trusted. Do not publish the `nat serve` port to the - host; publish only the proxy port. -- External traffic is protected as appropriate, typically with TLS at the proxy. +- The proxy authenticates every request and overwrites `X-User-ID` with the authenticated principal. It must not preserve, append, or forward a client-supplied value. +- `nat serve` is network isolated in such a way that all incoming network traffic is from the proxy. When configured, the header is authoritative for HTTP and WebSocket requests. A missing, empty, or repeated header is rejected. Other connection credentials and WebSocket auth messages cannot replace it. The setting cannot be combined From ed1716625e393cd87678b124cda5aacdd88dff5c Mon Sep 17 00:00:00 2001 From: David Gardner Date: Mon, 31 Aug 2026 16:58:50 -0700 Subject: [PATCH 4/6] Update phrasing Signed-off-by: David Gardner --- .../agents/auto-memory-wrapper/auto-memory-wrapper.md | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md index 936b84619f..88b8cb6025 100644 --- a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md +++ b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md @@ -216,11 +216,10 @@ Use this setting only for local testing or behind an authenticating reverse prox - The proxy authenticates every request and overwrites `X-User-ID` with the authenticated principal. It must not preserve, append, or forward a client-supplied value. - `nat serve` is network isolated in such a way that all incoming network traffic is from the proxy. -When configured, the header is authoritative for HTTP and WebSocket requests. A missing, empty, or repeated header is -rejected. Other connection credentials and WebSocket auth messages cannot replace it. The setting cannot be combined +When configured, the header is authoritative for HTTP and WebSocket requests. A missing, empty, or repeated header is rejected. Other connection credentials and WebSocket auth messages cannot replace it. The setting cannot be combined with `accepted_identity_credentials` or `identity_authentication`. -After configuring the trusted boundary, a request forwarded by the proxy can contain: +For local testing, this can be simulated with: ```bash curl -X POST http://localhost:8000/chat \ @@ -355,8 +354,7 @@ workflow: 2. **Memory backends are interchangeable** - Works with any implementation of `MemoryEditor` interface 3. **No memory tools needed** - The wrapped agent does not need explicit memory tools configured 4. **Transparent to inner agent** - The wrapped agent is unaware of memory operations -5. **Trusted identity headers are opt-in** - Never enable `identity_header` on a server that untrusted clients or - untrusted containers can reach directly. +5. **Trusted identity headers are opt-in** - Never enable `identity_header` on a server that untrusted clients can reach directly. --- From 51ab101cdfff1f464b300b7dbb6fbfeb571a7dce Mon Sep 17 00:00:00 2001 From: David Gardner Date: Mon, 31 Aug 2026 17:03:03 -0700 Subject: [PATCH 5/6] Update phrasing Signed-off-by: David Gardner --- docs/source/reference/rest-api/websockets.md | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/docs/source/reference/rest-api/websockets.md b/docs/source/reference/rest-api/websockets.md index fb1ffeaa70..6d2fa214db 100644 --- a/docs/source/reference/rest-api/websockets.md +++ b/docs/source/reference/rest-api/websockets.md @@ -112,8 +112,7 @@ general: Each JWT provider requires `issuer_url`, `jwks_uri`, and `audience`. It can also require `scopes` and configure `timeout` and `leeway`. Add another named provider and select it in `identity_authentication` to accept JWT tokens from another issuer. The issuer claim selects the matching provider; an unknown issuer or a failed signature, time, audience, or scope check returns an authentication error and does not restore workflow state. If `identity_authentication` is omitted, JWT tokens retain the existing decode-only behavior. Configuring `identity_authentication` while excluding `jwt` from `accepted_identity_credentials` is invalid. -For local testing or a deployment where an authenticating reverse proxy is the only service that can reach -`nat serve`, the FastAPI front end can instead trust an upstream identity header: +For local testing or a deployment where an authenticating reverse proxy is the only service that can reach `nat serve`, the FastAPI front end can instead trust an upstream identity header: ```yaml general: @@ -122,11 +121,9 @@ general: identity_header: X-User-ID ``` -This setting is secure only if the proxy authenticates every connection, overwrites any client-supplied value, and the -toolkit port is not exposed outside a trusted backend network. Every container that can reach that network must also -be trusted. When enabled, the header is authoritative: the connection is rejected if it is missing, empty, or repeated, +This setting is secure only if the proxy authenticates every connection, overwrites any client-supplied value, and the server is in an isolated network environment. When enabled, the header is authoritative: the connection is rejected if it is missing, empty, or repeated, and an `auth_message` cannot replace the resolved identity. `identity_header` cannot be combined with -`accepted_identity_credentials` or `identity_authentication`. +`accepted_identity_credentials` or `identity_authentication`. Do not use this technique if `nat serve` is reachable by untrusted clients. ## Auth Message This message allows clients to authenticate over a WebSocket connection when header-based or From 3e4c401a3924400b40f1d8fc0001237b924c55cc Mon Sep 17 00:00:00 2001 From: David Gardner Date: Tue, 1 Sep 2026 08:40:50 -0700 Subject: [PATCH 6/6] Document the 401 error Signed-off-by: David Gardner --- .../src/nat/plugins/mcp/client/fastapi_routes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py index e4ce1168e2..5b70f484c6 100644 --- a/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py +++ b/packages/nvidia_nat_mcp/src/nat/plugins/mcp/client/fastapi_routes.py @@ -266,6 +266,9 @@ async def get_per_user_mcp_client_tool_list(request: Request, ) -> MCPClientTool 400: { "description": "No per-user workflow is configured" }, + 401: { + "description": "Required identity header is missing or invalid" + }, 500: { "description": "Internal Server Error" }