fix: prevent auth redirect loop when user cancels or declines consent

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

Author: thepatrickchin

packages/nvidia_nat_core/src/nat/front_ends/fastapi/html_snippets/auth_code_grant_cancelled.py

@@ -45,7 +45,9 @@
         (function () {
             var returnTo = RETURN_URL_PLACEHOLDER;
             if (returnTo) {
-                window.location.replace(returnTo);
+                var url = new URL(returnTo);
+                url.searchParams.set('oauth_auth_error', 'cancelled');
+                window.location.replace(url.toString());
             } else {
                 window.history.back();
             }
@@ -60,14 +62,13 @@
 
 
 def build_auth_redirect_cancelled_html(return_url: str | None = None) -> str:
-    """Build the authorization-cancelled HTML page.
-
-    Redirects back to the UI without the ``oauth_auth_completed`` query
-    parameter so the UI's cancellation-message branch handles it.
+    """Build the same-page authorization-cancelled HTML page.
 
     Args:
-        return_url: The UI origin to navigate back to. Falls back to
-            ``window.history.back()`` when not provided.
+        return_url: The URL to redirect to after cancellation. When
+            provided the page navigates there immediately with an ``oauth_auth_error``
+            query parameter so the UI can detect the cancellation and avoid a
+            pre-auth redirect loop; otherwise it falls back to ``window.history.back()``.
 
     Returns:
         An HTML string for the post-cancellation redirect page.

packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.py

@@ -104,7 +104,12 @@ async def _websocket_endpoint(websocket: WebSocket):
                                                               token_store=worker._oauth_token_store,
                                                               session_id=nat_session_id)
             handler.set_flow_handler(flow_handler)
-            await flow_handler.pre_authenticate(worker._config.authentication)
+            skip_pre_auth = websocket.query_params.get("skip_pre_auth") == "true"
+            if not skip_pre_auth:
+                try:
+                    await flow_handler.pre_authenticate(worker._config.authentication)
+                except Exception as e:
+                    logger.info("Pre-authentication did not complete: %s", e)
             await handler.run()
 
     return _websocket_endpoint