Add option to enable redirect-based OAuth flow#1835
Add option to enable redirect-based OAuth flow#1835thepatrickchin wants to merge 13 commits intoNVIDIA:developfrom
Conversation
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
- "same-page" may be misleading as the popup is what keeps the user on the same page. Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
WalkthroughAdds an opt-in redirect-based OAuth2 authorization flow via a new boolean config Changes
Sequence DiagramssequenceDiagram
participant Browser as Browser (Client)
participant UI as Frontend UI
participant WS as WebSocket Handler
participant OAuth as OAuth Server
rect rgba(0, 150, 100, 0.5)
note over Browser,UI: Popup-based OAuth (use_redirect_auth: false)
end
Browser->>WS: Open WebSocket
WS->>UI: Send OAuth prompt (use_redirect=false)
UI->>UI: Open authorization URL in popup
Browser->>OAuth: Authorize (popup)
OAuth-->>Browser: Redirect callback (code)
Browser->>WS: GET /auth/redirect?code=...
WS->>OAuth: Exchange code for token
OAuth-->>WS: Return token
WS->>UI: Send success message (popup closes)
sequenceDiagram
participant Browser as Browser (Client)
participant WS as WebSocket Handler
participant OAuth as OAuth Server
rect rgba(100, 0, 150, 0.5)
note over Browser,OAuth: Redirect-based OAuth (use_redirect_auth: true)
end
Browser->>WS: Open WebSocket
WS->>Browser: Send OAuth prompt (use_redirect=true) + return_url
Browser->>OAuth: Full-page redirect to authorization
OAuth-->>Browser: Redirect callback (code)
Browser->>WS: GET /auth/redirect?code=...
WS->>OAuth: Exchange code for token
OAuth-->>WS: Return token
WS->>Browser: Return success HTML (redirect to return_url with marker)
Browser->>Browser: Resume original session (redirect target)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@packages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py`:
- Around line 120-121: The current WebSocketFlowHandler starts a same-tab
redirect flow even when config.use_popup_auth is False and self._return_url is
None; add a guard in the handler before constructing FlowState to fail fast: if
config.use_popup_auth is False and self._return_url is None, raise a clear
exception (or return an HTTP/WebSocket error) so the flow does not start without
a validated return URL; update the block where return_url is set and
FlowState(config=config, return_url=return_url) is created (references:
FlowState, config.use_popup_auth, self._return_url, websocket_flow_handler) to
implement this validation and error handling.
In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py`:
- Around line 107-110: The redirect-mode failure paths currently leave the
browser on a raw /auth/redirect error page; update the exception handlers in the
auth redirect route (the function handling flow_state and the except blocks for
OAuthError, httpx.HTTPError, and generic Exception) to mirror the success
behavior: when flow_state.config exists and flow_state.config.use_popup_auth is
False use build_auth_redirect_success_html (or a new
build_auth_redirect_error_html) to construct a redirect-capable error HTML using
flow_state.return_url, otherwise fall back to the existing
AUTH_REDIRECT_SUCCESS_HTML/ERROR constant; ensure the same flow_state,
use_popup_auth, build_auth_redirect_* helper and AUTH_REDIRECT_* constants are
referenced so failures perform a UI redirect like successes.
In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.py`:
- Around line 75-81: The origin membership check currently uses a literal
comparison against worker.front_end_config.cors.allow_origins which fails when
allow_origins contains the wildcard "*" and ignores allow_origin_regex; update
the logic before constructing WebSocketAuthenticationFlowHandler so that: if
allow_origins includes "*" accept any non-empty origin, otherwise check origin
equality against the list; additionally, if
worker.front_end_config.cors.allow_origin_regex is set, test origin against that
regex (treat a matching regex as allowed). Use the same variables (origin,
allowed_origins, allow_origin_regex) and pass the computed return_url into
WebSocketAuthenticationFlowHandler as before.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 7940c9bf-52f2-4d57-886c-0a4ab5fae921
📒 Files selected for processing (15)
docs/source/components/auth/api-authentication.mdexamples/front_ends/simple_auth/Dockerfileexamples/front_ends/simple_auth/patches/oauth2-server.patchexamples/front_ends/simple_auth/src/nat_simple_auth/configs/config.ymlpackages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.pypackages/nvidia_nat_core/src/nat/data_models/interactive.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/html_snippets/auth_code_grant_cancelled.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/html_snippets/auth_code_grant_success.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.pypackages/nvidia_nat_core/tests/nat/builder/test_interactive.pypackages/nvidia_nat_core/tests/nat/front_ends/auth_flow_handlers/test_websocket_flow_handler.pypackages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_auth_code_grant_cancelled.pypackages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_auth_code_grant_success.py
...ages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py
Outdated
Show resolved
Hide resolved
Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
- Fail fast when redirect auth has no validated return URL. - Handle redirect-mode failures the same way as redirect-mode success. - Properly handle wildcard "*" in allow_origins Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py (1)
84-123: 🛠️ Refactor suggestion | 🟠 MajorUse
logger.exception()to preserve stack traces in these catch-and-respond blocks.The three exception handlers (lines 85, 99, 112) catch exceptions without re-raising them, so they must use
logger.exception()instead oflogger.error()to capture the full stack trace information. This is mandatory per coding guidelines.♻️ Suggested fix
- except OAuthError as e: - logger.error("OAuth error during token exchange for state %s: %s (%s)", state, e.error, e.description) + except OAuthError as e: + logger.exception("OAuth error during token exchange for state %s: %s (%s)", state, e.error, e.description) if not flow_state.future.done(): flow_state.future.set_exception( RuntimeError(f"Authorization server rejected request: {e.error} ({e.description})")) if flow_state.config and flow_state.config.use_redirect_auth: error_html = build_auth_redirect_error_html(flow_state.return_url) @@ - except httpx.HTTPError as e: - logger.error("Network error during token fetch for state %s: %s", state, e) + except httpx.HTTPError as e: + logger.exception("Network error during token fetch for state %s: %s", state, e) if not flow_state.future.done(): flow_state.future.set_exception(RuntimeError(f"Network error during token fetch: {e}")) if flow_state.config and flow_state.config.use_redirect_auth: error_html = build_auth_redirect_error_html(flow_state.return_url) @@ - except Exception as e: - logger.error("Unexpected error during authentication for state %s: %s", state, e) + except Exception as e: + logger.exception("Unexpected error during authentication for state %s: %s", state, e) if not flow_state.future.done(): flow_state.future.set_exception(RuntimeError(f"Authentication failed: {e}")) if flow_state.config and flow_state.config.use_redirect_auth: error_html = build_auth_redirect_error_html(flow_state.return_url)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py` around lines 84 - 123, Replace the three logger.error calls in the exception handlers that catch OAuthError, httpx.HTTPError, and the generic Exception with logger.exception so the stack traces are preserved; update the handlers around the symbols OAuthError, httpx.HTTPError, and the final except Exception (the blocks that reference flow_state, build_auth_redirect_error_html, and AUTH_REDIRECT_ERROR_HTML) to call logger.exception with the same message format and parameters instead of logger.error, leaving the rest of the error handling (setting flow_state.future, building error_html, and returning the HTMLResponse) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py`:
- Around line 53-71: The code currently treats every OAuth callback error as a
user cancellation; change the logic to only treat error == "access_denied" as a
cancellation and keep the existing cancellation responses
(build_auth_redirect_cancelled_html, AUTH_REDIRECT_CANCELLED_POPUP_HTML,
flow_state.future.set_exception(...)) only in that branch; for all other error
values log the error (include error_description), set flow_state.future with a
clear exception mentioning the actual error (e.g. RuntimeError(f"OAuth error:
{error} ({error_description})")), call worker._remove_flow(state) as now, and
forward the request to the normal error UX path instead of the cancelled HTML
(i.e. do not return the cancelled HTML for non-access_denied errors).
---
Outside diff comments:
In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py`:
- Around line 84-123: Replace the three logger.error calls in the exception
handlers that catch OAuthError, httpx.HTTPError, and the generic Exception with
logger.exception so the stack traces are preserved; update the handlers around
the symbols OAuthError, httpx.HTTPError, and the final except Exception (the
blocks that reference flow_state, build_auth_redirect_error_html, and
AUTH_REDIRECT_ERROR_HTML) to call logger.exception with the same message format
and parameters instead of logger.error, leaving the rest of the error handling
(setting flow_state.future, building error_html, and returning the HTMLResponse)
unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: afa34493-7660-47ca-8f78-9d28245e697c
📒 Files selected for processing (12)
docs/source/components/auth/api-authentication.mdexamples/front_ends/simple_auth/src/nat_simple_auth/configs/config.ymlpackages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.pypackages/nvidia_nat_core/src/nat/data_models/interactive.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/html_snippets/auth_code_grant_error.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.pypackages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/websocket.pypackages/nvidia_nat_core/tests/nat/builder/test_interactive.pypackages/nvidia_nat_core/tests/nat/front_ends/auth_flow_handlers/test_websocket_flow_handler.pypackages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_auth_code_grant_error.pypackages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_websocket_route_origin.py
✅ Files skipped from review due to trivial changes (2)
- docs/source/components/auth/api-authentication.md
- packages/nvidia_nat_core/tests/nat/builder/test_interactive.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/front_ends/simple_auth/src/nat_simple_auth/configs/config.yml
- packages/nvidia_nat_core/src/nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py
- packages/nvidia_nat_core/src/nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py
packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py
Outdated
Show resolved
Hide resolved
- Only treat access_denied as a user cancellation. Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py (1)
101-140:⚠️ Potential issue | 🟠 MajorUse
logger.exception()to preserve full stack traces in exception handlers.Lines 102, 116, and 129 catch exceptions and log them with
logger.error()but do not re-raise. Per coding guidelines, you must uselogger.exception()to capture the full stack trace. The broadexcept Exceptionat line 128 (Ruff BLE001) is a secondary concern—prioritize fixing the logging pattern first.Suggested fix
except OAuthError as e: - logger.error("OAuth error during token exchange for state %s: %s (%s)", state, e.error, e.description) + logger.exception("OAuth error during token exchange for state %s: %s (%s)", state, e.error, e.description) if not flow_state.future.done(): flow_state.future.set_exception( RuntimeError(f"Authorization server rejected request: {e.error} ({e.description})")) if flow_state.config and flow_state.config.use_redirect_auth: error_html = build_auth_redirect_error_html(flow_state.return_url) else: error_html = AUTH_REDIRECT_ERROR_HTML return HTMLResponse(content=error_html, status_code=200, headers={ "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache" }) except httpx.HTTPError as e: - logger.error("Network error during token fetch for state %s: %s", state, e) + logger.exception("Network error during token fetch for state %s: %s", state, e) if not flow_state.future.done(): flow_state.future.set_exception(RuntimeError(f"Network error during token fetch: {e}")) if flow_state.config and flow_state.config.use_redirect_auth: error_html = build_auth_redirect_error_html(flow_state.return_url) else: error_html = AUTH_REDIRECT_ERROR_HTML return HTMLResponse(content=error_html, status_code=200, headers={ "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache" }) except Exception as e: - logger.error("Unexpected error during authentication for state %s: %s", state, e) + logger.exception("Unexpected error during authentication for state %s: %s", state, e) if not flow_state.future.done(): flow_state.future.set_exception(RuntimeError(f"Authentication failed: {e}")) if flow_state.config and flow_state.config.use_redirect_auth: error_html = build_auth_redirect_error_html(flow_state.return_url) else: error_html = AUTH_REDIRECT_ERROR_HTML return HTMLResponse(content=error_html, status_code=200, headers={ "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache" })🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py` around lines 101 - 140, The exception handlers in the auth route use logger.error which loses stack traces; update the three handlers catching OAuthError, httpx.HTTPError, and the generic Exception in the function handling the token exchange so they call logger.exception(...) instead of logger.error(...), leaving the message text and interpolation symbols (state, e, e.error, e.description) intact; keep the existing flow_state.future.set_exception(...) behavior and the HTMLResponse construction unchanged—only replace the logging calls to use logger.exception to preserve full stack traces for OAuthError, httpx.HTTPError, and the broad Exception handler.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py`:
- Around line 60-71: The code currently chooses redirect-cancel HTML whenever
flow_state.return_url is set, which ignores the configured popup vs redirect
mode; update the branching to use the flow state's mode flag
(flow_state.use_redirect_auth) instead of presence of flow_state.return_url.
Specifically, in the error == "access_denied" block, return
build_auth_redirect_cancelled_html(flow_state.return_url) only when
flow_state.use_redirect_auth is True (and return_url exists), otherwise return
AUTH_REDIRECT_CANCELLED_POPUP_HTML; keep the existing future.set_exception
behavior unchanged and reference flow_state, flow_state.use_redirect_auth,
flow_state.return_url, build_auth_redirect_cancelled_html, and
AUTH_REDIRECT_CANCELLED_POPUP_HTML when making the fix.
In
`@packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_auth_redirect_route.py`:
- Around line 77-96: Add a new test function (mirroring
test_access_denied_popup_returns_cancelled_html) that constructs
FlowState(config=_POPUP_CONFIG, return_url=_RETURN_URL), obtains a worker via
_make_worker(flow_state), calls _get with {"state": "teststate", "error":
"access_denied"}, and asserts response.status_code == 200 and that
"AUTH_CANCELLED" is present in response.text while "AUTH_ERROR" and any redirect
markers like _RETURN_URL.replace("/", "\\u002f") or "oauth_auth_completed" are
not; use the same helpers (_make_worker, _get) and constants (_POPUP_CONFIG,
_RETURN_URL, FlowState) to ensure popup mode remains independent of return_url.
---
Outside diff comments:
In `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.py`:
- Around line 101-140: The exception handlers in the auth route use logger.error
which loses stack traces; update the three handlers catching OAuthError,
httpx.HTTPError, and the generic Exception in the function handling the token
exchange so they call logger.exception(...) instead of logger.error(...),
leaving the message text and interpolation symbols (state, e, e.error,
e.description) intact; keep the existing flow_state.future.set_exception(...)
behavior and the HTMLResponse construction unchanged—only replace the logging
calls to use logger.exception to preserve full stack traces for OAuthError,
httpx.HTTPError, and the broad Exception handler.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: b57f31f0-f84d-4b39-ba5d-4ed124604737
📒 Files selected for processing (3)
examples/front_ends/simple_auth/src/nat_simple_auth/configs/config.ymlpackages/nvidia_nat_core/src/nat/front_ends/fastapi/routes/auth.pypackages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_auth_redirect_route.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/front_ends/simple_auth/src/nat_simple_auth/configs/config.yml
packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_auth_redirect_route.py
Show resolved
Hide resolved
- Gate cancellation UX with use_redirect_auth, not return_url. - Add a popup-mode regression test when return_url is present. Signed-off-by: Patrick Chin <8509935+thepatrickchin@users.noreply.github.com>
Description
This PR adds a
use_redirect_authoption in auth provider configurations to enable redirect-based OAuth flows while preserving popup flow by default.This features depends on the following NAT-UI PR: NVIDIA/NeMo-Agent-Toolkit-UI#127
Closes #1834
By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
use_redirect_auth) to choose redirect vs popup; redirect mode supports returning users to a provided return URL and returns HTML responses for success/error/cancel flows.Bug Fixes
Documentation
use_redirect_authand its default behavior.Tests