Add Device Authorization Grant (RFC 8628) for headless/container environments

python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions.py

@@ -211,13 +211,26 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
+class _ReauthSucceeded(Exception):
+    """Internal signal that re-authentication succeeded and the caller should retry."""
+
+
 def _handle_connection_error_with_reauth(exc, login_func):
-    """Handle ConnectionError with reauthentication logic."""
+    """Handle ConnectionError with reauthentication logic.
+
+    If the error indicates an expired token, triggers re-authentication via
+    ``login_func`` and raises ``_ReauthSucceeded`` so the calling decorator
+    can retry the original operation.  Non-expired connection errors are
+    surfaced as ``ClickExceptionRed`` directly.
+    """
     if "expired" in str(exc).lower():
-        click.echo(click.style("Token is expired, triggering re-authentication", fg="red"))
+        click.echo(click.style("Token is expired, triggering re-authentication", fg="yellow"))
         config = exc.get_config()
-        login_func(config)
-        raise ClickExceptionRed("Please try again now") from None
+        try:
+            login_func(config)
+        except Exception as reauth_exc:
+            raise ClickExceptionRed(f"Re-authentication failed: {reauth_exc}") from None
+        raise _ReauthSucceeded() from None
     else:
         raise ClickExceptionRed(str(exc)) from None
 
@@ -239,26 +252,78 @@ def _handle_exception_group_with_reauth(eg, login_func) -> NoReturn:
     raise eg
 
 
+def _raise_if_mappable(exc: BaseException) -> None:
+    """Raise a user-friendly ClickException if *exc* maps to one, otherwise return."""
+    if cli_exc := _map_cli_exception(exc):
+        raise cli_exc from None
+
+
+def _try_reauth_or_handle(handler, exc, login_func, *, allow_reauth: bool) -> bool:
+    """Attempt re-auth via *handler*; return True if re-auth succeeded.
+
+    When *allow_reauth* is False the handler is called without catching
+    ``_ReauthSucceeded`` — any re-auth signal propagates as a normal error.
+    """
+    if allow_reauth:
+        try:
+            handler(exc, login_func)
+        except _ReauthSucceeded:
+            return True
+    else:
+        handler(exc, login_func)
+    return False
+
+
+def _call_with_exception_handling(func, args, kwargs, login_func, *, allow_reauth: bool):
+    """Call *func* and handle exceptions, optionally allowing re-authentication.
+
+    Returns ``(result, needs_retry)`` where *needs_retry* is ``True`` when
+    re-authentication succeeded and the caller should retry the call.
+    """
+    try:
+        return func(*args, **kwargs), False
+    except _ReauthSucceeded:
+        raise ClickExceptionRed("Unexpected re-auth signal") from None
+    except BaseExceptionGroup as eg:
+        if _try_reauth_or_handle(_handle_exception_group_with_reauth, eg, login_func, allow_reauth=allow_reauth):
+            return None, True
+    except (ConnectionError, JumpstarterException, click.ClickException) as e:
+        if isinstance(e, ConnectionError) and not allow_reauth:
+            raise ClickExceptionRed(str(e)) from None
+        if _try_reauth_or_handle(_handle_single_exception_with_reauth, e, login_func, allow_reauth=allow_reauth):
+            return None, True
+    except Exception as e:
+        _raise_if_mappable(e)
+        raise
+    except KeyboardInterrupt as e:
+        _raise_if_mappable(e)
+        raise
+    return None, False  # pragma: no cover — handlers above always raise or return
+
+
 def handle_exceptions_with_reauthentication(login_func):
-    """Decorator to handle exceptions in blocking functions, including those wrapped in BaseExceptionGroup."""
+    """Decorator to handle exceptions in blocking functions, including those wrapped in BaseExceptionGroup.
+
+    When a ``ConnectionError`` with an expired-token message is caught, the
+    decorator triggers re-authentication via *login_func* and **retries the
+    original call exactly once**.  If the retry also fails, the error is
+    surfaced normally — no infinite loop.
+    """
 
     def decorator(func):
         @wraps(func)
         def wrapped(*args, **kwargs):
-            try:
-                return func(*args, **kwargs)
-            except BaseExceptionGroup as eg:
-                _handle_exception_group_with_reauth(eg, login_func)
-            except (ConnectionError, JumpstarterException, click.ClickException) as e:
-                _handle_single_exception_with_reauth(e, login_func)
-            except Exception as e:
-                if cli_exc := _map_cli_exception(e):
-                    raise cli_exc from None
-                raise
-            except KeyboardInterrupt as e:
-                if cli_exc := _map_cli_exception(e):
-                    raise cli_exc from None
-                raise
+            result, needs_retry = _call_with_exception_handling(
+                func, args, kwargs, login_func, allow_reauth=True
+            )
+            if not needs_retry:
+                return result
+
+            click.echo(click.style("Re-authenticated, retrying...", fg="yellow"))
+            result, _ = _call_with_exception_handling(
+                func, args, kwargs, login_func, allow_reauth=False
+            )
+            return result
 
         return wrapped
 

python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions_test.py

@@ -1,5 +1,6 @@
 import ssl
 from json import JSONDecodeError
+from unittest.mock import MagicMock
 
 import click
 import pytest
@@ -10,6 +11,8 @@
     handle_exceptions_with_reauthentication,
 )
 
+from jumpstarter.common.exceptions import ConnectionError as JmpConnectionError
+
 
 @pytest.fixture
 def anyio_backend():
@@ -183,3 +186,79 @@ def grpc_precondition_no_details_fn():
 
     with pytest.raises(click.ClickException, match="precondition"):
         grpc_precondition_no_details_fn()
+
+
+# ---------------------------------------------------------------------------
+# Tests for automatic retry after successful re-authentication (NS-REQ-1/3)
+# ---------------------------------------------------------------------------
+
+
+def _make_expired_connection_error():
+    """Create a ConnectionError that looks like an expired-token error."""
+    exc = JmpConnectionError("token expired")
+    config = MagicMock(name="client_config")
+    exc.set_config(config)
+    return exc, config
+
+
+def test_reauth_retries_on_success() -> None:
+    """After successful re-auth the decorator retries and returns the result (TS-NS-1)."""
+    call_count = 0
+
+    def login_func(_config):
+        pass  # success
+
+    @handle_exceptions_with_reauthentication(login_func)
+    def fn():
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            exc, _ = _make_expired_connection_error()
+            raise exc
+        return "sentinel"
+
+    result = fn()
+    assert result == "sentinel"
+    assert call_count == 2
+
+
+def test_reauth_failure_raises_click_exception() -> None:
+    """If login_func raises, the decorator surfaces a ClickException (TS-NS-2)."""
+    call_count = 0
+
+    def login_func(_config):
+        raise RuntimeError("IdP unreachable")
+
+    @handle_exceptions_with_reauthentication(login_func)
+    def fn():
+        nonlocal call_count
+        call_count += 1
+        exc, _ = _make_expired_connection_error()
+        raise exc
+
+    with pytest.raises(click.ClickException, match="Re-authentication failed"):
+        fn()
+
+    # The wrapped function should have been called only once (no retry).
+    assert call_count == 1
+
+
+def test_reauth_retry_bounded_to_one_attempt() -> None:
+    """The decorator retries at most once; a second failure surfaces normally (TS-NS-3)."""
+    call_count = 0
+
+    def login_func(_config):
+        pass  # always succeeds
+
+    @handle_exceptions_with_reauthentication(login_func)
+    def fn():
+        nonlocal call_count
+        call_count += 1
+        exc, _ = _make_expired_connection_error()
+        raise exc
+
+    with pytest.raises(click.ClickException):
+        fn()
+
+    # Original call + exactly one retry = 2 total.
+    assert call_count == 2