feat: add pre_authenticate option for authorization providers

- This makes pre-authentication an opt-in feature
- When True, authentication for the provider is triggered at WebSocket connection time before any user message is submitted.
- When False (default), authentication only occurs when the workflow explicitly requires it.

Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>

Author: thepatrickchin

examples/front_ends/simple_auth/src/nat_simple_auth/configs/config.yml

@@ -56,6 +56,7 @@ authentication:
       - email
     client_id: ${NAT_OAUTH_CLIENT_ID}
     client_secret: ${NAT_OAUTH_CLIENT_SECRET}
+    pre_authenticate: false
     use_pkce: false
     use_redirect_auth: false
 

packages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py

@@ -41,6 +41,12 @@ class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth
                      "remains open. When True, the browser navigates to the OAuth login page directly and is "
                      "redirected back after authentication completes."))
 
+    pre_authenticate: bool = Field(
+        default=False,
+        description=("When True, authentication is triggered at WebSocket connection time before the user submits "
+                     "their first prompt. When False (default), authentication is deferred until the workflow first "
+                     "requires credentials. Only applies in nat serve mode."))
+
     authorization_kwargs: dict[str, str] | None = Field(description=("Additional keyword arguments for the "
                                                                      "authorization request."),
                                                         default=None)

packages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -78,11 +78,12 @@ async def authenticate(
     async def pre_authenticate(self, auth_providers: dict[str, AuthProviderBaseConfig]) -> None:
         """Run auth for every configured OAuth2 provider before the first user message.
 
+        Only providers with pre_authenticate option set in their config are processed.
         Returns immediately if tokens are already cached. Otherwise triggers the OAuth
         redirect so the user authenticates at page load rather than mid-workflow.
         """
         for provider_config in auth_providers.values():
-            if isinstance(provider_config, OAuth2AuthCodeFlowProviderConfig):
+            if isinstance(provider_config, OAuth2AuthCodeFlowProviderConfig) and provider_config.pre_authenticate:
                 await self.authenticate(provider_config, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
 
     def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig) -> AsyncOAuth2Client:

packages/nvidia_nat_core/tests/nat/front_ends/auth_flow_handlers/test_websocket_flow_handler.py

@@ -508,6 +508,12 @@ async def test_pre_authenticate_skips_non_oauth2_providers(noop_handler):
     await noop_handler.pre_authenticate({"my_api_key": api_key_config})
 
 
+async def test_pre_authenticate_skips_oauth2_provider_flag_false(noop_handler, minimal_oauth_config):
+    """pre_authenticate does not trigger auth for OAuth2 providers with pre_authenticate=False (the default)."""
+    # minimal_oauth_config has pre_authenticate=False (the default); if the guard were absent this would hang
+    await noop_handler.pre_authenticate({"my_provider": minimal_oauth_config})
+
+
 async def test_pre_authenticate_uses_cached_token(minimal_oauth_config):
     """pre_authenticate returns immediately without calling create_websocket_message on a cache hit."""
 
@@ -526,15 +532,17 @@ async def create_websocket_message(self, msg):
 
     ctx = AuthenticatedContext(headers={"Authorization": "Bearer cached-tok"}, metadata={"expires_at": None})
     store: dict = {}
+    # Enable pre_authenticate so the cache lookup is actually reached
+    active_config = minimal_oauth_config.model_copy(update={"pre_authenticate": True})
     handler = WebSocketAuthenticationFlowHandler(
         add_flow_cb=_noop_add,
         remove_flow_cb=_noop_remove,
         web_socket_message_handler=_CountingWSHandler(),
         token_store=store,
         session_id="sess-1",
     )
-    key = handler._token_cache_key(minimal_oauth_config)
+    key = handler._token_cache_key(active_config)
     store[key] = (ctx, time.time() + 3600)
 
-    await handler.pre_authenticate({"my_provider": minimal_oauth_config})
+    await handler.pre_authenticate({"my_provider": active_config})
     assert message_count[0] == 0, "pre_authenticate must not trigger OAuth when token is cached"