Job-wait loop and login can be stopped in JupyterLite; a stale token re-authenticates [AI-written]

config.yml

@@ -66,7 +66,7 @@ notebooks:
       - ipywidgets
       - tabulate==0.9.0
       - mat3ra-utils
-      - mat3ra-api-client
+      - mat3ra-api-client>=2026.10.8.post0
   - name: dataframe
     packages_pyodide:
       - jinja2
@@ -79,7 +79,7 @@ notebooks:
       - emfs:/drive/packages/pydantic-2.7.1-py3-none-any.whl
       - mat3ra-standata
       - mat3ra-ide
-      - mat3ra-api-client
+      - mat3ra-api-client>=2026.10.8.post0
       - requests
       - mat3ra-ade
       - mat3ra-mode

pyproject.toml

@@ -24,7 +24,7 @@ utils_standata = [
   "mat3ra-standata",
 ]
 api = [
-  "mat3ra-api-client",
+  "mat3ra-api-client>=2026.10.8.post0",
 ]
 materials = [
   # ase >=3.25.0 is required for supercell to be generated,

src/py/mat3ra/notebooks_utils/api/job.py

@@ -1,29 +1,38 @@
+import asyncio
 import datetime
 from collections import Counter
-from typing import List
+from typing import Any, List
 
+import requests
 from mat3ra.api_client import JobEndpoints
 from mat3ra.utils.extra.tabulate import pretty_print
 
-from ..core.entity.job.api import create_job, get_jobs_statuses_by_ids, save_files, submit_jobs
+from ..core.entity.job.api import create_job, get_jobs_statuses_by_ids_async, save_files, submit_jobs
 from ..pyodide.runtime import interruptible_polling_loop
 
 
 # Contains no external dependencies, only uses the API client,
 # so can be used in both regular Python and Pyodide environments.
-# In regular Python, the interruptible_polling_loop does not actually do anything,
-# so the function behaves like a normal polling loop.
 @interruptible_polling_loop()
-def wait_for_jobs_to_finish_async(endpoint: JobEndpoints, job_ids: List[str]) -> bool:
+async def wait_for_jobs_to_finish_async(
+    endpoint: JobEndpoints, job_ids: List[str], *, abort_signal: Any = None
+) -> bool:
     """
     Waits for jobs to finish and prints their statuses.
     A job is considered finished if it is not in "pre-submission", "submitted", or "active" status.
 
     Args:
         endpoint (JobEndpoints): Job endpoint object from the Exabyte API Client
         job_ids (list): list of job IDs to wait for
+        abort_signal: JS AbortSignal of the Abort control, passed in by the polling loop
     """
-    statuses = get_jobs_statuses_by_ids(endpoint, job_ids)
+    try:
+        statuses = await get_jobs_statuses_by_ids_async(endpoint, job_ids, abort_signal=abort_signal)
+    except (asyncio.TimeoutError, OSError) as error:
+        if isinstance(error, requests.HTTPError) and error.response.status_code < 500:
+            raise
+        print(f"status check failed: {error!r}, retrying")
+        return True
     counts = Counter(statuses)
     headers = ["TIME", "SUBMITTED-JOBS", "ACTIVE-JOBS", "FINISHED-JOBS", "ERRORED-JOBS"]
     now = datetime.datetime.now().strftime("%Y-%m-%d-%H:%M:%S")
@@ -42,7 +51,6 @@ def wait_for_jobs_to_finish_async(endpoint: JobEndpoints, job_ids: List[str]) ->
 
 __all__ = [
     "create_job",
-    "get_jobs_statuses_by_ids",
     "save_files",
     "submit_jobs",
     "wait_for_jobs_to_finish_async",

src/py/mat3ra/notebooks_utils/auth.py

@@ -13,26 +13,37 @@
 import inspect
 import os
 
-from mat3ra.api_client import ACCESS_TOKEN_ENV_VAR
+import requests
+from mat3ra.api_client import ACCESS_TOKEN_ENV_VAR, APIClient, AuthContext
 
 from .core.api.auth import authenticate_oidc, get_oidc_base_url, store_token_data_in_environment
 from .io import get_data
 from .ipython.ui import show_device_flow_popup
 from .primitive.environment import is_pyodide_environment
 from .pyodide.api.auth import authenticate_jupyterlite
-from .token_store import load_token, save_token
+from .token_store import delete_token, load_token, save_token
 
 REFRESH_TOKEN_ENV_VAR = "OIDC_REFRESH_TOKEN"
 
 
 async def _authenticate_oidc_with_cache(force=False):
     oidc_url = get_oidc_base_url()
-    cached = None if force else await load_token(oidc_url)
-
-    if cached:
-        store_token_data_in_environment(cached)
-        return
+    access_token = os.environ.get(ACCESS_TOKEN_ENV_VAR)
+    token_data = {"access_token": access_token} if access_token else await load_token(oidc_url)
 
+    if token_data and not force:
+        try:
+            APIClient.authenticate(access_token=token_data["access_token"]).list_accounts()
+            store_token_data_in_environment(token_data)
+            return
+        except KeyError:
+            pass
+        except requests.HTTPError as error:
+            if error.response.status_code != 401:
+                raise
+
+    await delete_token(oidc_url)
+    os.environ.pop(ACCESS_TOKEN_ENV_VAR, None)
     token_data = await authenticate_oidc(show_popup=show_device_flow_popup)
     await save_token(oidc_url, token_data)
 
@@ -61,5 +72,13 @@ async def authenticate(force=False, globals_dict=None):
 
     if data_from_host:
         await authenticate_jupyterlite(data_from_host)
-    elif ACCESS_TOKEN_ENV_VAR not in os.environ or force:
+    else:
         await _authenticate_oidc_with_cache(force)
+
+
+async def reauthenticate(auth_context: AuthContext) -> None:
+    """
+    Replaces an access token the platform rejected with a new one from the device login, set on `auth_context`.
+    """
+    await _authenticate_oidc_with_cache(force=True)
+    auth_context.access_token = os.environ[ACCESS_TOKEN_ENV_VAR]

src/py/mat3ra/notebooks_utils/core/api/auth.py

@@ -1,12 +1,23 @@
 import asyncio
 import os
 import time
-from typing import Callable, Optional
+import urllib.parse
+from typing import Any, Awaitable, Callable, Optional
 
 import requests
 from mat3ra.api_client import ACCESS_TOKEN_ENV_VAR, CLIENT_ID, SCOPE, APIEnv, build_oidc_base_url
 
+from ...primitive.environment import is_pyodide_environment
+from ...pyodide.runtime import run_interruptible_loop_async
+
+try:
+    from pyodide.http import pyfetch  # type: ignore
+except ImportError:
+    pyfetch = None
+
 REFRESH_TOKEN_ENV_VAR = "OIDC_REFRESH_TOKEN"
+TOKEN_REQUEST_TIMEOUT_SECONDS = 10
+FORM_HEADERS = {"Content-Type": "application/x-www-form-urlencoded"}
 
 
 def get_oidc_base_url() -> str:
@@ -50,32 +61,58 @@ def store_token_data_in_environment(token_data: dict) -> None:
         os.environ[REFRESH_TOKEN_ENV_VAR] = token_data["refresh_token"]
 
 
+def _request_token_data(token_url: str, form_data: dict) -> tuple:
+    response = requests.post(token_url, data=form_data, headers=FORM_HEADERS, timeout=TOKEN_REQUEST_TIMEOUT_SECONDS)
+    return response.status_code, response.json() if response.status_code < 500 else {}
+
+
+async def _request_token_data_with_fetch(token_url: str, form_data: dict, abort_signal: Any) -> tuple:
+    """
+    `_request_token_data` through the browser's fetch, which leaves the event loop free while the request is in flight.
+    """
+    body = urllib.parse.urlencode(form_data, doseq=True)
+    response = await pyfetch(token_url, method="POST", body=body, headers=FORM_HEADERS, signal=abort_signal)
+    return response.status, await response.json() if response.status < 500 else {}
+
+
 async def _poll_for_token_data(
     oidc_base_url: str,
     client_id: str,
     device_code: str,
     polling_interval_seconds: int,
     expires_in_seconds: int,
 ) -> dict:
+    """Polls for the token until the device login is confirmed; in pyodide, ESC stops it at once."""
+    token_url = f"{oidc_base_url}/token"
+    form_data = {
+        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+        "device_code": device_code,
+        "client_id": client_id,
+        "redirect_uris": [],
+        "response_types": [],
+        "token_endpoint_auth_method": "none",
+    }
     deadline_seconds = time.time() + expires_in_seconds
-    while time.time() < deadline_seconds:
-        token_response = requests.post(
-            f"{oidc_base_url}/token",
-            data={
-                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
-                "device_code": device_code,
-                "client_id": client_id,
-                "redirect_uris": [],
-                "response_types": [],
-                "token_endpoint_auth_method": "none",
-            },
-            headers={"Content-Type": "application/x-www-form-urlencoded"},
-            timeout=10,
-        )
-        if token_response.status_code == 200:
-            return token_response.json()
-        await asyncio.sleep(polling_interval_seconds)
-    raise Exception("Timeout waiting for authorization.")
+    token_data: dict = {}
+
+    def request_token_data(abort_signal: Any) -> Awaitable[tuple]:
+        if is_pyodide_environment():
+            return _request_token_data_with_fetch(token_url, form_data, abort_signal)
+        return asyncio.get_running_loop().run_in_executor(None, _request_token_data, token_url, form_data)
+
+    async def poll_step(abort_signal: Any) -> bool:
+        if time.time() >= deadline_seconds:
+            raise Exception("Timeout waiting for authorization.")
+        status, response_data = await asyncio.wait_for(request_token_data(abort_signal), TOKEN_REQUEST_TIMEOUT_SECONDS)
+        if status != 200 and status < 500 and response_data.get("error") not in ("authorization_pending", "slow_down"):
+            raise Exception(f"Device login failed ({status}): {response_data.get('error')}.")
+        token_data.update(response_data if status == 200 else {})
+        return not token_data
+
+    await run_interruptible_loop_async(
+        poll_step, polling_interval_seconds, show_button=False, abort_hint_text="Press ESC to cancel"
+    )
+    return token_data
 
 
 async def authenticate_oidc(

src/py/mat3ra/notebooks_utils/core/entity/job/api.py

@@ -1,10 +1,23 @@
+import asyncio
+import json
 import re
+import urllib.parse
 import urllib.request
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Awaitable, Dict, Iterable, List, Optional, Union
 
+import requests
 from mat3ra.api_client import APIClient, JobEndpoints
 
+from ....auth import reauthenticate
+from ....primitive.environment import is_pyodide_environment
+
+try:
+    from pyodide.http import pyfetch  # type: ignore
+except ImportError:
+    pyfetch = None
+
 MATERIALS_SET_ENTITY_CLASS = "Material"
+DEFAULT_STATUS_TIMEOUT_SECONDS = 30
 
 
 def save_files(job_id: str, job_endpoint: JobEndpoints, filename_on_cloud: str, filename_on_disk: str) -> None:
@@ -25,18 +38,55 @@ def save_files(job_id: str, job_endpoint: JobEndpoints, filename_on_cloud: str,
         outp.write(server_response.read())
 
 
-def get_jobs_statuses_by_ids(endpoint: JobEndpoints, job_ids: List[str]) -> List[str]:
+async def _list_jobs_with_fetch(endpoint: JobEndpoints, query: dict, projection: dict, abort_signal: Any) -> List[dict]:
+    """
+    `endpoint.list` through the browser's fetch, which leaves the event loop free while the request is in flight.
+    Raises `requests.HTTPError` on an error status.
+    """
+    parameters = urllib.parse.urlencode({"query": json.dumps(query), "projection": json.dumps(projection)})
+    url = urllib.parse.urljoin(endpoint.conn.preamble, f"{endpoint.name}?{parameters}")
+    response = await pyfetch(url, headers={**endpoint.headers, **endpoint.auth.get_headers()}, signal=abort_signal)
+    if not response.ok:
+        error_response = requests.Response()
+        error_response.status_code = response.status
+        raise requests.HTTPError(f"Error {response.status}.", response=error_response)
+    return (await response.json())["data"]
+
+
+async def get_jobs_statuses_by_ids_async(
+    endpoint: JobEndpoints,
+    job_ids: List[str],
+    timeout: float = DEFAULT_STATUS_TIMEOUT_SECONDS,
+    abort_signal: Any = None,
+) -> List[str]:
     """
-    Gets jobs statues by their IDs.
+    Gets jobs statuses by their IDs without blocking the event loop: through the browser's fetch in pyodide,
+    in a worker thread otherwise. A rejected access token (401) is replaced through the device login once.
 
     Args:
         endpoint (JobEndpoints): Job endpoint object from the Exabyte API Client
         job_ids (list): list of job IDs to get the status for
+        timeout (float): seconds to wait for the response before raising asyncio.TimeoutError
+        abort_signal: JS AbortSignal that aborts the fetch in pyodide
 
     Returns:
         list: list of job statuses
     """
-    jobs = endpoint.list({"_id": {"$in": job_ids}}, {"fields": {"status": 1}})
+    query = {"_id": {"$in": job_ids}}
+    projection = {"fields": {"status": 1}}
+
+    def request_jobs() -> Awaitable[List[dict]]:
+        if is_pyodide_environment():
+            return _list_jobs_with_fetch(endpoint, query, projection, abort_signal)
+        return asyncio.get_running_loop().run_in_executor(None, endpoint.list, query, projection)
+
+    try:
+        jobs = await asyncio.wait_for(request_jobs(), timeout)
+    except requests.HTTPError as error:
+        if error.response.status_code != 401 or not endpoint.auth.access_token:
+            raise
+        await reauthenticate(endpoint.auth)
+        jobs = await asyncio.wait_for(request_jobs(), timeout)
     return [job["status"] for job in jobs]