nvidia-nat 1.3.0a20250926__py3-none-any.whl → 1.3.0a20250928__py3-none-any.whl

nat/authentication/api_key/api_key_auth_provider.py

@@ -80,7 +80,7 @@ class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
 
         raise ValueError(f"Unsupported header auth scheme: {header_auth_scheme}")
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult | None:
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult | None:
         """
         Authenticate the user using the API key credentials.
 

nat/authentication/http_basic_auth/http_basic_auth_provider.py

@@ -38,7 +38,7 @@ class HTTPBasicAuthProvider(AuthProviderBase):
 
         self._authenticated_tokens: dict[str, AuthResult] = {}
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
         """
         Performs simple HTTP Authentication using the provided user ID.
         """

nat/authentication/interfaces.py

@@ -54,7 +54,7 @@ class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
         return self._config
 
     @abstractmethod
-    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
         """
         Perform the authentication process for the client.
 
@@ -62,6 +62,9 @@ class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
         target API service, which may include obtaining tokens, refreshing credentials,
         or completing multi-step authentication flows.
 
+        Args:
+            user_id: Optional user identifier for authentication
+            kwargs: Additional authentication parameters for example: http response (typically from a 401)
         Raises:
             NotImplementedError: Must be implemented by subclasses.
         """
@@ -71,7 +74,7 @@ class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
 
 class FlowHandlerBase(ABC):
     """
-    Handles front-end specifc flows for authentication clients.
+    Handles front-end specific flows for authentication clients.
 
     Each front end will define a FlowHandler that will implement the authenticate method.
 

nat/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -15,6 +15,7 @@
 
 from datetime import datetime
 from datetime import timezone
+from typing import Callable
 
 from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
 from pydantic import SecretStr
@@ -22,6 +23,7 @@ from pydantic import SecretStr
 from nat.authentication.interfaces import AuthProviderBase
 from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
 from nat.builder.context import Context
+from nat.data_models.authentication import AuthenticatedContext
 from nat.data_models.authentication import AuthFlowType
 from nat.data_models.authentication import AuthResult
 from nat.data_models.authentication import BearerTokenCred
@@ -32,7 +34,7 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
     def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
         super().__init__(config)
         self._authenticated_tokens: dict[str, AuthResult] = {}
-        self._context = Context.get()
+        self._auth_callback = None
 
     async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
         refresh_token = auth_result.raw.get("refresh_token")
@@ -62,7 +64,12 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
 
         return new_auth_result
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+    def _set_custom_auth_callback(self,
+                                  auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
+                                                          AuthenticatedContext]):
+        self._auth_callback = auth_callback
+
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
         if user_id is None and hasattr(Context.get(), "metadata") and hasattr(
                 Context.get().metadata, "cookies") and Context.get().metadata.cookies is not None:
             session_id = Context.get().metadata.cookies.get("nat-session", None)
@@ -80,7 +87,12 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
             if refreshed_auth_result:
                 return refreshed_auth_result
 
-        auth_callback = self._context.user_auth_callback
+        # Try getting callback from the context if that's not set, use the default callback
+        try:
+            auth_callback = Context.get().user_auth_callback
+        except RuntimeError:
+            auth_callback = self._auth_callback
+
         if not auth_callback:
             raise RuntimeError("Authentication callback not set on Context.")
 

nat/cli/commands/mcp/mcp.py

@@ -157,6 +157,66 @@ def print_tool(tool_dict: dict[str, str | None], detail: bool = False) -> None:
         click.echo("-" * 60)
 
 
+def _set_auth_defaults(auth: bool,
+                       url: str | None,
+                       auth_redirect_uri: str | None,
+                       auth_user_id: str | None,
+                       auth_scopes: str | None) -> tuple[str | None, str | None, list[str] | None]:
+    """Set default auth values when --auth flag is used.
+
+    Args:
+        auth: Whether --auth flag was used
+        url: MCP server URL
+        auth_redirect_uri: OAuth2 redirect URI
+        auth_user_id: User ID for authentication
+        auth_scopes: OAuth2 scopes (comma-separated string)
+
+    Returns:
+        Tuple of (auth_redirect_uri, auth_user_id, auth_scopes_list) with defaults applied
+    """
+    if auth:
+        auth_redirect_uri = auth_redirect_uri or "http://localhost:8000/auth/redirect"
+        auth_user_id = auth_user_id or url
+        auth_scopes = auth_scopes or ""
+
+    # Convert comma-separated string to list, stripping whitespace
+    auth_scopes_list = [scope.strip() for scope in auth_scopes.split(',')] if auth_scopes else None
+
+    return auth_redirect_uri, auth_user_id, auth_scopes_list
+
+
+async def _create_mcp_client_config(
+    builder,
+    server_cfg,
+    url: str | None,
+    transport: str,
+    auth_redirect_uri: str | None,
+    auth_user_id: str | None,
+    auth_scopes: list[str] | None,
+):
+    from nat.plugins.mcp.client_impl import MCPClientConfig
+
+    if url and transport == "streamable-http" and auth_redirect_uri:
+        try:
+            from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
+            auth_config = MCPOAuth2ProviderConfig(
+                server_url=url,
+                redirect_uri=auth_redirect_uri,
+                default_user_id=auth_user_id or url,
+                scopes=auth_scopes or [],
+            )
+            auth_provider_name = "mcp_oauth2_cli"
+            await builder.add_auth_provider(auth_provider_name, auth_config)
+            server_cfg.auth_provider = auth_provider_name
+        except ImportError:
+            click.echo(
+                "[WARNING] MCP OAuth2 authentication requires nvidia-nat-mcp package.",
+                err=True,
+            )
+
+    return MCPClientConfig(server=server_cfg)
+
+
 async def list_tools_via_function_group(
     command: str | None,
     url: str | None,
@@ -164,6 +224,9 @@ async def list_tools_via_function_group(
     transport: str = 'sse',
     args: list[str] | None = None,
     env: dict[str, str] | None = None,
+    auth_redirect_uri: str | None = None,
+    auth_user_id: str | None = None,
+    auth_scopes: list[str] | None = None,
 ) -> list[dict[str, str | None]]:
     """List tools by constructing the mcp_client function group and introspecting functions.
 
@@ -192,15 +255,24 @@ async def list_tools_via_function_group(
         args=args if transport == 'stdio' else None,
         env=env if transport == 'stdio' else None,
     )
+
     group_cfg = MCPClientConfig(server=server_cfg)
 
     tools: list[dict[str, str | None]] = []
 
     async with WorkflowBuilder() as builder:  # type: ignore
+        # Add auth provider if url is provided and auth_redirect_uri is given (only for streamable-http)
+        group_cfg = await _create_mcp_client_config(builder,
+                                                    server_cfg,
+                                                    url,
+                                                    transport,
+                                                    auth_redirect_uri,
+                                                    auth_user_id,
+                                                    auth_scopes)
         group = await builder.add_function_group("mcp_client", group_cfg)
 
         # Access functions exposed by the group
-        fns = group.get_accessible_functions()
+        fns = await group.get_accessible_functions()
 
         def to_tool_entry(full_name: str, fn_obj) -> dict[str, str | None]:
             # full_name like "mcp_client.<tool>"
@@ -324,7 +396,10 @@ async def ping_mcp_server(url: str,
                           transport: str = 'streamable-http',
                           command: str | None = None,
                           args: list[str] | None = None,
-                          env: dict[str, str] | None = None) -> MCPPingResult:
+                          env: dict[str, str] | None = None,
+                          auth_redirect_uri: str | None = None,
+                          auth_user_id: str | None = None,
+                          auth_scopes: list[str] | None = None) -> MCPPingResult:
     """Ping an MCP server to check if it's responsive.
 
     Args:
@@ -383,7 +458,13 @@ def mcp_client_command():
     """
     MCP client commands.
     """
-    return None
+    try:
+        from nat.runtime.loader import PluginTypes
+        from nat.runtime.loader import discover_and_register_plugins
+        discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
+    except ImportError:
+        click.echo("[WARNING] MCP client functionality requires nvidia-nat-mcp package.", err=True)
+        pass
 
 
 @mcp_client_command.group(name="tool", invoke_without_command=False, help="Inspect and call MCP tools.")
@@ -412,14 +493,36 @@ def mcp_client_tool_group():
 @click.option('--tool', default=None, help='Get details for a specific tool by name')
 @click.option('--detail', is_flag=True, help='Show full details for all tools')
 @click.option('--json-output', is_flag=True, help='Output tool metadata in JSON format')
+@click.option('--auth',
+              is_flag=True,
+              help='Enable OAuth2 authentication with default settings (streamable-http only, not with --direct)')
+@click.option('--auth-redirect-uri',
+              help='OAuth2 redirect URI for authentication (streamable-http only, not with --direct)')
+@click.option('--auth-user-id', help='User ID for authentication (streamable-http only, not with --direct)')
+@click.option('--auth-scopes', help='OAuth2 scopes (comma-separated, streamable-http only, not with --direct)')
 @click.pass_context
-def mcp_client_tool_list(ctx, direct, url, transport, command, args, env, tool, detail, json_output):
+def mcp_client_tool_list(ctx,
+                         direct,
+                         url,
+                         transport,
+                         command,
+                         args,
+                         env,
+                         tool,
+                         detail,
+                         json_output,
+                         auth,
+                         auth_redirect_uri,
+                         auth_user_id,
+                         auth_scopes):
     """List MCP tool names (default) or show detailed tool information.
 
     Use --detail for full output including descriptions and input schemas.
     If --tool is provided, always shows full output for that specific tool.
     Use --direct to bypass MCPBuilder and use raw MCP protocol.
     Use --json-output to get structured JSON data instead of formatted text.
+    Use --auth to enable auth with default settings (streamable-http only, not with --direct).
+    Use --auth-redirect-uri to enable auth for protected MCP servers (streamable-http only, not with --direct).
 
     Args:
         ctx (click.Context): Click context object for command invocation
@@ -428,13 +531,22 @@ def mcp_client_tool_list(ctx, direct, url, transport, command, args, env, tool,
         tool (str | None): Optional specific tool name to retrieve detailed info for
         detail (bool): Whether to show full details (description + schema) for all tools
         json_output (bool): Whether to output tool metadata in JSON format instead of text
+        auth (bool): Whether to enable OAuth2 authentication (streamable-http only, not with --direct)
+        auth_redirect_uri (str | None): redirect URI for auth (streamable-http only, not with --direct)
+        auth_user_id (str | None): User ID for authentication (streamable-http only, not with --direct)
+        auth_scopes (str | None): OAuth2 scopes (comma-separated, streamable-http only, not with --direct)
 
     Examples:
         nat mcp client tool list                           # List tool names only
         nat mcp client tool list --detail                  # Show all tools with full details
         nat mcp client tool list --tool my_tool            # Show details for specific tool
         nat mcp client tool list --json-output             # Get JSON format output
-        nat mcp client tool list --direct --url http://... # Use direct protocol with custom URL
+        nat mcp client tool list --direct --url http://... # Use direct protocol with custom URL (no auth)
+        nat mcp client tool list --url https://example.com/mcp/ --auth # With auth using defaults
+        nat mcp client tool list --url https://example.com/mcp/ --transport streamable-http \
+            --auth-redirect-uri http://localhost:8000/auth/redirect # With custom auth settings
+        nat mcp client tool list --url https://example.com/mcp/ --transport streamable-http \
+            --auth-redirect-uri http://localhost:8000/auth/redirect --auth-user-id myuser # With auth and user ID
     """
     if ctx.invoked_subcommand is not None:
         return
@@ -447,11 +559,28 @@ def mcp_client_tool_list(ctx, direct, url, transport, command, args, env, tool,
             click.echo("[ERROR] --url is required when using sse or streamable-http client type", err=True)
             return
 
+    # Set auth defaults if --auth flag is used
+    auth_redirect_uri, auth_user_id, auth_scopes_list = _set_auth_defaults(
+        auth, url, auth_redirect_uri, auth_user_id, auth_scopes
+    )
+
     stdio_args = args.split() if args else []
     stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
 
-    fetcher = list_tools_direct if direct else list_tools_via_function_group
-    tools = asyncio.run(fetcher(command, url, tool, transport, stdio_args, stdio_env))
+    if direct:
+        tools = asyncio.run(
+            list_tools_direct(command, url, tool_name=tool, transport=transport, args=stdio_args, env=stdio_env))
+    else:
+        tools = asyncio.run(
+            list_tools_via_function_group(command,
+                                          url,
+                                          tool_name=tool,
+                                          transport=transport,
+                                          args=stdio_args,
+                                          env=stdio_env,
+                                          auth_redirect_uri=auth_redirect_uri,
+                                          auth_user_id=auth_user_id,
+                                          auth_scopes=auth_scopes_list))
 
     if json_output:
         click.echo(json.dumps(tools, indent=2))
@@ -482,13 +611,20 @@ def mcp_client_tool_list(ctx, direct, url, transport, command, args, env, tool,
 @click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
 @click.option('--timeout', default=60, show_default=True, help='Timeout in seconds for ping request')
 @click.option('--json-output', is_flag=True, help='Output ping result in JSON format')
+@click.option('--auth-redirect-uri',
+              help='OAuth2 redirect URI for authentication (streamable-http only, not with --direct)')
+@click.option('--auth-user-id', help='User ID for authentication (streamable-http only, not with --direct)')
+@click.option('--auth-scopes', help='OAuth2 scopes (comma-separated, streamable-http only, not with --direct)')
 def mcp_client_ping(url: str,
                     transport: str,
                     command: str | None,
                     args: str | None,
                     env: str | None,
                     timeout: int,
-                    json_output: bool) -> None:
+                    json_output: bool,
+                    auth_redirect_uri: str | None,
+                    auth_user_id: str | None,
+                    auth_scopes: str | None) -> None:
     """Ping an MCP server to check if it's responsive.
 
     This command sends a ping request to the MCP server and measures the response time.
@@ -498,12 +634,16 @@ def mcp_client_ping(url: str,
         url (str): MCP server URL to ping (default: http://localhost:9901/mcp)
         timeout (int): Timeout in seconds for the ping request (default: 60)
         json_output (bool): Whether to output the result in JSON format
+        auth_redirect_uri (str | None): redirect URI for auth (streamable-http only, not with --direct)
+        auth_user_id (str | None): User ID for auth (streamable-http only, not with --direct)
+        auth_scopes (str | None): OAuth2 scopes (comma-separated, streamable-http only, not with --direct)
 
     Examples:
         nat mcp client ping                                    # Ping default server
         nat mcp client ping --url http://custom-server:9901/mcp # Ping custom server
         nat mcp client ping --timeout 10                      # Use 10 second timeout
         nat mcp client ping --json-output                     # Get JSON format output
+        nat mcp client ping --url https://example.com/mcp/ --transport streamable-http --auth # With auth
     """
     # Validate combinations similar to list command
     if not validate_transport_cli_args(transport, command, args, env):
@@ -512,7 +652,24 @@ def mcp_client_ping(url: str,
     stdio_args = args.split() if args else []
     stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
 
-    result = asyncio.run(ping_mcp_server(url, timeout, transport, command, stdio_args, stdio_env))
+    # Auth validation: if user_id or scopes provided, require redirect_uri
+    if (auth_user_id or auth_scopes) and not auth_redirect_uri:
+        click.echo("[ERROR] --auth-redirect-uri is required when using --auth-user-id or --auth-scopes", err=True)
+        return
+
+    # Parse auth scopes, stripping whitespace
+    auth_scopes_list = [scope.strip() for scope in auth_scopes.split(',')] if auth_scopes else None
+
+    result = asyncio.run(
+        ping_mcp_server(url,
+                        timeout,
+                        transport,
+                        command,
+                        stdio_args,
+                        stdio_env,
+                        auth_redirect_uri,
+                        auth_user_id,
+                        auth_scopes_list))
 
     if json_output:
         click.echo(result.model_dump_json(indent=2))
@@ -635,7 +792,10 @@ async def call_tool_and_print(command: str | None,
                               args: list[str] | None,
                               env: dict[str, str] | None,
                               tool_args: dict[str, Any] | None,
-                              direct: bool) -> str:
+                              direct: bool,
+                              auth_redirect_uri: str | None = None,
+                              auth_user_id: str | None = None,
+                              auth_scopes: list[str] | None = None) -> str:
     """Call an MCP tool either directly or via the function group and return output.
 
     When ``direct`` is True, uses the raw MCP protocol client (bypassing the
@@ -681,11 +841,25 @@ async def call_tool_and_print(command: str | None,
         args=args if transport == 'stdio' else None,
         env=env if transport == 'stdio' else None,
     )
+
     group_cfg = MCPClientConfig(server=server_cfg)
 
     async with WorkflowBuilder() as builder:  # type: ignore
+        # Add auth provider if url is provided and auth_redirect_uri is given (only for streamable-http)
+        if url and transport == 'streamable-http' and auth_redirect_uri:
+            try:
+                group_cfg = await _create_mcp_client_config(builder,
+                                                            server_cfg,
+                                                            url,
+                                                            transport,
+                                                            auth_redirect_uri,
+                                                            auth_user_id,
+                                                            auth_scopes)
+            except ImportError:
+                click.echo("[WARNING] MCP OAuth2 authentication requires nvidia-nat-mcp package.", err=True)
+
         group = await builder.add_function_group("mcp_client", group_cfg)
-        fns = group.get_accessible_functions()
+        fns = await group.get_accessible_functions()
         full = f"mcp_client.{tool_name}"
         fn = fns.get(full)
         if fn is None:
@@ -713,6 +887,13 @@ async def call_tool_and_print(command: str | None,
 @click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
 @click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
 @click.option('--json-args', default=None, help='Pass tool args as a JSON object string')
+@click.option('--auth',
+              is_flag=True,
+              help='Enable OAuth2 authentication with default settings (streamable-http only, not with --direct)')
+@click.option('--auth-redirect-uri',
+              help='OAuth2 redirect URI for authentication (streamable-http only, not with --direct)')
+@click.option('--auth-user-id', help='User ID for authentication (streamable-http only, not with --direct)')
+@click.option('--auth-scopes', help='OAuth2 scopes (comma-separated, streamable-http only, not with --direct)')
 def mcp_client_tool_call(tool_name: str,
                          direct: bool,
                          url: str | None,
@@ -720,7 +901,11 @@ def mcp_client_tool_call(tool_name: str,
                          command: str | None,
                          args: str | None,
                          env: str | None,
-                         json_args: str | None) -> None:
+                         json_args: str | None,
+                         auth: bool,
+                         auth_redirect_uri: str | None,
+                         auth_user_id: str | None,
+                         auth_scopes: str | None) -> None:
     """Call an MCP tool by name with optional JSON arguments.
 
     Validates transport parameters, parses ``--json-args`` into a dictionary,
@@ -737,13 +922,20 @@ def mcp_client_tool_call(tool_name: str,
         args (str | None): For ``stdio`` transport, space-separated command arguments.
         env (str | None): For ``stdio`` transport, space-separated ``KEY=VALUE`` pairs.
         json_args (str | None): JSON object string with tool arguments (e.g. '{"q": "hello"}').
+        auth_redirect_uri (str | None): redirect URI for auth (streamable-http only, not with --direct)
+        auth_user_id (str | None): User ID for authentication (streamable-http only, not with --direct)
+        auth_scopes (str | None): OAuth2 scopes (comma-separated, streamable-http only, not with --direct)
 
     Examples:
         nat mcp client tool call echo --json-args '{"text": "Hello"}'
         nat mcp client tool call search --direct --url http://localhost:9901/mcp \
-            --json-args '{"query": "NVIDIA"}'
+            --json-args '{"query": "NVIDIA"}' # Direct mode (no auth)
         nat mcp client tool call run --transport stdio --command mcp-server \
             --args "--flag1 --flag2" --env "ENV1=V1 ENV2=V2" --json-args '{}'
+        nat mcp client tool call search --url https://example.com/mcp/ --auth \
+            --json-args '{"query": "test"}' # With auth using defaults
+        nat mcp client tool call search --url https://example.com/mcp/ \
+            --transport streamable-http --json-args '{"query": "test"}' --auth
     """
     # Validate transport args
     if not validate_transport_cli_args(transport, command, args, env):
@@ -753,6 +945,11 @@ def mcp_client_tool_call(tool_name: str,
     stdio_args = args.split() if args else []
     stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
 
+    # Set auth defaults if --auth flag is used
+    auth_redirect_uri, auth_user_id, auth_scopes_list = _set_auth_defaults(
+        auth, url, auth_redirect_uri, auth_user_id, auth_scopes
+    )
+
     # Parse tool args
     arg_obj: dict[str, Any] = {}
     if json_args:
@@ -777,6 +974,9 @@ def mcp_client_tool_call(tool_name: str,
                 env=stdio_env,
                 tool_args=arg_obj,
                 direct=direct,
+                auth_redirect_uri=auth_redirect_uri,
+                auth_user_id=auth_user_id,
+                auth_scopes=auth_scopes_list,
             ))
         if output:
             click.echo(output)

nat/data_models/authentication.py

@@ -249,21 +249,3 @@ class AuthResult(BaseModel):
                 target_kwargs.setdefault(k, {}).update(v)
             else:
                 target_kwargs[k] = v
-
-
-class AuthReason(str, Enum):
-    """
-    Why the caller is asking for auth now.
-    """
-    NORMAL = "normal"
-    RETRY_AFTER_401 = "retry_after_401"
-
-
-class AuthRequest(BaseModel):
-    """
-    Authentication request payload for provider.authenticate(...).
-    """
-    model_config = ConfigDict(extra="forbid")
-
-    reason: AuthReason = Field(default=AuthReason.NORMAL, description="Purpose of this auth attempt.")
-    www_authenticate: str | None = Field(default=None, description="Raw WWW-Authenticate header from a 401 response.")

nat/runtime/session.py

@@ -111,7 +111,7 @@ class SessionManager:
             token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
 
         if isinstance(http_connection, WebSocket):
-            self.set_metadata_from_websocket(user_message_id, conversation_id)
+            self.set_metadata_from_websocket(http_connection, user_message_id, conversation_id)
 
         if isinstance(http_connection, Request):
             self.set_metadata_from_http_request(http_connection)
@@ -161,11 +161,31 @@ class SessionManager:
         if request.headers.get("user-message-id"):
             self._context_state.user_message_id.set(request.headers["user-message-id"])
 
-    def set_metadata_from_websocket(self, user_message_id: str | None, conversation_id: str | None) -> None:
+    def set_metadata_from_websocket(self,
+                                    websocket: WebSocket,
+                                    user_message_id: str | None,
+                                    conversation_id: str | None) -> None:
         """
         Extracts and sets user metadata for Websocket connections.
         """
 
+        # Extract cookies from WebSocket headers (similar to HTTP request)
+        if websocket and hasattr(websocket, 'scope') and 'headers' in websocket.scope:
+            cookies = {}
+            for header_name, header_value in websocket.scope.get('headers', []):
+                if header_name == b'cookie':
+                    cookie_header = header_value.decode('utf-8')
+                    # Parse cookie header: "name1=value1; name2=value2"
+                    for cookie in cookie_header.split(';'):
+                        cookie = cookie.strip()
+                        if '=' in cookie:
+                            name, value = cookie.split('=', 1)
+                            cookies[name.strip()] = value.strip()
+
+            # Set cookies in metadata (same as HTTP request)
+            self._context.metadata._request.cookies = cookies
+            self._context_state.metadata.set(self._context.metadata)
+
         if conversation_id is not None:
             self._context_state.conversation_id.set(conversation_id)
 

{nvidia_nat-1.3.0a20250926.dist-info → nvidia_nat-1.3.0a20250928.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat
-Version: 1.3.0a20250926
+Version: 1.3.0a20250928
 Summary: NVIDIA NeMo Agent toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation
@@ -305,6 +305,7 @@ Requires-Dist: nat_router_agent; extra == "examples"
 Requires-Dist: nat_semantic_kernel_demo; extra == "examples"
 Requires-Dist: nat_sequential_executor; extra == "examples"
 Requires-Dist: nat_simple_auth; extra == "examples"
+Requires-Dist: nat_simple_auth_mcp; extra == "examples"
 Requires-Dist: nat_simple_web_query; extra == "examples"
 Requires-Dist: nat_simple_web_query_eval; extra == "examples"
 Requires-Dist: nat_simple_calculator; extra == "examples"

{nvidia_nat-1.3.0a20250926.dist-info → nvidia_nat-1.3.0a20250928.dist-info}/RECORD

@@ -21,10 +21,10 @@ nat/agent/tool_calling_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NM
 nat/agent/tool_calling_agent/agent.py,sha256=4SIp29I56oznPRQu7B3HCoX53Ri3_o3BRRYNJjeBkF8,11006
 nat/agent/tool_calling_agent/register.py,sha256=ijiRfgDVtt2p7_q1YbIQZmUVV8-jf3yT18HwtKyReUI,6822
 nat/authentication/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
-nat/authentication/interfaces.py,sha256=FAYM-QXVUn3a_8bmAZ7kP-lmN_BrLW8mo6krZJ3e0ME,3314
+nat/authentication/interfaces.py,sha256=1J2CWEJ_n6CLA3_HD3XV28CSbyfxrPAHzr7Q4kKDFdc,3511
 nat/authentication/register.py,sha256=lFhswYUk9iZ53mq33fClR9UfjJPdjGIivGGNHQeWiYo,915
 nat/authentication/api_key/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-nat/authentication/api_key/api_key_auth_provider.py,sha256=DoN1MoeB_uOWQMHUz2OiyhwbFxQkj0FOifN2TFr4sds,4022
+nat/authentication/api_key/api_key_auth_provider.py,sha256=QGRZZilD2GryhRJoODfKqypH54IwQp0I-Cck40Dc7dM,4032
 nat/authentication/api_key/api_key_auth_provider_config.py,sha256=zfkxH3yvUSKKldRf1K4PPm0rJLXGH0GDH8xj7anPYGQ,5472
 nat/authentication/api_key/register.py,sha256=Mhv3WyZ9H7C2JN8VuPvwlsJEZrwXJCLXCIokkN9RrP0,1147
 nat/authentication/credential_validator/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
@@ -32,10 +32,10 @@ nat/authentication/credential_validator/bearer_token_validator.py,sha256=cwGENd_
 nat/authentication/exceptions/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
 nat/authentication/exceptions/api_key_exceptions.py,sha256=6wnz951BI77rFYuHxoHOthz-y5oE08uxsuM6G5EvOyM,1545
 nat/authentication/http_basic_auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nat/authentication/http_basic_auth/http_basic_auth_provider.py,sha256=vmukbeI32aK6Ulc-yoaFHszr6BqmCvZHIpKr4-arMhs,3442
+nat/authentication/http_basic_auth/http_basic_auth_provider.py,sha256=OXr5TV87SiZtzSK9i_E6WXWyVhWq2MfqO_SS1aZ3p6U,3452
 nat/authentication/http_basic_auth/register.py,sha256=N2VD0vw7cYABsLxsGXl5yw0htc8adkrB0Y_EMxKwFfk,1235
 nat/authentication/oauth2/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
-nat/authentication/oauth2/oauth2_auth_code_flow_provider.py,sha256=iFp_kImFwDrweFHGTg4n9RGybH3-wRSJkyh-FPOxwDA,4538
+nat/authentication/oauth2/oauth2_auth_code_flow_provider.py,sha256=TDZvGbpCLW6EDCI-a4Z9KnQyUEcRaJ5hBCPVyIuy-KY,5099
 nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py,sha256=e165ysd2pX2WTbV3_FQKEjEaa4TAXkJ7B98WUGbqnGE,2204
 nat/authentication/oauth2/oauth2_resource_server_config.py,sha256=ltcNp8Dwb2Q4tlwMN5Cl0B5pouTLtXRoV-QopfqV45M,5314
 nat/authentication/oauth2/register.py,sha256=7rXhf-ilgSS_bUJsd9pOOCotL1FM8dKUt3ke1TllKkQ,1228
@@ -83,7 +83,7 @@ nat/cli/commands/info/info.py,sha256=BGqshIEDpNRH9hM-06k-Gq-QX-qNddPICSWCN-ReC-g
 nat/cli/commands/info/list_channels.py,sha256=K97TE6wtikgImY-wAbFNi0HHUGtkvIFd2woaG06VkT0,1277
 nat/cli/commands/info/list_components.py,sha256=QlAJVONBA77xW8Lx6Autw5NTAZNy_VrJGr1GL9MfnHM,4532
 nat/cli/commands/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-nat/cli/commands/mcp/mcp.py,sha256=CrpOOZZGq1nSsuFU_6h_VBL5ZE92Yr30WV9c5iirtzY,33332
+nat/cli/commands/mcp/mcp.py,sha256=Phtxegf0Ww89NHrwj4dEZh3vaHxjm7RV7NZMByUGhpM,43860
 nat/cli/commands/object_store/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
 nat/cli/commands/object_store/object_store.py,sha256=_ivB-R30a-66fNy-fUzi58HQ0Ay0gYsGz7T1xXoRa3Y,8576
 nat/cli/commands/registry/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
@@ -113,7 +113,7 @@ nat/control_flow/router_agent/register.py,sha256=p15Jy05PjJhXRrjqWiEgy47zmc-CFCu
 nat/data_models/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
 nat/data_models/agent.py,sha256=IwDyb9Zc3R4Zd5rFeqt7q0EQswczAl5focxV9KozIzs,1625
 nat/data_models/api_server.py,sha256=Zl0eAd-yV9PD8vUH8eWRvXFcUBdY2tENKe73q-Uxxgg,25699
-nat/data_models/authentication.py,sha256=t9-2CQaIfk74DM3CLaJU5NAYGs7m6BjdFFm1hj8WwtY,9064
+nat/data_models/authentication.py,sha256=wpH8tpy_wt1aA3eKIgbc8HfYdr_g6AfM7-NfWBZh-KA,8528
 nat/data_models/common.py,sha256=nXXfGrjpxebzBUa55mLdmzePLt7VFHvTAc6Znj3yEv0,5875
 nat/data_models/component.py,sha256=b_hXOA8Gm5UNvlFkAhsR6kEvf33ST50MKtr5kWf75Ao,1894
 nat/data_models/component_ref.py,sha256=KFDWFVCcvJCfBBcXTh9f3R802EVHBtHXh9OdbRqFmdM,4747
@@ -407,7 +407,7 @@ nat/retriever/nemo_retriever/retriever.py,sha256=gi3_qJFqE-iqRh3of_cmJg-SwzaQ3z2
 nat/runtime/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
 nat/runtime/loader.py,sha256=obUdAgZVYCPGC0R8u3wcoKFJzzSPQgJvrbU4OWygtog,7953
 nat/runtime/runner.py,sha256=Kzm5GRrGUFMQ_fbLOCJumYc4R-JXdTm5tUw2yMMDJpE,6450
-nat/runtime/session.py,sha256=U3UHQpdCBkCiJetsWdq9r6wUEVDBa2gv1VQedE64kY8,6959
+nat/runtime/session.py,sha256=DG4cpVg6GCVFY0cGzZnz55eLj0LoK5Q9Vg3NgbTqOHM,8029
 nat/runtime/user_metadata.py,sha256=ce37NRYJWnMOWk6A7VAQ1GQztjMmkhMOq-uYf2gNCwo,3692
 nat/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nat/settings/global_settings.py,sha256=JSlYnW3CeGJYGxlaXLz5F9xXf72I5TUz-w6qngG5di0,12476
@@ -469,10 +469,10 @@ nat/utils/reactive/base/observer_base.py,sha256=6BiQfx26EMumotJ3KoVcdmFBYR_fnAss
 nat/utils/reactive/base/subject_base.py,sha256=UQOxlkZTIeeyYmG5qLtDpNf_63Y7p-doEeUA08_R8ME,2521
 nat/utils/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nat/utils/settings/global_settings.py,sha256=9JaO6pxKT_Pjw6rxJRsRlFCXdVKCl_xUKU2QHZQWWNM,7294
-nvidia_nat-1.3.0a20250926.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
-nvidia_nat-1.3.0a20250926.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-nvidia_nat-1.3.0a20250926.dist-info/METADATA,sha256=AsBqos4EXJEwd-G-QSIF4iXItfZ7ULCpgwh0QCVd1lQ,22862
-nvidia_nat-1.3.0a20250926.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat-1.3.0a20250926.dist-info/entry_points.txt,sha256=4jCqjyETMpyoWbCBf4GalZU8I_wbstpzwQNezdAVbbo,698
-nvidia_nat-1.3.0a20250926.dist-info/top_level.txt,sha256=lgJWLkigiVZuZ_O1nxVnD_ziYBwgpE2OStdaCduMEGc,8
-nvidia_nat-1.3.0a20250926.dist-info/RECORD,,
+nvidia_nat-1.3.0a20250928.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat-1.3.0a20250928.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat-1.3.0a20250928.dist-info/METADATA,sha256=wgL7swlr3yFBn22PDRQ4CdjY7nIE_q2IgV_bKmiuZVc,22918
+nvidia_nat-1.3.0a20250928.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat-1.3.0a20250928.dist-info/entry_points.txt,sha256=4jCqjyETMpyoWbCBf4GalZU8I_wbstpzwQNezdAVbbo,698
+nvidia_nat-1.3.0a20250928.dist-info/top_level.txt,sha256=lgJWLkigiVZuZ_O1nxVnD_ziYBwgpE2OStdaCduMEGc,8
+nvidia_nat-1.3.0a20250928.dist-info/RECORD,,