feat(cli): ship the PKCE login and auth commands

.github/workflows/test-cli.yaml

@@ -41,15 +41,8 @@ jobs:
     - name: show pip list
       run: pip list
 
-    - name: Login passthrough
-      run: lightning login # since we are providing the credentials, this should be a no-op
-      env:
-        LIGHTNING_API_KEY: dummy-api-key
-
     - name: Tests
       working-directory: python/tests
-      env: # the following value is just a dummy since the actual http calls are mocked out
-        LIGHTNING_API_KEY: dummy-api-key
       run: |
         # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003
         coverage run --source lightning_sdk -m pytest cli/ -vv -s

python/lightning_sdk/api/billing_api.py

@@ -27,9 +27,9 @@
 import requests
 
 from lightning_sdk.api.utils import (
-    _authenticate_and_get_auth_headers,
     _raise_for_download_status,
     cached_lightning_client,
+    credential_auth,
 )
 
 
@@ -309,7 +309,7 @@ def _download_activity_csv(
         r = requests.get(
             f"{self._client.api_client.configuration.host}{endpoint}",
             params=query_params,
-            headers=_authenticate_and_get_auth_headers(),
+            auth=credential_auth(),
             stream=True,
             allow_redirects=True,
         )

python/lightning_sdk/api/filesystem_api.py

@@ -12,7 +12,6 @@
     _DOWNLOAD_MAX_TRIES,
     _DOWNLOAD_READ_TIMEOUT_SECONDS,
     _RETRYABLE_DOWNLOAD_STATUSES,
-    _authenticate_and_get_auth_headers,
     _BlobUploader,
     _collect_download_results,
     _paged_tree_entries,
@@ -22,6 +21,8 @@
     _stream_download_to_file,
     _TransientDownloadError,
     cached_lightning_client,
+    credential_auth,
+    request_auth,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
@@ -31,7 +32,6 @@ class FilesystemApi:
 
     def __init__(self) -> None:
         self._client = cached_lightning_client()
-        self._auth_headers = _authenticate_and_get_auth_headers()
 
     @property
     def client(self) -> LightningClient:
@@ -75,7 +75,7 @@ def fetch_page(query_params: Dict[str, str]) -> Dict:
             r = requests.get(
                 f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/trees/{path}",
                 params=query_params,
-                headers=self._auth_headers,
+                auth=credential_auth(),
             )
             if r.status_code != 200:
                 raise _RemoteApiError(f"Failed to list files: {r.status_code}", status_code=r.status_code)
@@ -135,6 +135,7 @@ def _download_single_file(
         pbar: Optional[tqdm] = None,
         progress_bar: bool = False,
         expected_size: Optional[int] = None,
+        auth: Optional[object] = None,
     ) -> None:
         """Download a single artifact file.
 
@@ -157,14 +158,15 @@ def _download_single_file(
         """
         url = f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/blobs/{remote_path}"
         owned_pbar = None
+        request = request_auth(auth)
 
         @backoff.on_exception(backoff.expo, _TransientDownloadError, max_tries=_DOWNLOAD_MAX_TRIES)
         def download() -> None:
             nonlocal owned_pbar, pbar
             try:
                 r = requests.get(
                     url,
-                    headers=self._auth_headers,
+                    auth=request,
                     stream=True,
                     allow_redirects=True,
                     timeout=(_DOWNLOAD_CONNECT_TIMEOUT_SECONDS, _DOWNLOAD_READ_TIMEOUT_SECONDS),
@@ -253,6 +255,7 @@ def download_folder(
                 mininterval=1,
             )
 
+        auth = credential_auth()
         try:
             with ThreadPoolExecutor(max_workers=num_workers) as executor:
                 futures = [
@@ -263,6 +266,7 @@ def download_folder(
                         teamspace_id,
                         pbar,
                         expected_size=entry.get("size"),
+                        auth=auth,
                     )
                     for entry in files
                 ]
@@ -288,7 +292,7 @@ def delete_file(self, teamspace_id: str, remote_path: str) -> None:
         r = requests.delete(
             f"{self._client.api_client.configuration.host}"
             f"/v1/projects/{teamspace_id}/artifacts/blobs/{remote_path.strip('/')}",
-            headers=self._auth_headers,
+            auth=credential_auth(),
         )
         if r.status_code not in (200, 204):
             raise _RemoteApiError(
@@ -311,7 +315,7 @@ def delete_folder(self, teamspace_id: str, remote_path: str) -> None:
         r = requests.delete(
             f"{self._client.api_client.configuration.host}"
             f"/v1/projects/{teamspace_id}/artifacts/trees/{remote_path.strip('/')}",
-            headers=self._auth_headers,
+            auth=credential_auth(),
         )
         if r.status_code not in (200, 204):
             raise _RemoteApiError(

python/lightning_sdk/api/license_api.py

@@ -9,8 +9,8 @@ class LicenseApi:
     def __init__(self, login_token: str) -> None:
         self._cloud_url = _cloud_url()
         self._auth = Auth()
-        self._auth.token_login(login_token, save_token=True)
-        self._client = self._auth.create_api_client()
+        self._auth.token_login(login_token, save_token=False)
+        self._client = self._auth.create_api_client(override="auth_token")
         self._api = ProductLicenseServiceApi(self._client)
 
     def validate_license(self, license_key: str, product_id: str) -> bool:

python/lightning_sdk/api/llm_api.py

@@ -7,13 +7,11 @@
 import time
 import warnings
 from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
-from urllib.parse import urlencode
 
 from rich.prompt import Confirm
 
 from lightning_sdk.api.utils import cached_lightning_client
-from lightning_sdk.lightning_cloud import env
-from lightning_sdk.lightning_cloud.login import Auth, AuthServer
+from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login, start_login
 from lightning_sdk.lightning_cloud.openapi.models import (
     StreamResultOfV1ConversationResponseChunk,
     V1ConversationResponseChunk,
@@ -25,45 +23,27 @@
 LITAI_CODE = os.environ.get("LITAI_CODE", "x334uv8t7v")
 
 
-class _AuthServer(AuthServer):
-    def __init__(self, model: str, *args: Any, **kwargs: Any) -> None:
-        self._model = model
-        super().__init__(*args, **kwargs)
-
-    def get_auth_url(self, port: int) -> str:
-        redirect_uri = f"http://localhost:{port}/login-complete"
-        params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITAI_CODE, "litaimodel": self._model})
-        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
-
-
-class _AuthLitAI(Auth):
-    def __init__(self, model: str, shall_confirm: bool = False) -> None:
-        super().__init__()
-        self._model = model
-        self._shall_confirm = shall_confirm
-
-    def _run_server(self) -> None:
-        if self._shall_confirm:
-            proceed = Confirm.ask(
-                "[bold yellow]LitAI needs to authenticate with Lightning AI.[/bold yellow]\n"
-                "This will open a browser window for login.\n"
-                "Do you want to continue?",
-                default=True,
-            )
-            if not proceed:
-                raise RuntimeError(
-                    "Login cancelled. Please login at https://platform.lightning.ai/sign-up. "
-                    "Or run `lightning login` to login."
-                )
-        print("Opening browser for authentication...")
-        print("Please come back to the terminal after logging in.")
-        time.sleep(1)
-        _AuthServer(self._model).login_with_browser(self)
-
-
 def authenticate(model: str, shall_confirm: bool = True) -> None:
-    auth = _AuthLitAI(model, shall_confirm)
-    auth.authenticate()
+    with browser_login(False):
+        existing = CredentialSource().resolve()
+    if existing is not None:
+        return
+    if shall_confirm:
+        proceed = Confirm.ask(
+            "[bold yellow]LitAI needs to authenticate with Lightning AI.[/bold yellow]\n"
+            "This will open a browser window for login.\n"
+            "Do you want to continue?",
+            default=True,
+        )
+        if not proceed:
+            raise RuntimeError(
+                "Login cancelled. Please login at https://platform.lightning.ai/sign-up. "
+                "Or run `lightning login` to login."
+            )
+    print("Opening browser for authentication...")
+    print("Please come back to the terminal after logging in.")
+    time.sleep(1)
+    start_login(sign_in_params={"okbhrt": LITAI_CODE, "litaimodel": model})
 
 
 class LLMApi:

python/lightning_sdk/api/logs_api.py

@@ -22,8 +22,9 @@
 from urllib.parse import urlparse
 
 from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.lightning_cloud.credentials import CredentialSource
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.lightning_cloud.rest_client import _LOGIN_MESSAGE, LightningClient
 
 # While tailing, ``recv()`` is given this timeout so a quiet socket can be re-checked (the server
 # holds the connection open with heartbeats rather than closing it). It must stay below the
@@ -497,6 +498,7 @@ def follow(
         try:
             import websocket
             from websocket import (
+                WebSocketBadStatusException,
                 WebSocketConnectionClosedException,
                 WebSocketException,
                 WebSocketTimeoutException,
@@ -506,7 +508,6 @@ def follow(
                 "Following logs requires the 'websocket-client' package (pip install websocket-client)."
             ) from ex
 
-        auth_header = Auth().authenticate()
         url = _websocket_url(follow_url)
         # Poll `recv()` rather than blocking on it, so `stop` and `idle_timeout` get a chance to
         # run. The poll has to stay below the server's heartbeat interval: a heartbeat resets the
@@ -518,9 +519,11 @@ def follow(
         last_line_at = time.monotonic()
 
         backoff = _MIN_RECONNECT_BACKOFF
+        refreshed = False
         while True:
             ws = None
             try:
+                auth_header = Auth().auth_header
                 # The handshake gets a longer budget than the deliberately short read poll.
                 ws = websocket.create_connection(
                     url, header=[f"Authorization: {auth_header}"], timeout=_CONNECT_TIMEOUT
@@ -530,6 +533,7 @@ def follow(
                     on_socket(ws)
                 # Connected cleanly: reset the backoff so the next blip recovers fast
                 backoff = _MIN_RECONNECT_BACKOFF
+                refreshed = False
                 while True:
                     try:
                         message = ws.recv()
@@ -554,6 +558,15 @@ def follow(
                                 continue
                             recent.add(entry)
                         yield entry
+            except WebSocketBadStatusException as err:
+                # Reconnecting with a credential the server refused would loop forever.
+                if err.status_code == 401:
+                    if refreshed or not CredentialSource().refresh(auth_header):
+                        raise ConnectionError(_LOGIN_MESSAGE) from err
+                    refreshed = True
+                    continue
+                if err.status_code == 403 or not reconnect or idle_timeout is not None:
+                    raise
             except (WebSocketException, OSError):
                 if not reconnect or idle_timeout is not None:
                     raise

python/lightning_sdk/api/sandbox_api.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 from lightning_sdk.lightning_cloud import env as lightning_env
-from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login
 from lightning_sdk.lightning_cloud.openapi import JobsServiceApi, SandboxesServiceApi
 from lightning_sdk.lightning_cloud.openapi.models import (
     SandboxesServiceCreateSandboxDirectoryBody,
@@ -17,7 +17,7 @@
     V1SandboxSnapshot,
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
-from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.lightning_cloud.rest_client import LightningClient, credential_request_wrapper
 
 
 @dataclass
@@ -228,10 +228,14 @@ def _apply_auth_and_host(self) -> None:
     def _ensure_auth(self) -> None:
         if self._auth_configured:
             return
-        auth_header = Auth().authenticate()
-        if not auth_header:
-            raise RuntimeError("Missing credentials: set api_key via configure() or use lightning login / env vars.")
-        self._client.api_client.set_default_header("Authorization", auth_header)
+        source = CredentialSource()
+        wrapped = credential_request_wrapper(self._client.api_client.request, source)
+
+        def request(*args: Any, **kwargs: Any) -> Any:
+            with browser_login(False):
+                return wrapped(*args, **kwargs)
+
+        self._client.api_client.request = request
         self._auth_configured = True
 
     @property
@@ -243,11 +247,14 @@ def host(self) -> str:
         return str(self._client.api_client.configuration.host).rstrip("/")
 
     def auth_header(self) -> str:
-        """Return the configured Authorization header, authenticating lazily when needed."""
+        """Return the Authorization header the next request would send."""
         self._ensure_auth()
-        headers = getattr(self._client.api_client, "default_headers", {})
-        auth = headers.get("Authorization", "") if isinstance(headers, dict) else ""
-        return str(auth or "")
+        if self._config.get("api_key"):
+            headers = getattr(self._client.api_client, "default_headers", {})
+            auth = headers.get("Authorization", "") if isinstance(headers, dict) else ""
+            return str(auth or "")
+        with browser_login(False):
+            return CredentialSource().header()
 
     def sandboxes(self) -> SandboxesServiceApi:
         self._ensure_auth()