diff --git a/docs/source/reference/rest-api/websockets.md b/docs/source/reference/rest-api/websockets.md index 548e2323cb..7c0b8516a5 100644 --- a/docs/source/reference/rest-api/websockets.md +++ b/docs/source/reference/rest-api/websockets.md @@ -65,7 +65,7 @@ to the client. ## Reconnecting an Active Conversation -To resume an active workflow, reconnect with the workflow's `conversation_id` query parameter and an identity credential that resolves to the same user who started it. The server restores state only when both values match. A connection with another identity, or no identity, does not receive the existing workflow state or its pending Human-in-the-Loop prompt. +To resume an active workflow, reconnect using the `conversation_id` query parameter for that workflow and an identity credential that resolves to the same user who started the workflow. The server restores the workflow state only when both values match. A connection with another identity, or no identity, does not receive the existing workflow state or the pending Human-in-the-Loop prompt for that workflow. The server accepts the following identity credentials: @@ -76,6 +76,42 @@ The server accepts the following identity credentials: Clients can also send a JWT, API key, or Basic credentials through an `auth_message`. Send the message before expecting restoration; restoration occurs after authentication succeeds. +By default, all listed identity credential methods are accepted. To restrict WebSocket identity credentials, set `accepted_identity_credentials` in the FastAPI front-end configuration: + +```yaml +general: + front_end: + _type: fastapi + accepted_identity_credentials: + - session_cookie + - jwt +``` + +Supported values are `session_cookie`, `jwt`, `api_key`, and `basic`. The `api_key` value covers both Bearer API keys and the `X-API-Key` header. An empty list rejects every supplied identity credential. When a client supplies a disabled credential method, the server returns an authentication error and does not restore workflow state. + +JWT signature and claim verification is optional. Define one or more named JWT authentication providers, then select them with `identity_authentication` in the FastAPI front-end configuration. Each selected provider verifies JWT tokens from its configured issuer before the server resolves the user identity: + +```yaml +authentication: + corporate_jwt: + _type: jwt + issuer_url: https://identity.example.com + jwks_uri: https://identity.example.com/.well-known/jwks.json + audience: nemo-agent-toolkit + scopes: + - workflow:resume + +general: + front_end: + _type: fastapi + accepted_identity_credentials: + - jwt + identity_authentication: + - corporate_jwt +``` + +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. + ## 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/packages/nvidia_nat_core/src/nat/authentication/jwt/__init__.py b/packages/nvidia_nat_core/src/nat/authentication/jwt/__init__.py new file mode 100644 index 0000000000..bcd923c929 --- /dev/null +++ b/packages/nvidia_nat_core/src/nat/authentication/jwt/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/packages/nvidia_nat_core/src/nat/authentication/jwt/jwt_auth_provider.py b/packages/nvidia_nat_core/src/nat/authentication/jwt/jwt_auth_provider.py new file mode 100644 index 0000000000..981f6aaec3 --- /dev/null +++ b/packages/nvidia_nat_core/src/nat/authentication/jwt/jwt_auth_provider.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Named JWT identity verification provider.""" + +from nat.authentication.credential_validator.bearer_token_validator import BearerTokenValidator +from nat.authentication.interfaces import AuthProviderBase +from nat.authentication.jwt.jwt_auth_provider_config import JwtAuthProviderConfig +from nat.data_models.authentication import AuthResult +from nat.data_models.authentication import TokenValidationResult + + +class JwtAuthProvider(AuthProviderBase[JwtAuthProviderConfig]): + """Validate inbound JWTs using a named issuer policy.""" + + def __init__(self, config: JwtAuthProviderConfig) -> None: + super().__init__(config) + self._validator = BearerTokenValidator( + issuer=config.issuer_url, + audience=config.audience, + jwks_uri=config.jwks_uri, + scopes=config.scopes, + timeout=config.timeout, + leeway=config.leeway, + ) + + async def verify(self, token: str) -> TokenValidationResult: + """Verify a JWT against this provider's configured trust policy.""" + return await self._validator.verify(token) + + @property + def validator(self) -> BearerTokenValidator: + """Return the cached validator used by this named provider.""" + return self._validator + + async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult: + """Validate the supplied ``token`` through the authentication-provider interface.""" + token = kwargs.get("token") + if not isinstance(token, str) or not token: + raise ValueError("JWT authentication requires a token") + + result = await self.verify(token) + if not result.active: + raise ValueError("JWT verification failed") + return AuthResult(raw=result.model_dump()) diff --git a/packages/nvidia_nat_core/src/nat/authentication/jwt/jwt_auth_provider_config.py b/packages/nvidia_nat_core/src/nat/authentication/jwt/jwt_auth_provider_config.py new file mode 100644 index 0000000000..d07fac2eb5 --- /dev/null +++ b/packages/nvidia_nat_core/src/nat/authentication/jwt/jwt_auth_provider_config.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Configuration for validating inbound JWT identity credentials.""" + +from urllib.parse import urlparse + +from pydantic import Field +from pydantic import field_validator + +from nat.data_models.authentication import AuthProviderBaseConfig + + +class JwtAuthProviderConfig(AuthProviderBaseConfig, name="jwt"): + """Named JWT verification policy for inbound identity credentials.""" + + issuer_url: str = Field(description="Expected JWT issuer claim.") + jwks_uri: str = Field(description="Endpoint containing trusted public keys for signature verification.") + audience: str = Field(description="Expected JWT audience claim.") + scopes: list[str] = Field(default_factory=list, description="Scopes required in a verified JWT.") + timeout: float = Field(default=10.0, gt=0, description="HTTP timeout for JWKS requests.") + leeway: int = Field(default=60, ge=0, description="Clock-skew allowance for JWT time claims, in seconds.") + + @field_validator("issuer_url", "jwks_uri") + @classmethod + def require_secure_url(cls, value: str, info) -> str: + parsed = urlparse(value) + is_local_http = parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1", "::1"} + if not parsed.netloc or (parsed.scheme != "https" and not is_local_http): + raise ValueError(f"{info.field_name} must use HTTPS (HTTP is allowed only for localhost)") + return value diff --git a/packages/nvidia_nat_core/src/nat/authentication/jwt/register.py b/packages/nvidia_nat_core/src/nat/authentication/jwt/register.py new file mode 100644 index 0000000000..8332325db3 --- /dev/null +++ b/packages/nvidia_nat_core/src/nat/authentication/jwt/register.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nat.authentication.jwt.jwt_auth_provider_config import JwtAuthProviderConfig +from nat.builder.builder import Builder +from nat.cli.register_workflow import register_auth_provider + + +@register_auth_provider(config_type=JwtAuthProviderConfig) +async def jwt_auth_provider(config: JwtAuthProviderConfig, builder: Builder): + from nat.authentication.jwt.jwt_auth_provider import JwtAuthProvider + + yield JwtAuthProvider(config) diff --git a/packages/nvidia_nat_core/src/nat/authentication/register.py b/packages/nvidia_nat_core/src/nat/authentication/register.py index 7591279740..d0d2400efd 100644 --- a/packages/nvidia_nat_core/src/nat/authentication/register.py +++ b/packages/nvidia_nat_core/src/nat/authentication/register.py @@ -17,4 +17,5 @@ from nat.authentication.api_key import register as register_api_key from nat.authentication.http_basic_auth import register as register_http_basic_auth +from nat.authentication.jwt import register as register_jwt from nat.authentication.oauth2 import register as register_oauth2 diff --git a/packages/nvidia_nat_core/src/nat/data_models/api_server.py b/packages/nvidia_nat_core/src/nat/data_models/api_server.py index 37bdf39ff2..cfdbdcd366 100644 --- a/packages/nvidia_nat_core/src/nat/data_models/api_server.py +++ b/packages/nvidia_nat_core/src/nat/data_models/api_server.py @@ -35,7 +35,7 @@ from nat.data_models.interactive import HumanPrompt from nat.utils.type_converter import GlobalTypeConverter -FINISH_REASONS = frozenset({'stop', 'length', 'tool_calls', 'content_filter', 'function_call'}) +FINISH_REASONS = frozenset({"stop", "length", "tool_calls", "content_filter", "function_call"}) UNKNOWN_MODEL_SENTINEL = "unknown-model" @@ -44,6 +44,7 @@ class UserMessageContentRoleType(StrEnum): """ Enum representing chat message roles in API requests and responses. """ + USER = "user" ASSISTANT = "assistant" SYSTEM = "system" @@ -54,6 +55,7 @@ class Request(BaseModel): """ Request is a data model that represents HTTP request and WebSocket attributes. """ + model_config = ConfigDict(extra="forbid") method: str | None = Field(default=None, @@ -77,6 +79,7 @@ class ChatContentType(StrEnum): """ ChatContentType is an Enum that represents the type of Chat content. """ + TEXT = "text" IMAGE_URL = "image_url" INPUT_AUDIO = "input_audio" @@ -168,45 +171,55 @@ class ChatRequest(BaseModel): tool_choice: str | dict[str, typing.Any] | None = Field(default=None, description="Controls which tool is called") parallel_tool_calls: bool | None = Field(default=True, description="Whether to enable parallel function calling") user: str | None = Field(default=None, description="Unique identifier representing end-user") - model_config = ConfigDict(extra="allow", - json_schema_extra={ - "example": { - "model": "nvidia/nemotron", - "messages": [{ - "role": "user", "content": "who are you?" - }], - "temperature": 0.7, - "stream": False - } - }) + model_config = ConfigDict( + extra="allow", + json_schema_extra={ + "example": { + "model": "nvidia/nemotron", + "messages": [{ + "role": "user", "content": "who are you?" + }], + "temperature": 0.7, + "stream": False, + } + }, + ) @staticmethod - def from_string(data: str, - *, - model: str | None = None, - temperature: float | None = None, - max_tokens: int | None = None, - top_p: float | None = None) -> "ChatRequest": - - return ChatRequest(messages=[Message(content=data, role=UserMessageContentRoleType.USER)], - model=model, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p) + def from_string( + data: str, + *, + model: str | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + top_p: float | None = None, + ) -> "ChatRequest": + + return ChatRequest( + messages=[Message(content=data, role=UserMessageContentRoleType.USER)], + model=model, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) @staticmethod - def from_content(content: list[UserContent], - *, - model: str | None = None, - temperature: float | None = None, - max_tokens: int | None = None, - top_p: float | None = None) -> "ChatRequest": - - return ChatRequest(messages=[Message(content=content, role=UserMessageContentRoleType.USER)], - model=model, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p) + def from_content( + content: list[UserContent], + *, + model: str | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + top_p: float | None = None, + ) -> "ChatRequest": + + return ChatRequest( + messages=[Message(content=content, role=UserMessageContentRoleType.USER)], + model=model, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + ) class ChatRequestOrMessage(BaseModel): @@ -220,6 +233,7 @@ class ChatRequestOrMessage(BaseModel): Note: When `messages` is provided, extra fields are allowed to enable lossless round-trip conversion with ChatRequest. When `input_message` is provided, no extra fields are permitted. """ + model_config = ConfigDict( extra="allow", json_schema_extra={ @@ -232,7 +246,7 @@ class ChatRequestOrMessage(BaseModel): "role": "user", "content": "What can you do?" }], "model": "nvidia/nemotron", - "temperature": 0.7 + "temperature": 0.7, }, ], "oneOf": [ @@ -244,7 +258,7 @@ class ChatRequestOrMessage(BaseModel): }, }, "additionalProperties": { - "not": True, "errorMessage": 'remove additional property ${0#}' + "not": True, "errorMessage": "remove additional property ${0#}" }, }, { @@ -254,9 +268,9 @@ class ChatRequestOrMessage(BaseModel): "type": "array" }, }, - "additionalProperties": True + "additionalProperties": True, }, - ] + ], }, ) @@ -265,7 +279,8 @@ class ChatRequestOrMessage(BaseModel): input_message: str | None = Field( default=None, - description="A single input message to process. Useful for functions that do not require a conversation") + description="A single input message to process. Useful for functions that do not require a conversation", + ) @property def is_string(self) -> bool: @@ -295,12 +310,14 @@ class ChoiceMessage(BaseModel): class ChoiceDeltaToolCallFunction(BaseModel): """Function details within a streamed tool call delta (OpenAI-compatible).""" + name: str | None = None arguments: str | None = None class ChoiceDeltaToolCall(BaseModel): """Tool call delta for streaming responses (OpenAI-compatible).""" + index: int id: str | None = None type: str | None = None @@ -309,6 +326,7 @@ class ChoiceDeltaToolCall(BaseModel): class ChoiceDelta(BaseModel): """Delta object for streaming responses (OpenAI-compatible)""" + content: str | None = None role: UserMessageContentRoleType | None = None tool_calls: list[ChoiceDeltaToolCall] | None = None @@ -316,18 +334,21 @@ class ChoiceDelta(BaseModel): class ChoiceBase(BaseModel): """Base choice model with common fields for both streaming and non-streaming responses""" + model_config = ConfigDict(extra="allow") - finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None + finish_reason: typing.Literal["stop", "length", "tool_calls", "content_filter", "function_call"] | None = None index: int class ChatResponseChoice(ChoiceBase): """Choice model for non-streaming responses - contains message field""" + message: ChoiceMessage class ChatResponseChunkChoice(ChoiceBase): """Choice model for streaming responses - contains delta field""" + delta: ChoiceDelta @@ -381,19 +402,21 @@ class ChatResponse(ResponseBaseModelOutput): system_fingerprint: str | None = None service_tier: typing.Literal["scale", "default"] | None = None - @field_serializer('created') + @field_serializer("created") def serialize_created(self, created: datetime.datetime) -> int: """Serialize datetime to Unix timestamp for OpenAI compatibility""" return int(created.timestamp()) @staticmethod - def from_string(data: str, - *, - id_: str | None = None, - object_: str | None = None, - model: str | None = None, - created: datetime.datetime | None = None, - usage: Usage) -> "ChatResponse": + def from_string( + data: str, + *, + id_: str | None = None, + object_: str | None = None, + model: str | None = None, + created: datetime.datetime | None = None, + usage: Usage, + ) -> "ChatResponse": if id_ is None: id_ = str(uuid.uuid4()) @@ -404,17 +427,20 @@ def from_string(data: str, if created is None: created = datetime.datetime.now(datetime.UTC) - return ChatResponse(id=id_, - object=object_, - model=model, - created=created, - choices=[ - ChatResponseChoice(index=0, - message=ChoiceMessage(content=data, - role=UserMessageContentRoleType.ASSISTANT), - finish_reason="stop") - ], - usage=usage) + return ChatResponse( + id=id_, + object=object_, + model=model, + created=created, + choices=[ + ChatResponseChoice( + index=0, + message=ChoiceMessage(content=data, role=UserMessageContentRoleType.ASSISTANT), + finish_reason="stop", + ) + ], + usage=usage, + ) class ChatResponseChunk(ResponseBaseModelOutput): @@ -435,19 +461,21 @@ class ChatResponseChunk(ResponseBaseModelOutput): service_tier: typing.Literal["scale", "default"] | None = None usage: Usage | None = None - @field_serializer('created') + @field_serializer("created") def serialize_created(self, created: datetime.datetime) -> int: """Serialize datetime to Unix timestamp for OpenAI compatibility""" return int(created.timestamp()) @staticmethod - def from_string(data: str, - *, - id_: str | None = None, - created: datetime.datetime | None = None, - model: str | None = None, - object_: str | None = None, - finish_reason: str | None = None) -> "ChatResponseChunk": + def from_string( + data: str, + *, + id_: str | None = None, + created: datetime.datetime | None = None, + model: str | None = None, + object_: str | None = None, + finish_reason: str | None = None, + ) -> "ChatResponseChunk": if id_ is None: id_ = str(uuid.uuid4()) @@ -460,28 +488,32 @@ def from_string(data: str, final_finish_reason = finish_reason if finish_reason in FINISH_REASONS else None - return ChatResponseChunk(id=id_, - choices=[ - ChatResponseChunkChoice(index=0, - delta=ChoiceDelta( - content=data, - role=UserMessageContentRoleType.ASSISTANT), - finish_reason=final_finish_reason) - ], - created=created, - model=model, - object=object_) + return ChatResponseChunk( + id=id_, + choices=[ + ChatResponseChunkChoice( + index=0, + delta=ChoiceDelta(content=data, role=UserMessageContentRoleType.ASSISTANT), + finish_reason=final_finish_reason, + ) + ], + created=created, + model=model, + object=object_, + ) @staticmethod - def create_streaming_chunk(content: str, - *, - id_: str | None = None, - created: datetime.datetime | None = None, - model: str | None = None, - role: UserMessageContentRoleType | None = None, - finish_reason: str | None = None, - usage: Usage | None = None, - system_fingerprint: str | None = None) -> "ChatResponseChunk": + def create_streaming_chunk( + content: str, + *, + id_: str | None = None, + created: datetime.datetime | None = None, + model: str | None = None, + role: UserMessageContentRoleType | None = None, + finish_reason: str | None = None, + usage: Usage | None = None, + system_fingerprint: str | None = None, + ) -> "ChatResponseChunk": """Create an OpenAI-compatible streaming chunk""" if id_ is None: id_ = str(uuid.uuid4()) @@ -501,14 +533,17 @@ def create_streaming_chunk(content: str, index=0, delta=delta, finish_reason=typing.cast( - typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None, - final_finish_reason)) + typing.Literal["stop", "length", "tool_calls", "content_filter", "function_call"] | None, + final_finish_reason, + ), + ) ], created=created, model=model, object="chat.completion.chunk", usage=usage, - system_fingerprint=system_fingerprint) + system_fingerprint=system_fingerprint, + ) class ResponseIntermediateStep(ResponseBaseModelIntermediate): @@ -631,6 +666,7 @@ class WebSocketMessageType(StrEnum): """ WebSocketMessageType is an Enum that represents WebSocket Message types. """ + USER_MESSAGE = "user_message" RESPONSE_MESSAGE = "system_response_message" INTERMEDIATE_STEP_MESSAGE = "system_intermediate_message" @@ -646,6 +682,7 @@ class WorkflowSchemaType(StrEnum): """ WorkflowSchemaType is an Enum that represents Workkflow response types. """ + GENERATE_STREAM = "generate_stream" CHAT_STREAM = "chat_stream" GENERATE = "generate" @@ -656,6 +693,7 @@ class WebSocketMessageStatus(StrEnum): """ WebSocketMessageStatus is an Enum that represents the status of a WebSocket message. """ + IN_PROGRESS = "in_progress" COMPLETE = "complete" @@ -702,6 +740,7 @@ class WebSocketUserMessage(BaseModel): For more details, refer to the API documentation: docs/source/developer_guide/websockets.md """ + # Allow extra fields in the model_config to support derived models model_config = ConfigDict(extra="allow") @@ -721,6 +760,7 @@ class WebSocketUserInteractionResponseMessage(BaseModel): For more details, refer to the API documentation: docs/source/developer_guide/websockets.md """ + type: typing.Literal[WebSocketMessageType.USER_INTERACTION_MESSAGE] id: str = "default" thread_id: str = "default" @@ -735,36 +775,51 @@ class WebSocketUserInteractionResponseMessage(BaseModel): class AuthMethod(StrEnum): """Supported authentication methods for WebSocket auth messages.""" + JWT = "jwt" API_KEY = "api_key" BASIC = "basic" OAUTH_MODE_PREFERENCE = "oauth_mode_preference" +SessionCookieIdentityCredentialType: typing.TypeAlias = typing.Literal["session_cookie"] +JwtIdentityCredentialType: typing.TypeAlias = typing.Literal["jwt"] +ApiKeyIdentityCredentialType: typing.TypeAlias = typing.Literal["api_key"] +BasicIdentityCredentialType: typing.TypeAlias = typing.Literal["basic"] +IdentityCredentialType: typing.TypeAlias = (SessionCookieIdentityCredentialType + | JwtIdentityCredentialType + | ApiKeyIdentityCredentialType + | BasicIdentityCredentialType) + + class OAuthMode(StrEnum): """How the UI presents the OAuth login page.""" + REDIRECT = "redirect" POPUP = "popup" class JwtAuthPayload(BaseModel): """JWT Bearer token authentication payload.""" + model_config = ConfigDict(extra="forbid") - method: typing.Literal[AuthMethod.JWT] = Field(description="Authentication method discriminator.") + method: JwtIdentityCredentialType = Field(description="Authentication method discriminator.") token: SerializableSecretStr = Field(min_length=1, description="Encoded JWT Bearer token.") class ApiKeyAuthPayload(BaseModel): """API key authentication payload.""" + model_config = ConfigDict(extra="forbid") - method: typing.Literal[AuthMethod.API_KEY] = Field(description="Authentication method discriminator.") + method: ApiKeyIdentityCredentialType = Field(description="Authentication method discriminator.") token: SerializableSecretStr = Field(min_length=1, description="API key token.") class BasicAuthPayload(BaseModel): """Username/password authentication payload.""" + model_config = ConfigDict(extra="forbid") - method: typing.Literal[AuthMethod.BASIC] = Field(description="Authentication method discriminator.") + method: BasicIdentityCredentialType = Field(description="Authentication method discriminator.") username: str = Field(min_length=1, description="Username for basic authentication.") password: SerializableSecretStr = Field(min_length=1, description="Password for basic authentication.") @@ -775,6 +830,7 @@ class OAuthModePreferencePayload(BaseModel): This is a routing hint, not a credential: it carries no identity and does not resolve a ``user_id``. """ + model_config = ConfigDict(extra="forbid") method: typing.Literal[AuthMethod.OAUTH_MODE_PREFERENCE] = Field( description="Authentication message discriminator.") @@ -789,6 +845,7 @@ class OAuthModePreferencePayload(BaseModel): class WebSocketAuthMessage(BaseModel): """WebSocket authentication message for payload-based auth when header or cookie auth is not feasible.""" + model_config = ConfigDict(extra="forbid") type: typing.Literal[WebSocketMessageType.AUTH_MESSAGE] payload: AuthPayload @@ -797,12 +854,14 @@ class WebSocketAuthMessage(BaseModel): class AuthMessageStatus(StrEnum): """Outcome of a WebSocket authentication attempt.""" + SUCCESS = "success" ERROR = "error" class WebSocketAuthResponseMessage(BaseModel): """Server response to a WebSocket ``auth_message``.""" + model_config = ConfigDict(extra="forbid") type: typing.Literal[WebSocketMessageType.AUTH_RESPONSE] = WebSocketMessageType.AUTH_RESPONSE status: AuthMessageStatus = Field(description="Outcome of the authentication attempt.") @@ -822,6 +881,7 @@ class WebSocketSystemIntermediateStepMessage(BaseModel): For more details, refer to the API documentation: docs/source/developer_guide/websockets.md """ + # Allow extra fields in the model_config to support derived models model_config = ConfigDict(extra="allow") @@ -848,6 +908,7 @@ class WebSocketSystemResponseTokenMessage(BaseModel): For more details, refer to the API documentation: docs/source/developer_guide/websockets.md """ + # Allow extra fields in the model_config to support derived models model_config = ConfigDict(extra="allow") @@ -878,11 +939,12 @@ class WebSocketSystemInteractionMessage(BaseModel): For more details, refer to the API documentation: docs/source/developer_guide/websockets.md """ + # Allow extra fields in the model_config to support derived models model_config = ConfigDict(extra="allow") - type: typing.Literal[ - WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE + type: typing.Literal[WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = ( + WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE) id: str | None = "default" thread_id: str | None = "default" parent_id: str = "default" @@ -902,11 +964,12 @@ class WebSocketObservabilityTraceMessage(BaseModel): WebSocket message for observability trace ID. Sent once after the workflow completes to correlate the request with observability traces. """ + # Allow extra fields in the model_config to support derived models model_config = ConfigDict(extra="allow") - type: typing.Literal[ - WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE] = WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE + type: typing.Literal[WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE] = ( + WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE) id: str = "default" parent_id: str = "default" conversation_id: str | None = None @@ -1011,7 +1074,7 @@ def _nat_chat_response_to_string(data: ChatResponse) -> str: def _string_to_nat_chat_response(data: str) -> ChatResponse: - '''Converts a string to an ChatResponse object''' + """Converts a string to an ChatResponse object""" # Simulate usage prompt_tokens = 0 @@ -1039,7 +1102,7 @@ def _chat_response_chunk_to_string(data: ChatResponseChunk) -> str: def _string_to_nat_chat_response_chunk(data: str) -> ChatResponseChunk: - '''Converts a string to an ChatResponseChunk object''' + """Converts a string to an ChatResponseChunk object""" # Build and return the response return ChatResponseChunk.from_string(data) 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 e55ca78870..52313d2ea9 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 @@ -174,10 +174,14 @@ def _from_api_key(cls, api_key: str) -> "UserInfo": return cls(api_key=SecretStr(api_key)) @classmethod - def _from_jwt(cls, jwt_info: JwtUserInfo) -> "UserInfo": + def _from_jwt(cls, jwt_info: JwtUserInfo, *, issuer_scoped: bool = False) -> "UserInfo": identity: str | None = jwt_info.identity_claim if identity is None: raise ValueError("JWT contains no usable identity claim (sub, email, preferred_username)") + if issuer_scoped: + if not jwt_info.issuer: + raise ValueError("Verified JWT identity requires an issuer claim") + identity = f"{jwt_info.issuer}\x1f{identity}" instance: UserInfo = cls() object.__setattr__(instance, "_jwt", jwt_info) instance._set_user_id(identity) 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 7e098a55b3..ffcd4f387b 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 @@ -24,7 +24,10 @@ from pydantic import Field from pydantic import SerializeAsAny from pydantic import field_validator +from pydantic import model_validator +from nat.data_models.api_server import IdentityCredentialType +from nat.data_models.component_ref import AuthenticationRef from nat.data_models.component_ref import ObjectStoreRef from nat.data_models.evaluator import EvalInputItem from nat.data_models.front_end import FrontEndBaseConfig @@ -51,15 +54,17 @@ def _is_reserved(path: Path) -> bool: class EvaluateRequest(BaseModel): """Request model for the evaluate endpoint.""" + config_file: str = Field(description="Path to the configuration file for evaluation") job_id: str | None = Field(default=None, description="Unique identifier for the evaluation job") reps: int = Field(default=1, gt=0, description="Number of repetitions for the evaluation, defaults to 1") expiry_seconds: int = Field( default=3600, gt=0, - description="Optional time (in seconds) before the job expires. Clamped between 600 (10 min) and 86400 (24h).") + description="Optional time (in seconds) before the job expires. Clamped between 600 (10 min) and 86400 (24h).", + ) - @field_validator('job_id', mode='after') + @field_validator("job_id", mode="after") @classmethod def validate_job_id(cls, job_id: str): job_id = job_id.strip() @@ -75,7 +80,7 @@ def validate_job_id(cls, job_id: str): return job_id - @field_validator('config_file', mode='after') + @field_validator("config_file", mode="after") @classmethod def validate_config_file(cls, config_file: str): config_file = config_file.strip() @@ -98,22 +103,26 @@ def validate_config_file(cls, config_file: str): class BaseAsyncResponse(BaseModel): """Base model for async responses.""" + job_id: str = Field(description="Unique identifier for the job") status: str = Field(description="Current status of the job") class EvaluateResponse(BaseAsyncResponse): """Response model for the evaluate endpoint.""" + pass class AsyncGenerateResponse(BaseAsyncResponse): """Response model for the async generation endpoint.""" + pass class BaseAsyncStatusResponse(BaseModel): """Base model for async status responses.""" + job_id: str = Field(description="Unique identifier for the evaluation job") status: str = Field(description="Current status of the evaluation job") error: str | None = Field(default=None, description="Error message if the job failed") @@ -124,6 +133,7 @@ class BaseAsyncStatusResponse(BaseModel): class EvaluateStatusResponse(BaseAsyncStatusResponse): """Response model for the evaluate status endpoint.""" + config_file: str = Field(description="Path to the configuration file used for evaluation") output_path: str | None = Field(default=None, description="Path to the output file if the job completed successfully") @@ -132,17 +142,20 @@ class EvaluateStatusResponse(BaseAsyncStatusResponse): class AsyncGenerationStatusResponse(BaseAsyncStatusResponse): output: dict | None = Field( default=None, - description="Output of the generate request, this is only available if the job completed successfully.") + description="Output of the generate request, this is only available if the job completed successfully.", + ) class EvaluateItemRequest(BaseModel): """Request model for single-item evaluation endpoint.""" + item: EvalInputItem = Field(description="Single evaluation input item to evaluate") evaluator_name: str = Field(description="Name of the evaluator to use (must match config)") class EvaluateItemResponse(BaseModel): """Response model for single-item evaluation endpoint.""" + success: bool = Field(description="Whether the evaluation completed successfully") result: SerializeAsAny[BaseModel] | None = Field(default=None, description="Evaluation result if successful") error: str | None = Field(default=None, description="Error message if evaluation failed") @@ -154,7 +167,6 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"): """ class EndpointBase(BaseModel): - method: typing.Literal["GET", "POST", "PUT", "DELETE"] description: str path: str | None = Field( @@ -200,11 +212,13 @@ class CrossOriginResourceSharing(BaseModel): description="A permitted regex string to match against origins to make cross-origin requests", ) allow_methods: list[str] | None = Field( - default_factory=lambda: ['GET'], - description="A list of HTTP methods that should be allowed for cross-origin requests.") + default_factory=lambda: ["GET"], + description="A list of HTTP methods that should be allowed for cross-origin requests.", + ) allow_headers: list[str] | None = Field( default_factory=list, - description="A list of HTTP request headers that should be supported for cross-origin requests.") + description="A list of HTTP request headers that should be supported for cross-origin requests.", + ) allow_credentials: bool | None = Field( default=False, description="Indicate that cookies should be supported for cross-origin requests.", @@ -226,11 +240,13 @@ class CrossOriginResourceSharing(BaseModel): scheduler_address: str | None = Field( default=None, description=("Address of the Dask scheduler to use for async jobs. If None, a Dask local cluster is created. " - "Note: This requires the optional dask dependency to be installed.")) + "Note: This requires the optional dask dependency to be installed."), + ) db_url: str | None = Field( default=None, - description= - "SQLAlchemy database URL for storing async job metadata, if unset a temporary SQLite database is used.") + description=( + "SQLAlchemy database URL for storing async job metadata, if unset a temporary SQLite database is used."), + ) max_running_async_jobs: int = Field( default=10, description=( @@ -238,7 +254,8 @@ class CrossOriginResourceSharing(BaseModel): "misleading as the actual number of concurrent async jobs is: " "`max_running_async_jobs * dask_threads_per_worker`. " "This parameter is only used when scheduler_address is `None` and a Dask local cluster is created."), - ge=1) + ge=1, + ) dask_workers: typing.Literal["threads", "processes"] = Field( default="processes", description=( @@ -254,16 +271,38 @@ class CrossOriginResourceSharing(BaseModel): default="0", description=("Memory limit for each Dask worker. Can be 'auto', a memory string like '4GB' or a float " "representing a fraction of the system memory. Default is '0' which means no limit. " - "Refer to https://docs.dask.org/en/stable/deploying-python.html#reference for details.")) + "Refer to https://docs.dask.org/en/stable/deploying-python.html#reference for details."), + ) dask_threads_per_worker: int = Field( default=1, description=( "Number of threads to use per worker. This parameter is only used when the value is greater than 0 and " "scheduler_address is `None` and a local Dask cluster is created. When set to 0 the value uses the Dask " - "default.")) + "default."), + ) step_adaptor: StepAdaptorConfig = StepAdaptorConfig() + accepted_identity_credentials: list[IdentityCredentialType] | None = Field( + default=None, + description=("Identity credential methods accepted for WebSocket connections and auth messages. " + "If omitted, all supported methods are accepted."), + ) + + identity_authentication: list[AuthenticationRef] = Field( + default_factory=list, + description=("Named JWT authentication providers used to verify WebSocket identity credentials. " + "If omitted, JWT claims retain the existing decode-only behavior."), + ) + + @model_validator(mode="after") + def validate_jwt_identity_policy(self) -> typing.Self: + if (self.identity_authentication and self.accepted_identity_credentials is not None + and "jwt" not in self.accepted_identity_credentials): + raise ValueError( + "identity_authentication cannot be configured when jwt identity credentials are not accepted") + return self + workflow: typing.Annotated[EndpointBase, Field(description="Endpoint for the default workflow.")] = EndpointBase( method="POST", path="/v1/workflow", @@ -282,24 +321,27 @@ class CrossOriginResourceSharing(BaseModel): ) evaluate_item: typing.Annotated[EndpointBase, - Field(description="Endpoint for evaluating a single item.")] = EndpointBase( + Field(description="Endpoint for evaluating a single item.")] = (EndpointBase( method="POST", path="/evaluate/item", description="Evaluate a single item with a specified evaluator", - ) + )) oauth2_callback_path: str | None = Field( default="/auth/redirect", - description="OAuth2.0 authentication callback endpoint. If None, no OAuth2 callback endpoint is created.") + description="OAuth2.0 authentication callback endpoint. If None, no OAuth2 callback endpoint is created.", + ) endpoints: list[Endpoint] = Field( default_factory=list, description=("Additional endpoints to add to the FastAPI app which run functions within the NAT configuration. " - "Each endpoint must have a unique path.")) + "Each endpoint must have a unique path."), + ) cors: CrossOriginResourceSharing = Field( default_factory=CrossOriginResourceSharing, - description="Cross origin resource sharing configuration for the FastAPI app") + description="Cross origin resource sharing configuration for the FastAPI app", + ) use_gunicorn: bool = Field( default=False, @@ -317,13 +359,16 @@ class CrossOriginResourceSharing(BaseModel): description=( "Object store reference for the FastAPI app. If present, static files can be uploaded via a POST " "request to '/static' and files will be served from the object store. The files will be served from the " - "object store at '/static/{file_name}'.")) + "object store at '/static/{file_name}'."), + ) disable_legacy_routes: bool = Field( default=False, - description="Disable the legacy routes for the FastAPI app. If True, the legacy routes are disabled.") + description="Disable the legacy routes for the FastAPI app. If True, the legacy routes are disabled.", + ) enable_interactive_extensions: bool = Field( default=False, description=("Enable the interactive extensions for OpenAI API compatible endpoints." + - " If True, the interactive extensions are enabled.")) + " If True, the interactive extensions are enabled."), + ) 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 76a0dc6de1..7a1d42fb34 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 @@ -27,6 +27,7 @@ from pydantic import ValidationError from starlette.websockets import WebSocketDisconnect +from nat.authentication.credential_validator.bearer_token_validator import BearerTokenValidator from nat.authentication.interfaces import FlowHandlerBase from nat.data_models.api_server import AuthMessageStatus from nat.data_models.api_server import ChatRequest @@ -34,6 +35,7 @@ from nat.data_models.api_server import ChatResponseChunk from nat.data_models.api_server import Error from nat.data_models.api_server import ErrorTypes +from nat.data_models.api_server import IdentityCredentialType from nat.data_models.api_server import OAuthModePreferencePayload from nat.data_models.api_server import ResponseObservabilityTrace from nat.data_models.api_server import ResponsePayloadOutput @@ -63,6 +65,8 @@ from nat.front_ends.fastapi.response_helpers import generate_streaming_response 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 JwtVerificationError from nat.runtime.user_manager import UserManager if typing.TYPE_CHECKING: @@ -82,18 +86,23 @@ class UserInteraction(BaseModel): class WebSocketMessageHandler: - _HITL_TIMEOUT_GRACE_PERIOD_SECONDS: int = 5 - def __init__(self, - socket: WebSocket, - session_manager: SessionManager, - step_adaptor: StepAdaptor, - worker: "FastApiFrontEndPluginWorker"): + def __init__( + self, + socket: WebSocket, + session_manager: SessionManager, + step_adaptor: StepAdaptor, + worker: "FastApiFrontEndPluginWorker", + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + jwt_validators: typing.Mapping[str, BearerTokenValidator] | None = None, + ): self._socket: WebSocket = socket self._session_manager: SessionManager = session_manager self._step_adaptor: StepAdaptor = step_adaptor self._worker: FastApiFrontEndPluginWorker = worker + self._accepted_identity_credentials = accepted_identity_credentials + self._jwt_validators = jwt_validators self._message_validator: MessageValidator = MessageValidator() self._running_workflow_task: asyncio.Task | None = None @@ -104,6 +113,7 @@ def __init__(self, self._pending_observability_trace: ResponseObservabilityTrace | None = None self._user_id: str | None = None self._restoration_attempted: bool = False + self._connection_rejected: bool = False self._flow_handler: FlowHandlerBase | None = None @@ -169,13 +179,33 @@ async def _restore_execution_state(self) -> None: # Copy the original timeout so it is preserved for subsequent reconnections prompt_content = prompt_content.model_copy(update={"timeout": time_remaining_in_seconds}) - await self.create_websocket_message(data_model=prompt_content, - message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE, - status=WebSocketMessageStatus.IN_PROGRESS) + await self.create_websocket_message( + data_model=prompt_content, + message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE, + status=WebSocketMessageStatus.IN_PROGRESS, + ) async def __aenter__(self) -> "WebSocketMessageHandler": await self._socket.accept() - user_info = UserManager.extract_user_from_connection(self._socket) + try: + user_info = await UserManager.extract_user_from_connection_with_verification( + self._socket, + accepted_identity_credentials=self._accepted_identity_credentials, + jwt_validators=self._jwt_validators, + ) + except (IdentityCredentialNotAcceptedError, JwtVerificationError) as exc: + self._connection_rejected = True + response = WebSocketAuthResponseMessage( + status=AuthMessageStatus.ERROR, + payload=Error( + code=ErrorTypes.USER_AUTH_ERROR, + message="Authentication failed", + details=str(exc), + ), + ) + await self._socket.send_json(response.model_dump()) + await self._socket.close(code=1008, reason="Identity credential was rejected") + return self if user_info is not None: self._user_id = user_info.get_user_id() await self._restore_execution_state() @@ -213,25 +243,26 @@ async def run(self) -> None: Preflight auth runs concurrently with the receive loop so connection-level messages (e.g. ``oauth_mode_preference``) are handled while a preflight OAuth flow awaits user login. """ + if self._connection_rejected: + return + preflight_task: asyncio.Task = asyncio.create_task(self._run_preflight_auth()) try: while True: - try: - message: dict[str, Any] = await self._socket.receive_json() validated_message: BaseModel = await self._message_validator.validate_message(message) # Received a request to start a workflow - if (isinstance(validated_message, WebSocketUserMessage)): + if isinstance(validated_message, WebSocketUserMessage): await self.process_workflow_request(validated_message) - elif (isinstance(validated_message, WebSocketAuthMessage)): + elif isinstance(validated_message, WebSocketAuthMessage): await self._process_auth_message(validated_message) - elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)): + elif isinstance(validated_message, WebSocketUserInteractionResponseMessage): user_content = await self._process_websocket_user_interaction_response_message(validated_message ) assert self._user_interaction is not None @@ -277,13 +308,18 @@ async def _process_auth_message(self, message: WebSocketAuthMessage) -> None: """Resolve user identity, or record a non-identity routing hint (OAuth mode).""" if isinstance(message.payload, OAuthModePreferencePayload): from nat.front_ends.fastapi.auth_flow_handlers import websocket_flow_handler + if isinstance(self._flow_handler, websocket_flow_handler.WebSocketAuthenticationFlowHandler): self._flow_handler.set_oauth_mode(message.payload.mode) return identity_resolved = False try: - user_info: UserInfo = UserManager._from_auth_payload(message.payload) + user_info: UserInfo = await UserManager.from_auth_payload_with_verification( + message.payload, + accepted_identity_credentials=self._accepted_identity_credentials, + jwt_validators=self._jwt_validators, + ) self._user_id = user_info.get_user_id() identity_resolved = True response: WebSocketAuthResponseMessage = WebSocketAuthResponseMessage( @@ -351,44 +387,54 @@ async def process_workflow_request(self, user_message_as_validated_type: WebSock def _done_callback(_task: asyncio.Task): if self._running_workflow_task is _task: self._running_workflow_task = None - if self._running_workflow_task is None and _user_id and _conversation_id and \ - self._worker.get_conversation_handler(_user_id, _conversation_id) is self: + if (self._running_workflow_task is None and _user_id and _conversation_id + and self._worker.get_conversation_handler(_user_id, _conversation_id) is self): self._worker.remove_conversation_handler(_user_id, _conversation_id) # Only the *_STREAM schemas stream; others aggregate a single result. Streaming a # non-streaming schema converts chunks to the single output schema and raises. - streaming = self._workflow_schema_type in (WorkflowSchemaType.CHAT_STREAM, - WorkflowSchemaType.GENERATE_STREAM) + streaming = self._workflow_schema_type in ( + WorkflowSchemaType.CHAT_STREAM, + WorkflowSchemaType.GENERATE_STREAM, + ) self._running_workflow_task = asyncio.create_task( - self._run_workflow(payload=message_content, - user_message_id=self._message_parent_id, - conversation_id=self._conversation_id, - streaming=streaming, - result_type=self._schema_output_mapping[self._workflow_schema_type], - output_type=self._schema_output_mapping[self._workflow_schema_type])) + self._run_workflow( + payload=message_content, + user_message_id=self._message_parent_id, + conversation_id=self._conversation_id, + streaming=streaming, + result_type=self._schema_output_mapping[self._workflow_schema_type], + output_type=self._schema_output_mapping[self._workflow_schema_type], + )) self._running_workflow_task.add_done_callback(_done_callback) except ValueError as e: logger.exception("User message content not found: %s", str(e)) - await self.create_websocket_message(data_model=Error(code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT, - message="User message content could not be found", - details=str(e)), - message_type=WebSocketMessageType.ERROR_MESSAGE, - status=WebSocketMessageStatus.IN_PROGRESS) + await self.create_websocket_message( + data_model=Error( + code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT, + message="User message content could not be found", + details=str(e), + ), + message_type=WebSocketMessageType.ERROR_MESSAGE, + status=WebSocketMessageStatus.IN_PROGRESS, + ) except RuntimeError as e: logger.exception("Internal workflow initialization error: %s", str(e)) - await self.create_websocket_message(data_model=Error(code=ErrorTypes.WORKFLOW_ERROR, - message=type(e).__name__, - details=str(e)), - message_type=WebSocketMessageType.ERROR_MESSAGE, - status=WebSocketMessageStatus.IN_PROGRESS) - - async def create_websocket_message(self, - data_model: BaseModel, - message_type: str | None = None, - status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS) -> None: + await self.create_websocket_message( + data_model=Error(code=ErrorTypes.WORKFLOW_ERROR, message=type(e).__name__, details=str(e)), + message_type=WebSocketMessageType.ERROR_MESSAGE, + status=WebSocketMessageStatus.IN_PROGRESS, + ) + + async def create_websocket_message( + self, + data_model: BaseModel, + message_type: str | None = None, + status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS, + ) -> None: """ Creates a websocket message that will be ready for routing based on message type or data model. @@ -405,8 +451,8 @@ async def create_websocket_message(self, message_schema: type[BaseModel] = await self._message_validator.get_message_schema_by_type(message_type) - if hasattr(data_model, 'id'): - message_id: str = str(getattr(data_model, 'id')) + if hasattr(data_model, "id"): + message_id: str = str(getattr(data_model, "id")) else: message_id = str(uuid.uuid4()) @@ -419,7 +465,8 @@ async def create_websocket_message(self, parent_id=self._message_parent_id, conversation_id=self._conversation_id, content=content, - status=status) + status=status, + ) elif issubclass(message_schema, WebSocketSystemIntermediateStepMessage): message = await self._message_validator.create_system_intermediate_step_message( @@ -427,7 +474,8 @@ async def create_websocket_message(self, parent_id=await self._message_validator.get_intermediate_step_parent_id(data_model), conversation_id=self._conversation_id, content=content, - status=status) + status=status, + ) elif issubclass(message_schema, WebSocketSystemInteractionMessage): message = await self._message_validator.create_system_interaction_message( @@ -435,14 +483,16 @@ async def create_websocket_message(self, parent_id=self._message_parent_id, conversation_id=self._conversation_id, content=content, - status=status) + status=status, + ) elif issubclass(message_schema, WebSocketObservabilityTraceMessage): message = await self._message_validator.create_observability_trace_message( message_id=message_id, parent_id=self._message_parent_id, conversation_id=self._conversation_id, - content=content) + content=content, + ) elif isinstance(content, Error): raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}") @@ -450,7 +500,7 @@ async def create_websocket_message(self, elif issubclass(message_schema, Error): raise TypeError(f"Invalid message type: {message_type}") - elif (message is None): + elif message is None: raise ValueError( f"Message type could not be resolved by input data model: {data_model.model_dump_json()}") @@ -459,10 +509,11 @@ async def create_websocket_message(self, message = await self._message_validator.create_system_response_token_message( message_type=WebSocketMessageType.ERROR_MESSAGE, conversation_id=self._conversation_id, - content=Error(code=ErrorTypes.WORKFLOW_ERROR, message=type(e).__name__, details=str(e))) + content=Error(code=ErrorTypes.WORKFLOW_ERROR, message=type(e).__name__, details=str(e)), + ) finally: - if (message is not None): + if message is not None: await self._socket.send_json(message.model_dump()) async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanResponse: @@ -486,12 +537,13 @@ async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanRe started_at=time.monotonic()) try: - await self.create_websocket_message(data_model=prompt.content, - message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE, - status=WebSocketMessageStatus.IN_PROGRESS) - - if (isinstance(prompt.content, HumanPromptNotification)): + await self.create_websocket_message( + data_model=prompt.content, + message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE, + status=WebSocketMessageStatus.IN_PROGRESS, + ) + if isinstance(prompt.content, HumanPromptNotification): return HumanResponseNotification() backend_timeout_in_seconds: int | None = (prompt.content.timeout + self._HITL_TIMEOUT_GRACE_PERIOD_SECONDS @@ -512,31 +564,36 @@ async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanRe # Delete the future from the outstanding human prompts dictionary self._user_interaction = None - async def _run_workflow(self, - payload: typing.Any, - user_message_id: str | None = None, - conversation_id: str | None = None, - streaming: bool = True, - result_type: type | None = None, - output_type: type | None = None) -> None: + async def _run_workflow( + self, + payload: typing.Any, + user_message_id: str | None = None, + conversation_id: str | None = None, + streaming: bool = True, + result_type: type | None = None, + output_type: type | None = None, + ) -> None: _cancelled = False try: auth_callback = self._flow_handler.authenticate if self._flow_handler else None - async with self._session_manager.session(user_id=self._user_id, - user_message_id=user_message_id, - conversation_id=conversation_id, - http_connection=self._socket, - user_input_callback=self.human_interaction_callback, - user_authentication_callback=auth_callback) as session: + async with self._session_manager.session( + user_id=self._user_id, + user_message_id=user_message_id, + conversation_id=conversation_id, + http_connection=self._socket, + user_input_callback=self.human_interaction_callback, + user_authentication_callback=auth_callback, + ) as session: self._session_manager._context.metadata._request.payload = self._user_message_payload - async for value in generate_streaming_response(payload, - session=session, - streaming=streaming, - step_adaptor=self._step_adaptor, - result_type=result_type, - output_type=output_type): - + async for value in generate_streaming_response( + payload, + session=session, + streaming=streaming, + step_adaptor=self._step_adaptor, + result_type=result_type, + output_type=output_type, + ): # Store observability trace to send after completion message if isinstance(value, ResponseObservabilityTrace): if self._pending_observability_trace is None: @@ -554,23 +611,26 @@ async def _run_workflow(self, except Exception as e: logger.exception("Unhandled workflow error") - await self.create_websocket_message(data_model=Error(code=ErrorTypes.WORKFLOW_ERROR, - message=type(e).__name__, - details=str(e)), - message_type=WebSocketMessageType.ERROR_MESSAGE, - status=WebSocketMessageStatus.IN_PROGRESS) + await self.create_websocket_message( + data_model=Error(code=ErrorTypes.WORKFLOW_ERROR, message=type(e).__name__, details=str(e)), + message_type=WebSocketMessageType.ERROR_MESSAGE, + status=WebSocketMessageStatus.IN_PROGRESS, + ) finally: try: if not _cancelled: - await self.create_websocket_message(data_model=SystemResponseContent(), - message_type=WebSocketMessageType.RESPONSE_MESSAGE, - status=WebSocketMessageStatus.COMPLETE) + await self.create_websocket_message( + data_model=SystemResponseContent(), + message_type=WebSocketMessageType.RESPONSE_MESSAGE, + status=WebSocketMessageStatus.COMPLETE, + ) # Send observability trace after completion message if self._pending_observability_trace is not None: await self.create_websocket_message( data_model=self._pending_observability_trace, - message_type=WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE) + message_type=WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE, + ) finally: self._pending_observability_trace = None 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 3dda7548ac..8e5828950e 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 @@ -21,6 +21,8 @@ from fastapi import FastAPI from starlette.websockets import WebSocket +from nat.authentication.credential_validator.bearer_token_validator import BearerTokenValidator +from nat.authentication.jwt.jwt_auth_provider import JwtAuthProvider from nat.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import WebSocketAuthenticationFlowHandler from nat.front_ends.fastapi.message_handler import WebSocketMessageHandler from nat.runtime.session import SESSION_COOKIE_NAME @@ -29,7 +31,7 @@ logger = logging.getLogger(__name__) # Only allow URL-safe characters in session IDs (alphanumeric, hyphen, underscore, period, tilde). -_SAFE_SESSION_ID_RE = re.compile(r'^[A-Za-z0-9\-_.~]+$') +_SAFE_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9\-_.~]+$") def _is_origin_allowed(origin: str | None, allowed_origins: list[str], allow_origin_regex: str | None) -> bool: @@ -51,7 +53,22 @@ def _is_origin_allowed(origin: str | None, allowed_origins: list[str], allow_ori return False -def websocket_endpoint(*, worker: Any, session_manager: SessionManager): +async def _get_jwt_validators(worker: Any, session_manager: SessionManager) -> dict[str, BearerTokenValidator]: + validators: dict[str, BearerTokenValidator] = {} + for provider_name in worker.front_end_config.identity_authentication: + provider = await session_manager.shared_builder.get_auth_provider(provider_name) + if not isinstance(provider, JwtAuthProvider): + raise ValueError(f"Identity authentication provider '{provider_name}' must have _type: jwt") + + issuer = provider.config.issuer_url + if issuer in validators: + raise ValueError(f"Duplicate JWT issuer configured for identity authentication: {issuer}") + validators[issuer] = provider.validator + return validators + + +def websocket_endpoint(*, worker: Any, session_manager: SessionManager, jwt_validators: dict[str, + BearerTokenValidator]): """Build websocket endpoint handler with auth-flow integration.""" async def _websocket_endpoint(websocket: WebSocket): @@ -90,7 +107,14 @@ async def _websocket_endpoint(websocket: WebSocket): websocket.scope["headers"] = headers - async with WebSocketMessageHandler(websocket, session_manager, worker.get_step_adaptor(), worker) as handler: + async with WebSocketMessageHandler( + websocket, + session_manager, + worker.get_step_adaptor(), + worker, + accepted_identity_credentials=worker.front_end_config.accepted_identity_credentials, + jwt_validators=jwt_validators, + ) as handler: origin = websocket.headers.get("origin") allowed_origins = worker.front_end_config.cors.allow_origins or [] allow_origin_regex = worker.front_end_config.cors.allow_origin_regex @@ -113,8 +137,12 @@ async def add_websocket_routes( ): """Add websocket route for an endpoint.""" if endpoint.websocket_path: - app.add_api_websocket_route(endpoint.websocket_path, - websocket_endpoint( - worker=worker, - session_manager=session_manager, - )) + jwt_validators = await _get_jwt_validators(worker, session_manager) + app.add_api_websocket_route( + endpoint.websocket_path, + websocket_endpoint( + worker=worker, + session_manager=session_manager, + jwt_validators=jwt_validators, + ), + ) 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 d67a252fda..2b1d0a4d47 100644 --- a/packages/nvidia_nat_core/src/nat/runtime/user_manager.py +++ b/packages/nvidia_nat_core/src/nat/runtime/user_manager.py @@ -25,10 +25,12 @@ from pydantic import SecretStr from starlette.requests import Request +from nat.authentication.credential_validator.bearer_token_validator import BearerTokenValidator from nat.authentication.jwt_utils import decode_jwt_claims_unverified from nat.data_models.api_server import ApiKeyAuthPayload from nat.data_models.api_server import AuthPayload from nat.data_models.api_server import BasicAuthPayload +from nat.data_models.api_server import IdentityCredentialType from nat.data_models.api_server import JwtAuthPayload from nat.data_models.user_info import BasicUserInfo from nat.data_models.user_info import JwtUserInfo @@ -37,11 +39,23 @@ logger = logging.getLogger(__name__) +class IdentityCredentialNotAcceptedError(ValueError): + """Raised when a supplied identity credential method is disabled by policy.""" + + +class JwtVerificationError(ValueError): + """Raised when an enabled JWT verification policy rejects a token.""" + + class UserManager: """Stateless resolver that creates ``UserInfo`` from HTTP/WebSocket connections.""" @classmethod - def extract_user_from_connection(cls, connection: Request | WebSocket) -> UserInfo | None: + def extract_user_from_connection( + cls, + connection: Request | WebSocket, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + ) -> UserInfo | None: """Resolve an HTTP/WebSocket connection into a ``UserInfo``. Args: @@ -57,22 +71,79 @@ def extract_user_from_connection(cls, connection: Request | WebSocket) -> UserIn """ cookie: str | None = cls._get_session_cookie(connection) if cookie: + cls._ensure_identity_credential_accepted("session_cookie", accepted_identity_credentials) return cls._user_info_from_session_cookie(cookie) auth_header: str | None = cls._get_auth_header(connection) if auth_header: - resolved: UserInfo | None = cls._resolve_from_auth_header(auth_header) + resolved: UserInfo | None = cls._resolve_from_auth_header(auth_header, accepted_identity_credentials) if resolved is not None: return resolved api_key: str | None = cls._get_api_key_header(connection) if api_key: + cls._ensure_identity_credential_accepted("api_key", accepted_identity_credentials) + return UserInfo._from_api_key(api_key) + + return None + + @classmethod + async def extract_user_from_connection_with_verification( + cls, + connection: Request | WebSocket, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + jwt_validators: typing.Mapping[str, BearerTokenValidator] | None = None, + ) -> UserInfo | None: + cookie = cls._get_session_cookie(connection) + if cookie: + cls._ensure_identity_credential_accepted("session_cookie", accepted_identity_credentials) + return cls._user_info_from_session_cookie(cookie) + + auth_header = cls._get_auth_header(connection) + if auth_header: + resolved = await cls._resolve_from_auth_header_with_verification( + auth_header, + accepted_identity_credentials=accepted_identity_credentials, + jwt_validators=jwt_validators, + ) + if resolved is not None: + return resolved + + api_key = cls._get_api_key_header(connection) + if api_key: + cls._ensure_identity_credential_accepted("api_key", accepted_identity_credentials) return UserInfo._from_api_key(api_key) return None @classmethod - def _resolve_from_auth_header(cls, auth_header: str) -> UserInfo | None: + async def _resolve_from_auth_header_with_verification( + cls, + auth_header: str, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + jwt_validators: typing.Mapping[str, BearerTokenValidator] | None = None, + ) -> UserInfo | None: + parts = auth_header.strip().split(maxsplit=1) + if len(parts) != 2: + return None + + scheme, credential = parts[0].lower(), parts[1] + if not credential: + return None + + if scheme == "bearer" and credential.count(".") == 2: + cls._ensure_identity_credential_accepted("jwt", accepted_identity_credentials) + claims = await cls._decode_verified_jwt(credential, jwt_validators) + return cls._user_info_from_jwt(claims, issuer_scoped=bool(jwt_validators)) + + return cls._resolve_from_auth_header(auth_header, accepted_identity_credentials) + + @classmethod + def _resolve_from_auth_header( + cls, + auth_header: str, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + ) -> UserInfo | None: """Parse an ``Authorization`` header and resolve identity by scheme. Args: @@ -98,17 +169,23 @@ def _resolve_from_auth_header(cls, auth_header: str) -> UserInfo | None: if scheme == "bearer": if credential.count(".") == 2: + cls._ensure_identity_credential_accepted("jwt", accepted_identity_credentials) claims: dict[str, typing.Any] = decode_jwt_claims_unverified(credential) return cls._user_info_from_jwt(claims) + cls._ensure_identity_credential_accepted("api_key", accepted_identity_credentials) return UserInfo._from_api_key(credential) if scheme == "basic": + cls._ensure_identity_credential_accepted("basic", accepted_identity_credentials) return cls._user_info_from_basic_auth(credential) return None @staticmethod - def _from_auth_payload(payload: AuthPayload) -> UserInfo: + def _from_auth_payload( + payload: AuthPayload, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + ) -> UserInfo: """Resolve a ``UserInfo`` from a WebSocket auth message payload. This is an identity resolver, not an authenticator. JWTs are decoded @@ -127,17 +204,20 @@ def _from_auth_payload(payload: AuthPayload) -> UserInfo: ValueError: If the payload cannot be resolved to a valid user identity. """ if isinstance(payload, JwtAuthPayload): + UserManager._ensure_identity_credential_accepted("jwt", accepted_identity_credentials) raw_token: str = payload.token.get_secret_value() claims: dict[str, typing.Any] = decode_jwt_claims_unverified(raw_token) return UserManager._user_info_from_jwt(claims) if isinstance(payload, ApiKeyAuthPayload): + UserManager._ensure_identity_credential_accepted("api_key", accepted_identity_credentials) token_value: str = payload.token.get_secret_value() if not token_value: raise ValueError("API key token is empty") return UserInfo._from_api_key(token_value) if isinstance(payload, BasicAuthPayload): + UserManager._ensure_identity_credential_accepted("basic", accepted_identity_credentials) return UserInfo(basic_user=BasicUserInfo( username=payload.username, password=payload.password, @@ -145,6 +225,52 @@ def _from_auth_payload(payload: AuthPayload) -> UserInfo: typing.assert_never(payload) + @staticmethod + async def from_auth_payload_with_verification( + payload: AuthPayload, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None = None, + jwt_validators: typing.Mapping[str, BearerTokenValidator] | None = None, + ) -> UserInfo: + """Resolve an auth payload, verifying JWTs when a verifier is configured.""" + if isinstance(payload, JwtAuthPayload) and jwt_validators: + UserManager._ensure_identity_credential_accepted("jwt", accepted_identity_credentials) + raw_token = payload.token.get_secret_value() + claims = await UserManager._decode_verified_jwt(raw_token, jwt_validators) + return UserManager._user_info_from_jwt(claims, issuer_scoped=True) + + return UserManager._from_auth_payload(payload, accepted_identity_credentials) + + @staticmethod + async def _decode_verified_jwt( + token: str, + jwt_validators: typing.Mapping[str, BearerTokenValidator] | None, + ) -> dict[str, typing.Any]: + claims = decode_jwt_claims_unverified(token) + if not jwt_validators: + return claims + + issuer = claims.get("iss") + if not isinstance(issuer, str) or not issuer: + raise JwtVerificationError("JWT verification requires a non-empty issuer claim") + + jwt_validator = jwt_validators.get(issuer) + if jwt_validator is None: + raise JwtVerificationError(f"JWT issuer is not accepted: {issuer}") + + result = await jwt_validator.verify(token) + if not result.active: + raise JwtVerificationError("JWT verification failed") + return claims + + @staticmethod + def _ensure_identity_credential_accepted( + credential_type: IdentityCredentialType, + accepted_identity_credentials: typing.Collection[IdentityCredentialType] | None, + ) -> None: + """Reject a recognized credential type when it is disabled by policy.""" + 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_session_cookie(connection: Request | WebSocket) -> str | None: """Extract the ``nat-session`` cookie value from a Request or WebSocket.""" @@ -207,7 +333,7 @@ def _user_info_from_session_cookie(cookie_value: str) -> UserInfo: return UserInfo._from_session_cookie(cookie_value) @staticmethod - def _user_info_from_jwt(claims: dict[str, typing.Any]) -> UserInfo: + def _user_info_from_jwt(claims: dict[str, typing.Any], *, issuer_scoped: bool = False) -> UserInfo: """Build a ``UserInfo`` from decoded JWT claims. Registered claims (``sub``, ``iss``, ``aud``, ``exp``, ``iat``) follow @@ -224,8 +350,8 @@ def _user_info_from_jwt(claims: dict[str, typing.Any]) -> UserInfo: if not has_identity: raise ValueError("JWT contains no usable identity claim (sub, email, preferred_username)") - given_name: str | None = (claims.get("given_name") if isinstance(claims.get("given_name"), str) else None) - family_name: str | None = (claims.get("family_name") if isinstance(claims.get("family_name"), str) else None) + given_name: str | None = claims.get("given_name") if isinstance(claims.get("given_name"), str) else None + family_name: str | None = claims.get("family_name") if isinstance(claims.get("family_name"), str) else None if not given_name and not family_name: raw_name: typing.Any = claims.get("name") if isinstance(raw_name, str) and raw_name.strip(): @@ -271,7 +397,7 @@ def _user_info_from_jwt(claims: dict[str, typing.Any]) -> UserInfo: if isinstance(claims.get("azp"), str) or isinstance(claims.get("client_id"), str) else None), claims=claims, ) - return UserInfo._from_jwt(jwt_info) + return UserInfo._from_jwt(jwt_info, issuer_scoped=issuer_scoped) @staticmethod def _user_info_from_basic_auth(b64_credential: str) -> UserInfo: diff --git a/packages/nvidia_nat_core/tests/nat/authentication/test_jwt_auth_provider.py b/packages/nvidia_nat_core/tests/nat/authentication/test_jwt_auth_provider.py new file mode 100644 index 0000000000..ccbb063190 --- /dev/null +++ b/packages/nvidia_nat_core/tests/nat/authentication/test_jwt_auth_provider.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import AsyncMock +from unittest.mock import patch + +import pytest + +from nat.authentication.jwt.jwt_auth_provider import JwtAuthProvider +from nat.authentication.jwt.jwt_auth_provider_config import JwtAuthProviderConfig +from nat.data_models.authentication import TokenValidationResult + + +def _provider() -> JwtAuthProvider: + return JwtAuthProvider( + JwtAuthProviderConfig( + issuer_url="https://identity.example.com", + jwks_uri="https://identity.example.com/jwks.json", + audience="nat-api", + )) + + +async def test_authenticate_returns_verified_token_result(): + provider = _provider() + validation_result = TokenValidationResult( + client_id="nat-client", + subject="verified-user", + issuer="https://identity.example.com", + token_type="bearer", + active=True, + ) + with patch.object(provider, "verify", AsyncMock(return_value=validation_result)) as verify: + result = await provider.authenticate(token="signed-token") + + verify.assert_awaited_once_with("signed-token") + assert result.raw == validation_result.model_dump() + + +@pytest.mark.parametrize("token", [None, "", 123], ids=["missing", "empty", "non-string"]) +async def test_authenticate_rejects_invalid_token_input(token): + provider = _provider() + + with pytest.raises(ValueError, match="requires a token"): + await provider.authenticate(token=token) + + +async def test_authenticate_rejects_inactive_token(): + provider = _provider() + validation_result = TokenValidationResult(client_id=None, token_type="bearer", active=False) + + with patch.object(provider, "verify", AsyncMock(return_value=validation_result)) as verify: + with pytest.raises(ValueError, match="JWT verification failed"): + await provider.authenticate(token="rejected-token") + + verify.assert_awaited_once_with("rejected-token") diff --git a/packages/nvidia_nat_core/tests/nat/authentication/test_jwt_auth_provider_config.py b/packages/nvidia_nat_core/tests/nat/authentication/test_jwt_auth_provider_config.py new file mode 100644 index 0000000000..8589991796 --- /dev/null +++ b/packages/nvidia_nat_core/tests/nat/authentication/test_jwt_auth_provider_config.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from pydantic import ValidationError + +from nat.authentication.jwt.jwt_auth_provider_config import JwtAuthProviderConfig + + +def test_jwt_auth_provider_requires_complete_verification_policy(): + with pytest.raises(ValidationError): + JwtAuthProviderConfig(issuer_url="https://identity.example.com") + + +def test_jwt_auth_provider_accepts_complete_verification_policy(): + config = JwtAuthProviderConfig( + issuer_url="https://identity.example.com", + jwks_uri="https://identity.example.com/jwks.json", + audience="nat-api", + scopes=["workflow:resume"], + ) + + assert config.issuer_url == "https://identity.example.com" + assert config.jwks_uri == "https://identity.example.com/jwks.json" + assert config.audience == "nat-api" + assert config.scopes == ["workflow:resume"] + + +@pytest.mark.parametrize("field", ["issuer_url", "jwks_uri"]) +def test_jwt_auth_provider_rejects_insecure_remote_urls(field: str): + values = { + "issuer_url": "https://identity.example.com", + "jwks_uri": "https://identity.example.com/jwks.json", + "audience": "nat-api", + } + values[field] = "http://identity.example.com" + + with pytest.raises(ValidationError, match="must use HTTPS"): + JwtAuthProviderConfig(**values) 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 3bcc7522a3..95eeea626b 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 @@ -24,14 +24,14 @@ "description": "all values provided", "path": "/test", "websocket_path": "/ws", - "openai_api_path": "/openai" + "openai_api_path": "/openai", } ENDPOINT_BASE_REQUIRED_VALUES = {"method": "POST", "description": "only required values"} -ENDPOINT_ALL_VALUES = ENDPOINT_BASE_ALL_VALUES | {'function_name': 'apples'} +ENDPOINT_ALL_VALUES = ENDPOINT_BASE_ALL_VALUES | {"function_name": "apples"} -ENDPOINT_REQUIRED_VALUES = ENDPOINT_BASE_REQUIRED_VALUES | {'function_name': 'oranges'} +ENDPOINT_REQUIRED_VALUES = ENDPOINT_BASE_REQUIRED_VALUES | {"function_name": "oranges"} CORS_ALL_VALUES = { "allow_origins": ["http://example.com", "https://example.com"], @@ -40,7 +40,7 @@ "allow_headers": ["Content-Type"], "allow_credentials": True, "expose_headers": ["X-Custom-Header"], - "max_age": 3600 + "max_age": 3600, } CORS_REQUIRED_VALUES = {} @@ -71,12 +71,12 @@ def _test_model_instantiation(model_class, model_kwargs): """ model = model_class(**model_kwargs) assert model.model_fields_set == model_kwargs.keys() - for (key, expected_value) in model_kwargs.items(): + for key, expected_value in model_kwargs.items(): actual_value = getattr(model, key) if isinstance(actual_value, BaseModel) and isinstance(expected_value, dict): _test_model_instantiation(actual_value.__class__, expected_value) elif isinstance(actual_value, list) and isinstance(expected_value, list): - for (i, v) in enumerate(actual_value): + for i, v in enumerate(actual_value): if isinstance(v, BaseModel) and isinstance(expected_value[i], dict): _test_model_instantiation(v.__class__, expected_value[i]) else: @@ -87,8 +87,11 @@ def _test_model_instantiation(model_class, model_kwargs): return model -@pytest.mark.parametrize("endpoint_kwargs", [ENDPOINT_BASE_ALL_VALUES.copy(), ENDPOINT_BASE_REQUIRED_VALUES.copy()], - ids=["all-values", "required-values"]) +@pytest.mark.parametrize( + "endpoint_kwargs", + [ENDPOINT_BASE_ALL_VALUES.copy(), ENDPOINT_BASE_REQUIRED_VALUES.copy()], + ids=["all-values", "required-values"], +) def test_endpoint_base(endpoint_kwargs: dict): _test_model_instantiation(FastApiFrontEndConfig.EndpointBase, endpoint_kwargs) @@ -98,8 +101,11 @@ def test_endpoint_base_invalid_method(): FastApiFrontEndConfig.EndpointBase(method="INVALID", description="test") -@pytest.mark.parametrize("endpoint_kwargs", [ENDPOINT_ALL_VALUES.copy(), ENDPOINT_REQUIRED_VALUES.copy()], - ids=["all-values", "required-values"]) +@pytest.mark.parametrize( + "endpoint_kwargs", + [ENDPOINT_ALL_VALUES.copy(), ENDPOINT_REQUIRED_VALUES.copy()], + ids=["all-values", "required-values"], +) def test_endpoint(endpoint_kwargs: dict): _test_model_instantiation(FastApiFrontEndConfig.Endpoint, endpoint_kwargs) @@ -119,8 +125,10 @@ def test_cross_origin_resource_sharing(cors_kwargs: dict): @pytest.mark.parametrize( - "config_kwargs", [FAST_API_FRONT_END_CONFIG_ALL_VALUES.copy(), FAST_API_FRONT_END_CONFIG_REQUIRES_VALUES.copy()], - ids=["all-values", "required-values"]) + "config_kwargs", + [FAST_API_FRONT_END_CONFIG_ALL_VALUES.copy(), FAST_API_FRONT_END_CONFIG_REQUIRES_VALUES.copy()], + ids=["all-values", "required-values"], +) def test_fast_api_front_end_config(config_kwargs: dict): model = _test_model_instantiation(FastApiFrontEndConfig, config_kwargs) @@ -139,5 +147,51 @@ def test_fast_api_front_end_config(config_kwargs: dict): assert isinstance(model.endpoints, list) assert isinstance(model.cors, FastApiFrontEndConfig.CrossOriginResourceSharing) assert isinstance(model.use_gunicorn, bool) - assert (isinstance(model.runner_class, str) or model.runner_class is None) - assert (isinstance(model.object_store, str) or model.object_store is None) + assert isinstance(model.runner_class, str) or model.runner_class is None + assert isinstance(model.object_store, str) or model.object_store is None + + +def test_accepted_identity_credentials_default_preserves_all_methods(): + config = FastApiFrontEndConfig() + + assert config.accepted_identity_credentials is None + + +def test_accepted_identity_credentials_accepts_supported_method_names(): + config = FastApiFrontEndConfig(accepted_identity_credentials=["session_cookie", "jwt", "api_key", "basic"]) + + assert config.accepted_identity_credentials == ["session_cookie", "jwt", "api_key", "basic"] + + +def test_accepted_identity_credentials_accepts_empty_list(): + config = FastApiFrontEndConfig(accepted_identity_credentials=[]) + + assert config.accepted_identity_credentials == [] + + +def test_accepted_identity_credentials_rejects_unknown_method(): + with pytest.raises(ValueError, match="accepted_identity_credentials"): + FastApiFrontEndConfig(accepted_identity_credentials=["unknown"]) + + +def test_identity_authentication_defaults_to_decode_only_behavior(): + config = FastApiFrontEndConfig() + + assert config.identity_authentication == [] + + +def test_identity_authentication_accepts_named_providers(): + config = FastApiFrontEndConfig( + accepted_identity_credentials=["jwt"], + identity_authentication=["corporate_jwt", "partner_jwt"], + ) + + assert config.identity_authentication == ["corporate_jwt", "partner_jwt"] + + +def test_identity_authentication_rejects_policy_that_disables_jwt(): + with pytest.raises(ValueError, match="identity_authentication cannot be configured when jwt identity credentials"): + FastApiFrontEndConfig( + accepted_identity_credentials=["session_cookie"], + 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 bee8546794..9709a88e21 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 @@ -14,14 +14,18 @@ # limitations under the License. import asyncio +import base64 +import json from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import patch +from starlette.websockets import WebSocket from starlette.websockets import WebSocketDisconnect from nat.data_models.api_server import ApiKeyAuthPayload from nat.data_models.api_server import AuthMessageStatus +from nat.data_models.api_server import JwtAuthPayload from nat.data_models.api_server import OAuthMode from nat.data_models.api_server import OAuthModePreferencePayload from nat.data_models.api_server import WebSocketAuthMessage @@ -31,9 +35,19 @@ from nat.front_ends.fastapi.message_handler import WebSocketMessageHandler -def _make_message_handler() -> tuple[WebSocketMessageHandler, AsyncMock, WebSocketAuthenticationFlowHandler]: +def _make_jwt(claims: dict) -> str: + """Build a minimal unsigned JWT for exercising the verifier boundary.""" + header = base64.urlsafe_b64encode(json.dumps({"alg": "none", "typ": "JWT"}).encode()).rstrip(b"=").decode() + payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=").decode() + return f"{header}.{payload}." + + +def _make_message_handler( + accepted_identity_credentials=None, + jwt_validators=None, +) -> tuple[WebSocketMessageHandler, MagicMock, WebSocketAuthenticationFlowHandler]: """Build a WebSocketMessageHandler with a mockable socket and a real flow handler.""" - socket = AsyncMock() + socket = MagicMock(spec=WebSocket) session_manager = MagicMock() session_manager.get_workflow_single_output_schema.return_value = None session_manager.get_workflow_streaming_output_schema.return_value = None @@ -42,6 +56,8 @@ def _make_message_handler() -> tuple[WebSocketMessageHandler, AsyncMock, WebSock session_manager=session_manager, step_adaptor=MagicMock(), worker=MagicMock(), + accepted_identity_credentials=accepted_identity_credentials, + jwt_validators=jwt_validators, ) flow_handler = WebSocketAuthenticationFlowHandler( add_flow_cb=AsyncMock(), @@ -60,8 +76,10 @@ async def test_context_manager_resolves_connection_identity_before_restoration() restore = AsyncMock() handler._restore_execution_state = restore - with patch("nat.front_ends.fastapi.message_handler.UserManager.extract_user_from_connection", - return_value=user_info): + with patch( + "nat.front_ends.fastapi.message_handler.UserManager.extract_user_from_connection_with_verification", + return_value=user_info, + ): await handler.__aenter__() socket.accept.assert_awaited_once() @@ -69,6 +87,66 @@ async def test_context_manager_resolves_connection_identity_before_restoration() restore.assert_awaited_once() +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"]) + restore = AsyncMock() + handler._restore_execution_state = restore + + from nat.runtime.user_manager import IdentityCredentialNotAcceptedError + + with patch( + "nat.front_ends.fastapi.message_handler.UserManager.extract_user_from_connection_with_verification", + side_effect=IdentityCredentialNotAcceptedError("Identity credential type 'session_cookie' is not accepted"), + ): + await handler.__aenter__() + + restore.assert_not_awaited() + response = socket.send_json.await_args.args[0] + assert response["status"] == AuthMessageStatus.ERROR + socket.close.assert_awaited_once_with(code=1008, reason="Identity credential was rejected") + + +async def test_context_manager_accepts_verified_connection_jwt_before_restoration(): + """An active verifier result permits identity resolution and owned-state restoration.""" + issuer = "https://identity.example.com" + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=True)) + handler, socket, _ = _make_message_handler(jwt_validators={issuer: jwt_validator}) + token = _make_jwt({"iss": issuer, "sub": "verified-user"}) + socket.scope = {"headers": [(b"authorization", f"Bearer {token}".encode())]} + socket.query_params = {"conversation_id": "conversation-a"} + restore = AsyncMock() + handler._restore_execution_state = restore + + await handler.__aenter__() + + jwt_validator.verify.assert_awaited_once_with(token) + assert handler._user_id is not None + restore.assert_awaited_once() + + +async def test_context_manager_rejects_unverified_connection_jwt_without_restoration(): + """An inactive verifier result cannot establish identity or restore conversation state.""" + issuer = "https://identity.example.com" + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=False)) + handler, socket, _ = _make_message_handler(jwt_validators={issuer: jwt_validator}) + token = _make_jwt({"iss": issuer, "sub": "unverified-user"}) + socket.scope = {"headers": [(b"authorization", f"Bearer {token}".encode())]} + socket.query_params = {"conversation_id": "conversation-a"} + restore = AsyncMock() + handler._restore_execution_state = restore + + await handler.__aenter__() + + assert handler._user_id is None + restore.assert_not_awaited() + response = socket.send_json.await_args.args[0] + assert response["status"] == AuthMessageStatus.ERROR + socket.close.assert_awaited_once_with(code=1008, reason="Identity credential was rejected") + + async def test_anonymous_connection_does_not_attempt_conversation_lookup(): """A conversation ID cannot restore state without a resolved user identity.""" handler, socket, _ = _make_message_handler() @@ -91,7 +169,10 @@ async def test_successful_auth_message_attempts_owned_restoration_once(): payload=ApiKeyAuthPayload(method="api_key", token="test-api-key"), ) - with patch("nat.front_ends.fastapi.message_handler.UserManager._from_auth_payload", return_value=user_info): + with patch( + "nat.front_ends.fastapi.message_handler.UserManager.from_auth_payload_with_verification", + return_value=user_info, + ): await handler._process_auth_message(msg) await handler._process_auth_message(msg) @@ -108,8 +189,10 @@ async def test_failed_auth_message_does_not_attempt_restoration(): payload=ApiKeyAuthPayload(method="api_key", token="test-api-key"), ) - with patch("nat.front_ends.fastapi.message_handler.UserManager._from_auth_payload", - side_effect=ValueError("invalid credential")): + with patch( + "nat.front_ends.fastapi.message_handler.UserManager.from_auth_payload_with_verification", + side_effect=ValueError("invalid credential"), + ): await handler._process_auth_message(msg) restore.assert_not_awaited() @@ -117,6 +200,51 @@ async def test_failed_auth_message_does_not_attempt_restoration(): assert response["status"] == AuthMessageStatus.ERROR +async def test_disabled_auth_message_preserves_existing_identity_and_does_not_restore(): + """A disabled auth message cannot replace identity or trigger restoration.""" + handler, socket, _ = _make_message_handler(accepted_identity_credentials=["jwt"]) + handler._user_id = "existing-user" + restore = AsyncMock() + handler._restore_execution_state = restore + msg = WebSocketAuthMessage( + type=WebSocketMessageType.AUTH_MESSAGE, + payload=ApiKeyAuthPayload(method="api_key", token="test-api-key"), + ) + + await handler._process_auth_message(msg) + + assert handler._user_id == "existing-user" + restore.assert_not_awaited() + response = socket.send_json.await_args.args[0] + assert response["status"] == AuthMessageStatus.ERROR + assert response["payload"]["details"] == "Identity credential type 'api_key' is not accepted" + + +async def test_rejected_auth_message_jwt_preserves_identity_and_does_not_restore(): + """A rejected JWT auth message cannot replace identity or restore conversation state.""" + issuer = "https://identity.example.com" + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=False)) + handler, socket, _ = _make_message_handler(jwt_validators={issuer: jwt_validator}) + handler._user_id = "existing-user" + socket.query_params = {"conversation_id": "conversation-a"} + restore = AsyncMock() + handler._restore_execution_state = restore + token = _make_jwt({"iss": issuer, "sub": "unverified-user"}) + msg = WebSocketAuthMessage( + type=WebSocketMessageType.AUTH_MESSAGE, + payload=JwtAuthPayload(method="jwt", token=token), + ) + + await handler._process_auth_message(msg) + + assert handler._user_id == "existing-user" + restore.assert_not_awaited() + response = socket.send_json.await_args.args[0] + assert response["status"] == AuthMessageStatus.ERROR + assert response["payload"]["details"] == "JWT verification failed" + + async def test_process_auth_message_sets_oauth_mode_and_sends_no_response(): """An oauth_mode_preference payload updates the flow handler's mode and emits no auth response.""" handler, socket, flow_handler = _make_message_handler() diff --git a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_jwt_providers.py b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_jwt_providers.py new file mode 100644 index 0000000000..d82f74b713 --- /dev/null +++ b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_jwt_providers.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +import pytest + +from nat.authentication.jwt.jwt_auth_provider import JwtAuthProvider +from nat.authentication.jwt.jwt_auth_provider_config import JwtAuthProviderConfig +from nat.front_ends.fastapi.routes.websocket import _get_jwt_validators + + +def _provider(issuer: str) -> JwtAuthProvider: + return JwtAuthProvider( + JwtAuthProviderConfig( + issuer_url=issuer, + jwks_uri=f"{issuer}/jwks.json", + audience="nat-api", + )) + + +async def test_get_jwt_validators_resolves_named_providers_by_issuer(): + corporate = _provider("https://corporate.example.com") + partner = _provider("https://partner.example.com") + worker = MagicMock() + worker.front_end_config.identity_authentication = ["corporate_jwt", "partner_jwt"] + session_manager = MagicMock() + session_manager.shared_builder.get_auth_provider = AsyncMock(side_effect=[corporate, partner]) + + validators = await _get_jwt_validators(worker, session_manager) + + assert validators == { + "https://corporate.example.com": corporate.validator, + "https://partner.example.com": partner.validator, + } + + +async def test_get_jwt_validators_rejects_non_jwt_provider(): + worker = MagicMock() + worker.front_end_config.identity_authentication = ["not_jwt"] + session_manager = MagicMock() + session_manager.shared_builder.get_auth_provider = AsyncMock(return_value=MagicMock()) + + with pytest.raises(ValueError, match="must have _type: jwt"): + await _get_jwt_validators(worker, session_manager) + + +async def test_get_jwt_validators_rejects_duplicate_issuers(): + first = _provider("https://identity.example.com") + second = _provider("https://identity.example.com") + worker = MagicMock() + worker.front_end_config.identity_authentication = ["first", "second"] + session_manager = MagicMock() + session_manager.shared_builder.get_auth_provider = AsyncMock(side_effect=[first, second]) + + with pytest.raises(ValueError, match="Duplicate JWT issuer"): + await _get_jwt_validators(worker, session_manager) diff --git a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_route_origin.py b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_route_origin.py index de357e0409..e29a65d83f 100644 --- a/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_route_origin.py +++ b/packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_route_origin.py @@ -13,9 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + import pytest +from starlette.websockets import WebSocket from nat.front_ends.fastapi.routes.websocket import _is_origin_allowed +from nat.front_ends.fastapi.routes.websocket import websocket_endpoint +from nat.runtime.session import SESSION_COOKIE_NAME @pytest.mark.parametrize( @@ -47,6 +53,31 @@ ("http://localhost:3000", [], r"http://localhost:\d+", True), # Both list and regex configured; list matches first ("http://localhost:3000", ["http://localhost:3000"], r"http://other\.com", True), - ]) + ], +) def test_is_origin_allowed(origin, allowed_origins, allow_origin_regex, expected): assert _is_origin_allowed(origin, allowed_origins, allow_origin_regex) is expected + + +async def test_legacy_session_query_uses_configured_identity_policy(): + websocket = MagicMock(spec=WebSocket) + websocket.query_params = {"session": "legacy-session"} + websocket.scope = {"headers": []} + websocket.headers.get.return_value = None + websocket.accept = AsyncMock() + websocket.send_json = AsyncMock() + websocket.close = AsyncMock() + + worker = MagicMock() + worker.front_end_config.accepted_identity_credentials = [] + worker.front_end_config.cors.allow_origins = [] + worker.front_end_config.cors.allow_origin_regex = None + session_manager = MagicMock() + + endpoint = websocket_endpoint(worker=worker, session_manager=session_manager, jwt_validators={}) + await endpoint(websocket) + + assert (b"cookie", f"{SESSION_COOKIE_NAME}=legacy-session".encode()) in websocket.scope["headers"] + websocket.send_json.assert_awaited_once() + websocket.close.assert_awaited_once_with(code=1008, reason="Identity credential was rejected") + worker.get_conversation_handler.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 66d02110f9..17c22e3c56 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 @@ -33,6 +33,8 @@ from nat.data_models.user_info import JwtUserInfo 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 JwtVerificationError from nat.runtime.user_manager import UserManager @@ -87,16 +89,16 @@ def test_deterministic_uuid_from_cookie(self): req1 = _mock_request(cookies={SESSION_COOKIE_NAME: "same-cookie"}) req2 = _mock_request(cookies={SESSION_COOKIE_NAME: "same-cookie"}) - assert UserManager.extract_user_from_connection(req1).get_user_id() == \ - UserManager.extract_user_from_connection(req2).get_user_id() + assert (UserManager.extract_user_from_connection(req1).get_user_id() == + UserManager.extract_user_from_connection(req2).get_user_id()) def test_different_cookies_different_uuids(self): """Input: two Requests with different cookie values. Asserts they produce different user_ids.""" req1 = _mock_request(cookies={SESSION_COOKIE_NAME: "cookie-a"}) req2 = _mock_request(cookies={SESSION_COOKIE_NAME: "cookie-b"}) - assert UserManager.extract_user_from_connection(req1).get_user_id() != \ - UserManager.extract_user_from_connection(req2).get_user_id() + assert (UserManager.extract_user_from_connection(req1).get_user_id() + != UserManager.extract_user_from_connection(req2).get_user_id()) class TestFromConnectionRequestJwt: @@ -223,6 +225,86 @@ def test_websocket_jwt_returns_user_info(self): assert isinstance(details, JwtUserInfo) assert details.email == "ws@example.com" + async def test_omitted_verifier_preserves_decode_only_behavior(self): + token = _make_jwt({"sub": "decode-only-user"}) + ws = _mock_websocket(auth_header=f"Bearer {token}") + + info = await UserManager.extract_user_from_connection_with_verification(ws, jwt_validators=None) + + assert info is not None + assert info.get_user_details().subject == "decode-only-user" + + async def test_active_verifier_accepts_websocket_jwt(self): + issuer = "https://identity.example.com" + token = _make_jwt({"iss": issuer, "sub": "verified-user"}) + ws = _mock_websocket(auth_header=f"Bearer {token}") + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=True)) + + info = await UserManager.extract_user_from_connection_with_verification( + ws, + jwt_validators={issuer: jwt_validator}, + ) + + jwt_validator.verify.assert_awaited_once_with(token) + assert info is not None + assert info.get_user_details().subject == "verified-user" + + async def test_inactive_verifier_rejects_websocket_jwt(self): + issuer = "https://identity.example.com" + token = _make_jwt({"iss": issuer, "sub": "unverified-user"}) + ws = _mock_websocket(auth_header=f"Bearer {token}") + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=False)) + + with pytest.raises(JwtVerificationError, match="JWT verification failed"): + await UserManager.extract_user_from_connection_with_verification( + ws, + jwt_validators={issuer: jwt_validator}, + ) + + async def test_unknown_issuer_rejects_websocket_jwt(self): + token = _make_jwt({"iss": "https://unknown.example.com", "sub": "unknown-user"}) + ws = _mock_websocket(auth_header=f"Bearer {token}") + + with pytest.raises(JwtVerificationError, match="JWT issuer is not accepted"): + await UserManager.extract_user_from_connection_with_verification( + ws, + jwt_validators={"https://identity.example.com": MagicMock()}, + ) + + async def test_missing_issuer_rejects_websocket_jwt_when_verification_is_configured(self): + token = _make_jwt({"sub": "missing-issuer-user"}) + ws = _mock_websocket(auth_header=f"Bearer {token}") + + with pytest.raises(JwtVerificationError, match="non-empty issuer claim"): + await UserManager.extract_user_from_connection_with_verification( + ws, + jwt_validators={"https://identity.example.com": MagicMock()}, + ) + + async def test_verified_jwt_identity_is_scoped_by_issuer(self): + first_issuer = "https://first.example.com" + second_issuer = "https://second.example.com" + validators = {} + for issuer in (first_issuer, second_issuer): + validator = MagicMock() + validator.verify = AsyncMock(return_value=MagicMock(active=True)) + validators[issuer] = validator + + first = await UserManager.extract_user_from_connection_with_verification( + _mock_websocket(auth_header=f"Bearer {_make_jwt({'iss': first_issuer, 'sub': 'shared-subject'})}"), + jwt_validators=validators, + ) + second = await UserManager.extract_user_from_connection_with_verification( + _mock_websocket(auth_header=f"Bearer {_make_jwt({'iss': second_issuer, 'sub': 'shared-subject'})}"), + jwt_validators=validators, + ) + + assert first is not None + assert second is not None + assert first.get_user_id() != second.get_user_id() + class TestFromConnectionPriority: """extract_user_from_connection prefers session cookie over JWT when both are present.""" @@ -247,6 +329,108 @@ def test_websocket_cookie_takes_precedence_over_jwt(self): info: UserInfo = UserManager.extract_user_from_connection(ws) assert info.get_user_details() == "ws-cookie-user" + def test_disabled_cookie_does_not_fall_through_to_allowed_jwt(self): + """A disabled higher-priority credential rejects the connection instead of changing identity.""" + token: str = _make_jwt({"sub": "jwt-user"}) + ws = _mock_websocket( + cookie_header=f"{SESSION_COOKIE_NAME}=ws-cookie-user", + auth_header=f"Bearer {token}", + ) + + with pytest.raises(IdentityCredentialNotAcceptedError, match="session_cookie"): + UserManager.extract_user_from_connection(ws, accepted_identity_credentials=["jwt"]) + + async def test_async_resolver_does_not_fall_through_from_disabled_cookie_to_verified_jwt(self): + issuer = "https://identity.example.com" + token = _make_jwt({"iss": issuer, "sub": "jwt-user"}) + ws = _mock_websocket( + cookie_header=f"{SESSION_COOKIE_NAME}=ws-cookie-user", + auth_header=f"Bearer {token}", + ) + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=True)) + + with pytest.raises(IdentityCredentialNotAcceptedError, match="session_cookie"): + await UserManager.extract_user_from_connection_with_verification( + ws, + accepted_identity_credentials=["jwt"], + jwt_validators={issuer: jwt_validator}, + ) + + jwt_validator.verify.assert_not_awaited() + + +class TestAcceptedIdentityCredentials: + """Configured credential methods limit WebSocket identity resolution.""" + + @pytest.mark.parametrize( + ("websocket", "accepted_method"), + [ + (_mock_websocket(cookie_header=f"{SESSION_COOKIE_NAME}=session"), "session_cookie"), + (_mock_websocket(auth_header=f"Bearer {_make_jwt({'sub': 'jwt-user'})}"), "jwt"), + (_mock_websocket(auth_header="Bearer api-key"), "api_key"), + (_mock_websocket(api_key_header="api-key"), "api_key"), + (_mock_websocket(auth_header="Basic dXNlcjpwYXNz"), "basic"), + ], + ids=["session-cookie", "jwt", "bearer-api-key", "x-api-key", "basic"], + ) + def test_enabled_connection_credential_is_accepted(self, websocket, accepted_method): + info = UserManager.extract_user_from_connection(websocket, accepted_identity_credentials=[accepted_method]) + + assert info is not None + + @pytest.mark.parametrize( + ("websocket", "disabled_method"), + [ + (_mock_websocket(cookie_header=f"{SESSION_COOKIE_NAME}=session"), "session_cookie"), + (_mock_websocket(auth_header=f"Bearer {_make_jwt({'sub': 'jwt-user'})}"), "jwt"), + (_mock_websocket(auth_header="Bearer api-key"), "api_key"), + (_mock_websocket(api_key_header="api-key"), "api_key"), + (_mock_websocket(auth_header="Basic dXNlcjpwYXNz"), "basic"), + ], + ids=["session-cookie", "jwt", "bearer-api-key", "x-api-key", "basic"], + ) + def test_disabled_connection_credential_is_rejected(self, websocket, disabled_method): + with pytest.raises(IdentityCredentialNotAcceptedError, match=disabled_method): + UserManager.extract_user_from_connection(websocket, accepted_identity_credentials=[]) + + @pytest.mark.parametrize( + ("websocket", "accepted_method"), + [ + (_mock_websocket(cookie_header=f"{SESSION_COOKIE_NAME}=session"), "session_cookie"), + (_mock_websocket(auth_header=f"Bearer {_make_jwt({'sub': 'jwt-user'})}"), "jwt"), + (_mock_websocket(auth_header="Bearer api-key"), "api_key"), + (_mock_websocket(api_key_header="api-key"), "api_key"), + (_mock_websocket(auth_header="Basic dXNlcjpwYXNz"), "basic"), + ], + ids=["session-cookie", "jwt", "bearer-api-key", "x-api-key", "basic"], + ) + async def test_enabled_connection_credential_is_accepted_by_async_resolver(self, websocket, accepted_method): + info = await UserManager.extract_user_from_connection_with_verification( + websocket, + accepted_identity_credentials=[accepted_method], + ) + + assert info is not None + + @pytest.mark.parametrize( + ("websocket", "disabled_method"), + [ + (_mock_websocket(cookie_header=f"{SESSION_COOKIE_NAME}=session"), "session_cookie"), + (_mock_websocket(auth_header=f"Bearer {_make_jwt({'sub': 'jwt-user'})}"), "jwt"), + (_mock_websocket(auth_header="Bearer api-key"), "api_key"), + (_mock_websocket(api_key_header="api-key"), "api_key"), + (_mock_websocket(auth_header="Basic dXNlcjpwYXNz"), "basic"), + ], + ids=["session-cookie", "jwt", "bearer-api-key", "x-api-key", "basic"], + ) + async def test_disabled_connection_credential_is_rejected_by_async_resolver(self, websocket, disabled_method): + with pytest.raises(IdentityCredentialNotAcceptedError, match=disabled_method): + await UserManager.extract_user_from_connection_with_verification( + websocket, + accepted_identity_credentials=[], + ) + class TestFromConnectionNoCredential: """extract_user_from_connection with missing or invalid credentials.""" @@ -296,8 +480,7 @@ def test_jwt_payload_deterministic_uuid(self): p1 = JwtAuthPayload(method="jwt", token=SecretStr(token)) p2 = JwtAuthPayload(method="jwt", token=SecretStr(token)) - assert UserManager._from_auth_payload(p1).get_user_id() == \ - UserManager._from_auth_payload(p2).get_user_id() + assert UserManager._from_auth_payload(p1).get_user_id() == UserManager._from_auth_payload(p2).get_user_id() def test_jwt_payload_invalid_token_raises(self): """Input: JWT payload with non-JWT string. Asserts raises ValueError matching "malformed".""" @@ -317,6 +500,38 @@ def test_jwt_payload_no_identity_claim_raises(self): with pytest.raises(ValueError, match="no usable identity claim"): UserManager._from_auth_payload(payload) + async def test_active_verifier_accepts_jwt_payload(self): + issuer = "https://identity.example.com" + token = _make_jwt({"iss": issuer, "sub": "verified-payload-user"}) + payload = JwtAuthPayload(method="jwt", token=SecretStr(token)) + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=True)) + + info = await UserManager.from_auth_payload_with_verification(payload, jwt_validators={issuer: jwt_validator}) + + jwt_validator.verify.assert_awaited_once_with(token) + assert info.get_user_details().subject == "verified-payload-user" + + async def test_inactive_verifier_rejects_jwt_payload(self): + issuer = "https://identity.example.com" + token = _make_jwt({"iss": issuer, "sub": "unverified-payload-user"}) + payload = JwtAuthPayload(method="jwt", token=SecretStr(token)) + jwt_validator = MagicMock() + jwt_validator.verify = AsyncMock(return_value=MagicMock(active=False)) + + with pytest.raises(JwtVerificationError, match="JWT verification failed"): + await UserManager.from_auth_payload_with_verification(payload, jwt_validators={issuer: jwt_validator}) + + async def test_missing_issuer_rejects_jwt_payload_when_verification_is_configured(self): + token = _make_jwt({"sub": "missing-issuer-user"}) + payload = JwtAuthPayload(method="jwt", token=SecretStr(token)) + + with pytest.raises(JwtVerificationError, match="non-empty issuer claim"): + await UserManager.from_auth_payload_with_verification( + payload, + jwt_validators={"https://identity.example.com": MagicMock()}, + ) + class TestFromAuthPayloadApiKey: """_from_auth_payload resolves UserInfo from an ApiKeyAuthPayload.""" @@ -334,8 +549,7 @@ def test_api_key_deterministic_uuid(self): p1 = ApiKeyAuthPayload(method="api_key", token=SecretStr("same-key")) p2 = ApiKeyAuthPayload(method="api_key", token=SecretStr("same-key")) - assert UserManager._from_auth_payload(p1).get_user_id() == \ - UserManager._from_auth_payload(p2).get_user_id() + assert UserManager._from_auth_payload(p1).get_user_id() == UserManager._from_auth_payload(p2).get_user_id() def test_api_key_empty_token_raises(self): """Input: API key payload with empty token. Asserts raises ValidationError (min_length=1).""" @@ -361,16 +575,28 @@ def test_basic_payload_deterministic_uuid(self): p1 = BasicAuthPayload(method="basic", username="bob", password=SecretStr("pass")) p2 = BasicAuthPayload(method="basic", username="bob", password=SecretStr("pass")) - assert UserManager._from_auth_payload(p1).get_user_id() == \ - UserManager._from_auth_payload(p2).get_user_id() + assert UserManager._from_auth_payload(p1).get_user_id() == UserManager._from_auth_payload(p2).get_user_id() def test_basic_different_users_different_uuids(self): """Input: two different basic payloads. Asserts they produce different user_ids.""" p1 = BasicAuthPayload(method="basic", username="alice", password=SecretStr("pass")) p2 = BasicAuthPayload(method="basic", username="bob", password=SecretStr("pass")) - assert UserManager._from_auth_payload(p1).get_user_id() != \ - UserManager._from_auth_payload(p2).get_user_id() + assert UserManager._from_auth_payload(p1).get_user_id() != UserManager._from_auth_payload(p2).get_user_id() + + +@pytest.mark.parametrize( + ("payload", "disabled_method"), + [ + (JwtAuthPayload(method="jwt", token=SecretStr(_make_jwt({"sub": "user"}))), "jwt"), + (ApiKeyAuthPayload(method="api_key", token=SecretStr("api-key")), "api_key"), + (BasicAuthPayload(method="basic", username="user", password=SecretStr("password")), "basic"), + ], + ids=["jwt", "api-key", "basic"], +) +async def test_disabled_auth_payload_is_rejected(payload, disabled_method): + with pytest.raises(IdentityCredentialNotAcceptedError, match=disabled_method): + await UserManager.from_auth_payload_with_verification(payload, accepted_identity_credentials=[]) class TestHandlerProcessAuthMessage: @@ -378,6 +604,7 @@ class TestHandlerProcessAuthMessage: def _make_handler(self): from nat.front_ends.fastapi.message_handler import WebSocketMessageHandler + mock_socket = MagicMock(spec=WebSocket) mock_socket.send_json = AsyncMock() handler = WebSocketMessageHandler( @@ -396,6 +623,7 @@ def _last_sent_payload(self, handler) -> dict: async def test_jwt_auth_message_sets_user_id(self): """Input: valid JWT auth message. Asserts handler._user_id is set and success response sent.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() token: str = _make_jwt({"sub": "ws-auth-user", "email": "ws@auth.io"}) msg = WebSocketAuthMessage( @@ -417,6 +645,7 @@ async def test_jwt_auth_message_sets_user_id(self): async def test_api_key_auth_message_sets_user_id(self): """Input: API key auth message. Asserts handler._user_id is set and success response sent.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -432,6 +661,7 @@ async def test_api_key_auth_message_sets_user_id(self): async def test_basic_auth_message_sets_user_id(self): """Input: basic auth message. Asserts handler._user_id is set and success response sent.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -446,6 +676,7 @@ async def test_basic_auth_message_sets_user_id(self): async def test_invalid_jwt_leaves_user_id_none_and_sends_failure(self): """Input: malformed JWT auth message. Asserts user_id stays None and error response sent.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -464,6 +695,7 @@ async def test_invalid_jwt_leaves_user_id_none_and_sends_failure(self): async def test_api_key_auth_success_response_contains_user_id(self): """Input: API key auth message. Asserts response user_id matches handler._user_id.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -477,6 +709,7 @@ async def test_api_key_auth_success_response_contains_user_id(self): async def test_basic_auth_success_response_contains_user_id(self): """Input: basic auth message. Asserts response user_id matches handler._user_id.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -490,6 +723,7 @@ async def test_basic_auth_success_response_contains_user_id(self): async def test_auth_message_user_id_matches_direct_resolution(self): """The handler-stored user_id must match a direct _from_auth_payload call.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() token: str = _make_jwt({"sub": "consistency-check", "email": "c@c.io"}) payload = JwtAuthPayload(method="jwt", token=SecretStr(token)) @@ -524,6 +758,7 @@ async def test_user_id_forwarded_to_session(self): async def test_success_response_payload_is_none(self): """Input: valid JWT auth message. Asserts success response payload is None (no error).""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() token: str = _make_jwt({"sub": "u", "email": "a@b.com"}) msg = WebSocketAuthMessage( @@ -538,6 +773,7 @@ async def test_success_response_payload_is_none(self): async def test_error_response_user_id_is_none(self): """Input: malformed JWT auth message. Asserts error response user_id is None.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -551,6 +787,7 @@ async def test_error_response_user_id_is_none(self): async def test_error_response_has_details(self): """Input: malformed JWT auth message. Asserts error response contains non-empty details string.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() msg = WebSocketAuthMessage( type="auth_message", @@ -565,6 +802,7 @@ async def test_error_response_has_details(self): async def test_second_auth_message_overrides_user_id(self): """Input: two auth messages for different users. Asserts second overrides first user_id.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() token_a: str = _make_jwt({"sub": "user-a", "email": "user-a@x.com"}) @@ -592,6 +830,7 @@ async def test_second_auth_message_overrides_user_id(self): async def test_auth_then_workflow_passes_user_id(self): """Input: auth message then _run_workflow. Asserts session is called with the resolved user_id.""" from nat.data_models.api_server import WebSocketAuthMessage + handler = self._make_handler() token: str = _make_jwt({"sub": "flow-user", "email": "flow@x.com"}) msg = WebSocketAuthMessage( @@ -848,6 +1087,7 @@ class TestContextVarPropagation: def test_context_var_set_and_readable(self): """Input: set user_id to "test-user". Asserts get() returns "test-user".""" from nat.builder.context import ContextState + state: ContextState = ContextState.get() token = state.user_id.set("test-user") try: @@ -858,6 +1098,7 @@ def test_context_var_set_and_readable(self): def test_context_var_reset_restores_previous(self): """Input: set "user-a", then "user-b", then reset. Asserts get() returns "user-a" after reset.""" from nat.builder.context import ContextState + state: ContextState = ContextState.get() token_a = state.user_id.set("user-a") try: @@ -1021,8 +1262,8 @@ def test_basic_auth_deterministic_uuid(self): req1 = _mock_request(headers={"authorization": f"Basic {b64}"}) req2 = _mock_request(headers={"authorization": f"Basic {b64}"}) - assert UserManager.extract_user_from_connection(req1).get_user_id() == \ - UserManager.extract_user_from_connection(req2).get_user_id() + assert (UserManager.extract_user_from_connection(req1).get_user_id() == + UserManager.extract_user_from_connection(req2).get_user_id()) def test_basic_auth_matches_direct_construction(self): """Input: Basic auth via header matches UserInfo(basic_user=...) with same creds.""" @@ -1082,8 +1323,8 @@ def test_api_key_deterministic_uuid(self): req1 = _mock_request(headers={"authorization": "Bearer sk-key-xyz"}) req2 = _mock_request(headers={"authorization": "Bearer sk-key-xyz"}) - assert UserManager.extract_user_from_connection(req1).get_user_id() == \ - UserManager.extract_user_from_connection(req2).get_user_id() + assert (UserManager.extract_user_from_connection(req1).get_user_id() == + UserManager.extract_user_from_connection(req2).get_user_id()) def test_api_key_matches_from_api_key_factory(self): """Input: API key via Bearer header matches UserInfo._from_api_key with same key.""" @@ -1131,8 +1372,8 @@ def test_x_api_key_deterministic_uuid(self): req1 = _mock_request(headers={"x-api-key": "nvapi-stable"}) req2 = _mock_request(headers={"x-api-key": "nvapi-stable"}) - assert UserManager.extract_user_from_connection(req1).get_user_id() == \ - UserManager.extract_user_from_connection(req2).get_user_id() + assert (UserManager.extract_user_from_connection(req1).get_user_id() == + UserManager.extract_user_from_connection(req2).get_user_id()) def test_x_api_key_matches_bearer_api_key(self): """Input: Same key via X-API-Key and Bearer. Asserts same user_id."""