feat(mlflow): add StaticTokenAuth for HTTP bearer-token servers

src/anemoi/utils/mlflow/auth.py

@@ -70,6 +70,7 @@ def __bool__(self) -> bool:
 class ServerConfig(BaseModel):
     refresh_token: str | None = None
     refresh_expires: int = 0
+    static_token: str | None = None
 
     @field_validator("refresh_expires", mode="before")
     def to_int(cls, value: float | int) -> int:
@@ -126,6 +127,8 @@ def normalise_urls(self) -> ServerStore:
 class AuthBase(ABC):
     """Base class for authentication implementations."""
 
+    _enabled: bool = False
+
     @abstractmethod
     def __init__(self, *args, **kwargs):
         pass
@@ -166,6 +169,18 @@ def user_info(self) -> UserInfo:
         return UserInfo()
 
 
+def enabled_guard(fn: Callable) -> Callable:
+    """Decorator to call or ignore a method based on the instance's `_enabled` flag."""
+
+    @wraps(fn)
+    def _wrapper(self: AuthBase, *args, **kwargs) -> Callable | None:
+        if self._enabled:
+            return fn(self, *args, **kwargs)
+        return None
+
+    return _wrapper
+
+
 class TokenAuth(AuthBase):
     """Manage authentication with a keycloak token server."""
 
@@ -270,22 +285,12 @@ def load_config() -> dict:
         last = {}
         for url, cfg in store.items():
             if cfg.refresh_expires > last.get("refresh_expires", 0):
-                last = dict(url=url, **cfg.model_dump())
+                # exclude fields added after this deprecated API to preserve its contract
+                last = dict(url=url, **cfg.model_dump(exclude={"static_token"}))
 
         return last
 
-    def enabled(fn: Callable) -> Callable:  # noqa: N805
-        """Decorator to call or ignore a function based on the `enabled` flag."""
-
-        @wraps(fn)
-        def _wrapper(self: TokenAuth, *args, **kwargs) -> Callable | None:
-            if self._enabled:
-                return fn(self, *args, **kwargs)
-            return None
-
-        return _wrapper
-
-    @enabled
+    @enabled_guard
     def login(self, force_credentials: bool = False, **kwargs: dict) -> None:
         """Acquire a new refresh token and save it to disk.
 
@@ -333,7 +338,7 @@ def login(self, force_credentials: bool = False, **kwargs: dict) -> None:
 
         self.log.info("✅ Successfully logged in to MLflow. Happy logging!")
 
-    @enabled
+    @enabled_guard
     def authenticate(self, **kwargs: dict) -> None:
         """Check the access token and refresh it if necessary. A new refresh token will also be acquired upon refresh.
 
@@ -367,7 +372,7 @@ def authenticate(self, **kwargs: dict) -> None:
 
         os.environ[self.target_env_var] = self.access_token
 
-    @enabled
+    @enabled_guard
     def save(self, **kwargs: dict) -> None:
         """Save the latest refresh token to disk."""
         del kwargs  # unused
@@ -445,3 +450,141 @@ def _request(self, path: str, payload: dict) -> dict:
         except HTTPError:
             self.log.exception("HTTP error occurred")
             raise
+
+
+class StaticTokenAuth(AuthBase):
+    """Authentication with a static, pre-issued bearer token.
+
+    Unlike `TokenAuth`, this class does not perform any refresh flow or talk to a
+    token server. It simply stores the supplied token in the target environment
+    variable (by default `MLFLOW_TRACKING_TOKEN`) so that MLflow sends it as an
+    `Authorization: Bearer <token>` header on every request.
+
+    The token is persisted to (and loaded from) the same on-disk store used by
+    `TokenAuth` (`~/.anemoi/mlflow-token.json`), keyed by server URL. This means a
+    token supplied once via `login()` is remembered across sessions, and multiple
+    servers can each hold their own static token.
+
+    This is useful for MLflow servers that accept a long-lived HTTP access token.
+
+    Note
+    ----
+    If the server sits behind a proxy/gateway that enforces per-method
+    permissions, the token must allow the HTTP methods the MLflow client uses.
+    The client issues several read operations as ``POST``, so a read-only token
+    may be insufficient.
+    """
+
+    _config_file = TokenAuth._config_file
+
+    def __init__(
+        self,
+        url: str | None,
+        token: str | None = None,
+        enabled: bool = True,
+        target_env_var: str = "MLFLOW_TRACKING_TOKEN",
+    ) -> None:
+        """Initialise the static token authentication object.
+
+        Parameters
+        ----------
+        url : str | None
+            URL of the MLflow server the token belongs to. Used as the storage key.
+        token : str | None, optional
+            The static bearer token to use. If not provided, a previously saved token
+            for this URL is loaded from disk. By default None.
+        enabled : bool, optional
+            Set this to False to turn off authentication, by default True.
+        target_env_var : str, optional
+            The environment variable to store the access token in,
+            by default `MLFLOW_TRACKING_TOKEN`.
+
+        """
+        self.access_token = token
+        self._enabled = enabled
+        self.target_env_var = target_env_var
+        self.log = logging.getLogger(__name__)
+
+        if not url:
+            self.url = None
+            assert not enabled, "URL must be provided if authentication is enabled."
+            return
+
+        self.url = url.rstrip("/")
+
+        # load a previously saved token if none was supplied
+        if self.access_token is None:
+            config = TokenAuth._get_store().get(self.url)
+            if config is not None:
+                self.access_token = config.static_token
+
+    def __call__(self) -> None:
+        self.authenticate()
+
+    @enabled_guard
+    def save(self, **kwargs: dict) -> None:
+        """Persist the static token to disk, keyed by server URL."""
+        del kwargs  # unused
+        if not self.access_token:
+            self.log.warning("No token to save.")
+            return
+
+        with CONFIG_LOCK:
+            store = TokenAuth._get_store()
+            existing = store.get(self.url)
+            config = existing.model_copy() if existing is not None else ServerConfig()
+            config.static_token = self.access_token
+            store.update(self.url, config)
+            save_config(self._config_file, store.model_dump())
+
+    @enabled_guard
+    def login(self, force_credentials: bool = False, **kwargs: dict) -> None:
+        """Acquire a static token and save it to disk.
+
+        If a token was supplied at construction (or a valid one is already on disk) it
+        is reused. Otherwise, or when `force_credentials` is set, the user is prompted
+        to paste one interactively.
+
+        Parameters
+        ----------
+        force_credentials : bool, optional
+            Force a prompt for a new token even if one is available, by default False.
+        kwargs : dict
+            Additional keyword arguments (unused).
+
+        """
+        del kwargs  # unused
+        self.log.info("🌐 Using static token authentication for %s", self.url)
+
+        if force_credentials or not self.access_token:
+            self.log.info("📝 Please paste your MLflow access token")
+            self.log.info("📝 (you will not see the output, just press enter after pasting):")
+            self.access_token = getpass("Access Token: ")
+
+        if not self.access_token:
+            msg = "❌ No token provided. Please try again."
+            raise RuntimeError(msg)
+
+        self.save()
+        self.authenticate()
+        self.log.info("✅ Static token stored. Happy logging!")
+
+    @enabled_guard
+    def authenticate(self, **kwargs: dict) -> None:
+        """Store the static token in the target environment variable."""
+        del kwargs  # unused
+        if not self.access_token:
+            msg = "You are not logged in to MLflow. Please log in first."
+            raise RuntimeError(msg)
+        os.environ[self.target_env_var] = self.access_token
+
+    def user_info(self) -> UserInfo:
+        """Get user info embedded in the token, if it is a JWT."""
+        if not self._enabled or not self.access_token:
+            return UserInfo()
+
+        try:
+            return UserInfo.from_jwt(self.access_token)
+        except Exception as e:
+            self.log.exception("Failed to decode access token.", exc_info=e)
+            return UserInfo()

src/anemoi/utils/mlflow/client.py

@@ -19,6 +19,7 @@
         "The `mlflow` package is required to use AnemoiMLflowclient. Please install it with `pip install mlflow`."
     )
 
+from .auth import StaticTokenAuth
 from .auth import TokenAuth
 from .utils import health_check
 
@@ -31,6 +32,7 @@ def __init__(
         tracking_uri: str,
         *args,
         authentication: bool = False,
+        static_token: str | bool | None = None,
         check_health: bool = True,
         **kwargs,
     ) -> None:
@@ -42,6 +44,12 @@ def __init__(
             The URI of the MLflow tracking server.
         authentication : bool, optional
             Enable token authentication, by default False
+        static_token : str | bool | None, optional
+            Use a static, pre-issued bearer token instead of the interactive refresh-token flow.
+            Pass the token string to use it directly, or pass ``True`` to load a previously
+            saved static token for this server from disk (``~/.anemoi/mlflow-token.json``).
+            When set, the token is stored/loaded via `StaticTokenAuth`. Requires
+            `authentication=True` to take effect. By default None (use `TokenAuth`).
         check_health : bool, optional
             Check the health of the MLflow server on init, by default True
         *args : Any
@@ -50,7 +58,11 @@ def __init__(
             Additional keyword arguments to pass to the MLflow client.
 
         """
-        self.anemoi_auth = TokenAuth(tracking_uri, enabled=authentication)
+        if static_token is not None and static_token is not False:
+            token = static_token if isinstance(static_token, str) else None
+            self.anemoi_auth = StaticTokenAuth(tracking_uri, token=token, enabled=authentication)
+        else:
+            self.anemoi_auth = TokenAuth(tracking_uri, enabled=authentication)
         if check_health:
             super().__getattribute__("anemoi_auth").authenticate()
             health_check(tracking_uri)

tests/test_mlflow_auth.py

@@ -19,6 +19,7 @@
 from anemoi.utils.mlflow.auth import NoAuth
 from anemoi.utils.mlflow.auth import ServerConfig
 from anemoi.utils.mlflow.auth import ServerStore
+from anemoi.utils.mlflow.auth import StaticTokenAuth
 from anemoi.utils.mlflow.auth import TokenAuth
 from anemoi.utils.mlflow.auth import UserInfo
 
@@ -198,6 +199,7 @@ def test_legacy_format(mocker: pytest.MockerFixture) -> None:
         legacy_config["url"]: {
             "refresh_token": legacy_config["refresh_token"],
             "refresh_expires": legacy_config["refresh_expires"],
+            "static_token": None,
         }
     }
     mocker.patch(
@@ -349,3 +351,116 @@ def test_user_info(mocker: pytest.MockerFixture) -> None:
         assert info.name is None
         assert info.email is None
         assert info.username is None
+
+
+def static_mocks(mocker: pytest.MockerFixture, config: dict | None = None) -> pytest.Mock:
+    """Patch the on-disk store helpers for StaticTokenAuth tests."""
+    mocker.patch(
+        "anemoi.utils.mlflow.auth.load_raw_config",
+        return_value=config or {},
+    )
+    save = mocker.patch("anemoi.utils.mlflow.auth.save_config")
+    mocker.patch("os.environ")
+    return save
+
+
+def test_static_token_direct(mocker: pytest.MockerFixture) -> None:
+    save = static_mocks(mocker)
+    auth = StaticTokenAuth("https://test.url", token="my-static-token")
+
+    assert auth.access_token == "my-static-token"  # noqa: S105
+
+    auth.authenticate()
+    os.environ.__setitem__.assert_called_once_with("MLFLOW_TRACKING_TOKEN", "my-static-token")
+
+    # a direct token is not saved until login/save is called
+    save.assert_not_called()
+
+
+def test_static_token_loaded_from_disk(mocker: pytest.MockerFixture) -> None:
+    config = {"https://test.url": {"static_token": "saved-token"}}
+    static_mocks(mocker, config=config)
+
+    auth = StaticTokenAuth("https://test.url")
+    assert auth.access_token == "saved-token"  # noqa: S105
+
+
+def test_static_token_save(mocker: pytest.MockerFixture) -> None:
+    save = static_mocks(mocker)
+    auth = StaticTokenAuth("https://test.url", token="my-static-token")
+    auth.save()
+
+    _, saved = save.call_args.args
+    assert saved["https://test.url"]["static_token"] == "my-static-token"
+
+
+def test_static_token_save_preserves_refresh_token(mocker: pytest.MockerFixture) -> None:
+    # a server that already has a refresh token should keep it when a static token is added
+    config = {"https://test.url": {"refresh_token": "keep-me", "refresh_expires": 123}}
+    save = static_mocks(mocker, config=config)
+
+    auth = StaticTokenAuth("https://test.url", token="my-static-token")
+    auth.save()
+
+    _, saved = save.call_args.args
+    assert saved["https://test.url"]["static_token"] == "my-static-token"
+    assert saved["https://test.url"]["refresh_token"] == "keep-me"
+
+
+def test_static_token_login_prompts_when_missing(mocker: pytest.MockerFixture) -> None:
+    save = static_mocks(mocker)
+    mocker.patch("anemoi.utils.mlflow.auth.getpass", return_value="pasted-token")
+
+    auth = StaticTokenAuth("https://test.url")
+    auth.login()
+
+    assert auth.access_token == "pasted-token"  # noqa: S105
+    _, saved = save.call_args.args
+    assert saved["https://test.url"]["static_token"] == "pasted-token"
+    os.environ.__setitem__.assert_called_once_with("MLFLOW_TRACKING_TOKEN", "pasted-token")
+
+
+def test_static_token_login_force_credentials(mocker: pytest.MockerFixture) -> None:
+    static_mocks(mocker)
+    mocker.patch("anemoi.utils.mlflow.auth.getpass", return_value="new-token")
+
+    auth = StaticTokenAuth("https://test.url", token="old-token")
+    auth.login(force_credentials=True)
+
+    assert auth.access_token == "new-token"  # noqa: S105
+
+
+def test_static_token_not_logged_in(mocker: pytest.MockerFixture) -> None:
+    static_mocks(mocker)
+    auth = StaticTokenAuth("https://test.url")
+    pytest.raises(RuntimeError, auth.authenticate)
+
+
+def test_static_token_disabled(mocker: pytest.MockerFixture) -> None:
+    static_mocks(mocker)
+    auth = StaticTokenAuth("https://test.url", token="my-static-token", enabled=False)
+    auth.authenticate()
+    os.environ.__setitem__.assert_not_called()
+
+
+def test_static_token_target_env_var(mocker: pytest.MockerFixture) -> None:
+    static_mocks(mocker)
+    auth = StaticTokenAuth(
+        "https://test.url",
+        token="my-static-token",
+        target_env_var="MLFLOW_TEST_ENV_VAR",
+    )
+    auth.authenticate()
+    os.environ.__setitem__.assert_called_once_with("MLFLOW_TEST_ENV_VAR", "my-static-token")
+
+
+def test_static_token_user_info(mocker: pytest.MockerFixture) -> None:
+    static_mocks(mocker)
+    payload = {"name": "John Anemoi", "preferred_username": "anemoi1234", "email": "john@example.com"}
+    token = f"e30.{base64.b64encode(json.dumps(payload).encode()).decode()}.e30"
+
+    auth = StaticTokenAuth("https://test.url", token=token)
+    info = auth.user_info()
+    assert info
+    assert info.name == "John Anemoi"
+    assert info.username == "anemoi1234"