feat(cli): OAuth 2.0 Device Authorization Grant, silent token refresh, and TLS fixes

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

@@ -4,7 +4,6 @@
 import types
 from functools import wraps
 from types import TracebackType
-from typing import NoReturn
 
 import click
 
@@ -205,45 +204,62 @@
 
 
 def _handle_connection_error_with_reauth(exc, login_func):
-    """Handle ConnectionError with reauthentication logic."""
+    """Handle ConnectionError with reauthentication logic.
+
+    Returns True if re-auth succeeded and the caller should retry.
+    Raises if re-auth is not applicable or fails.
+    """
     if "expired" in str(exc).lower():
-        click.echo(click.style("Token is expired, triggering re-authentication", fg="red"))
         config = exc.get_config()
         login_func(config)
-        raise ClickExceptionRed("Please try again now") from None
+        return True
     else:
         raise ClickExceptionRed(str(exc)) from None
 
 
 def _handle_single_exception_with_reauth(exc, login_func):
-    """Handle a single exception (may raise)."""
+    """Handle a single exception (may raise).
+
+    Returns True if re-auth succeeded and the caller should retry.
+    """
     if isinstance(exc, ConnectionError):
-        _handle_connection_error_with_reauth(exc, login_func)
+        return _handle_connection_error_with_reauth(exc, login_func)
     elif cli_exc := _map_cli_exception(exc):
         raise cli_exc from None
-    # Not handled: fall through
+    return False
+
 
+def _handle_exception_group_with_reauth(eg, login_func):
+    """Handle exceptions wrapped in BaseExceptionGroup.
 
-def _handle_exception_group_with_reauth(eg, login_func) -> NoReturn:
-    """Handle exceptions wrapped in BaseExceptionGroup."""
+    Returns True if re-auth succeeded and the caller should retry.
+    """
     for exc in leaf_exceptions(eg, fix_tracebacks=False):
-        _handle_single_exception_with_reauth(exc, login_func)
+        if _handle_single_exception_with_reauth(exc, login_func):
+            return True
     # If no handled exceptions, re-raise the original group
     raise eg
 
 
 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 the wrapped function raises a token-expired error, the decorator re-authenticates
+    via login_func and retries the function exactly once.
+    """
 
     def decorator(func):
         @wraps(func)
         def wrapped(*args, **kwargs):
+            retry = False
             try:
                 return func(*args, **kwargs)
             except BaseExceptionGroup as eg:
-                _handle_exception_group_with_reauth(eg, login_func)
+                if _handle_exception_group_with_reauth(eg, login_func):
+                    retry = True
             except (ConnectionError, JumpstarterException, click.ClickException) as e:
-                _handle_single_exception_with_reauth(e, login_func)
+                if _handle_single_exception_with_reauth(e, login_func):
+                    retry = True
             except Exception as e:
                 if cli_exc := _map_cli_exception(e):
                     raise cli_exc from None
@@ -253,6 +269,18 @@
                     raise cli_exc from None
                 raise
 
+            if retry:
+                try:
+                    return func(*args, **kwargs)
+                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
+
         return wrapped
 
     return decorator

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

@@ -137,6 +137,53 @@ def grpc_auth_fn():
         grpc_auth_fn()
 
 
+def test_handle_exceptions_with_reauth_retries_on_expired_token() -> None:
+    """After successful re-auth, the command is retried automatically."""
+    from jumpstarter.common.exceptions import ConnectionError
+
+    call_count = [0]
+
+    def login_func(config):
+        config.token = "new_token"
+
+    @handle_exceptions_with_reauthentication(login_func)
+    def command_fn(config=None):
+        call_count[0] += 1
+        if call_count[0] == 1:
+            exc = ConnectionError("token is expired")
+            exc.set_config(config)
+            raise exc
+        return "result"
+
+    config = type("Config", (), {"token": "old_token"})()
+    result = command_fn(config=config)
+
+    assert result == "result"
+    assert call_count[0] == 2
+
+
+def test_handle_exceptions_with_reauth_does_not_retry_twice() -> None:
+    """If retry also fails with expired token, the error propagates."""
+    from jumpstarter.common.exceptions import ConnectionError
+
+    login_calls = [0]
+
+    def login_func(config):
+        login_calls[0] += 1
+
+    @handle_exceptions_with_reauthentication(login_func)
+    def always_expired_fn(config=None):
+        exc = ConnectionError("token is expired")
+        exc.set_config(config)
+        raise exc
+
+    config = type("Config", (), {"token": "tok"})()
+    with pytest.raises(click.ClickException):
+        always_expired_fn(config=config)
+
+    assert login_calls[0] == 1
+
+
 def test_handle_exceptions_maps_grpc_invalid_argument() -> None:
     class MockGrpcError(Exception):
         def code(self):

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

@@ -2,21 +2,43 @@
 import os
 import ssl
 import time
+import warnings
 from dataclasses import dataclass
 from functools import wraps
 from typing import ClassVar
 
 import aiohttp
+import anyio
 import certifi
 import click
 from aiohttp import web
 from anyio import create_memory_object_stream
 from anyio.to_thread import run_sync
+
+# Authlib still pulls deprecated authlib.jose on first OAuth2Session use, not only at
+# import time; a transient catch_warnings() around the import does not suppress that.
+warnings.filterwarnings(
+    "ignore",
+    message=r"authlib\.jose module is deprecated.*",
+    module=r"authlib\..*",
+)
+
+# When the user opts into --insecure-tls, we set verify=False on the requests session.
+# urllib3 emits InsecureRequestWarning for every such request; suppress it since the
+# user has already acknowledged the risk.
+warnings.filterwarnings(
+    "ignore",
+    message=r"Unverified HTTPS request is being made.*",
+    module=r"urllib3\..*",
+)
 from authlib.integrations.requests_client import OAuth2Session
 from joserfc.jws import extract_compact
 from yarl import URL
 
-from jumpstarter.config.env import JMP_OIDC_CALLBACK_PORT
+from jumpstarter.config.env import JMP_OIDC_CALLBACK_PORT, JMP_OIDC_DEVICE_FLOW
+
+DEVICE_FLOW_POLL_INTERVAL = 5
+DEVICE_FLOW_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code"
 
 
 def _get_ssl_context() -> ssl.SSLContext:
@@ -46,6 +68,12 @@
         default=True,
         help="Request offline_access scope (refresh token)",
     )
+    @click.option(
+        "--device-flow/--no-device-flow",
+        "device_flow",
+        default=None,
+        help="Use OAuth 2.0 Device Authorization Grant (auto-detected in Dev Spaces)",
+    )
     @wraps(f)
     def wrapper(*args, **kwds):
         return f(*args, **kwds)
@@ -177,6 +205,84 @@
             lambda: client.fetch_token(config["token_endpoint"], authorization_response=authorization_response)
         )
 
+    async def device_code_grant(self):
+        config = await self.configuration()
+
+        device_endpoint = config.get("device_authorization_endpoint")
+        if not device_endpoint:
+            raise click.ClickException(
+                "OIDC provider does not support Device Authorization Grant. "
+                "Ensure 'oauth2DeviceAuthorizationGrantEnabled' is enabled on the OIDC client."
+            )
+
+        ssl_context: ssl.SSLContext | bool = False if self.insecure_tls else _get_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.post(
+                device_endpoint,
+                data={
+                    "client_id": self.client_id,
+                    "scope": " ".join(self._scopes()),
+                },
+            ) as resp:
+                if resp.status != 200:
+                    body = await resp.text()
+                    raise click.ClickException(f"Device authorization request failed (HTTP {resp.status}): {body}")
+                device_data = await resp.json()
+
+        device_code = device_data["device_code"]
+        user_code = device_data["user_code"]
+        verification_uri = device_data.get("verification_uri_complete") or device_data["verification_uri"]
+        expires_in = device_data.get("expires_in", 600)
+        interval = device_data.get("interval", DEVICE_FLOW_POLL_INTERVAL)
+
+        print(
+            f"\nOpen the following URL in your browser and enter the code:"
+            f"\n  {verification_uri}"
+            f"\n  Code: {user_code}\n"
+        )
+
+        token_endpoint = config["token_endpoint"]
+        deadline = time.monotonic() + expires_in
+
+        while time.monotonic() < deadline:
+            await anyio.sleep(interval)
+
+            client = self.client()
+            try:
+                token = await run_sync(
+                    lambda c=client: c.fetch_token(
+                        token_endpoint,
+                        grant_type=DEVICE_FLOW_GRANT_TYPE,
+                        device_code=device_code,
+                    )
+                )
+                return token
+            except Exception as e:
+                error_msg = str(e)
+                if "authorization_pending" in error_msg:
+                    continue
+                if "slow_down" in error_msg:
+                    interval += 5
+                    continue
+                if "expired_token" in error_msg or "expired" in error_msg:
+                    raise click.ClickException("Device code expired. Please try again.") from e
+                if "access_denied" in error_msg:
+                    raise click.ClickException("Authorization was denied by the user.") from e
+                raise click.ClickException(f"Device flow token request failed: {e}") from e
+
+        raise click.ClickException("Device code expired (timeout). Please try again.")
+
+
+def should_use_device_flow(device_flow_flag: bool | None = None) -> bool:
+    if device_flow_flag is not None:
+        return device_flow_flag
+    if os.environ.get(JMP_OIDC_DEVICE_FLOW, "").strip() == "1":
+        return True
+    if os.environ.get("VSCODE_INJECTION", "").strip() == "1":
+        return True
+    return False
+
 
 def decode_jwt(token: str):
     try: