fix: blueapi login/logout when using multiple blueapis

docs/how-to/1.2 authn-authz.md

@@ -14,7 +14,9 @@ Here is an example configuration for authenticating to p46-blueapi:
 :language: yaml
 ```
 
-- **auth_token_path**: (Optional) Specify where to save the token. If omitted, the default is `~/.cache/blueapi_cache` or `$XDG_CACHE_HOME/blueapi_cache` if `XDG_CACHE_HOME` is set.
+- **auth_token_path**: (Optional) Specify where to save the token. If omitted, a separate cache is kept for each server, named after its host (and port, if given): e.g. `~/.cache/blueapi_cache.i22-blueapi.diamond.ac.uk` or `~/.cache/blueapi_cache.localhost_8000`. `$XDG_CACHE_HOME` is used instead of `~/.cache` if set.
+
+  Access tokens are refreshed shortly before they expire. If the authentication server cannot be reached the login is kept and an error explains what happened; the cache is only removed once the authentication server rejects the refresh token, at which point you need to log in again.
 
 ---
 

src/blueapi/cli/cli.py

@@ -31,7 +31,11 @@
     UnauthorisedAccessError,
     UnknownPlanError,
 )
-from blueapi.client.session import SessionCacheManager, SessionManager
+from blueapi.client.session import (
+    AuthenticationServerUnavailableError,
+    SessionManager,
+    default_cache_path,
+)
 from blueapi.config import (
     ApplicationConfig,
     ConfigLoader,
@@ -291,6 +295,8 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             raise ClickException(
                 "Access denied. Please check your login status and try again."
             ) from e
+        except AuthenticationServerUnavailableError as e:
+            raise ClickException(str(e)) from e
 
     return wrapper
 
@@ -604,23 +610,11 @@ def login(obj: dict) -> None:
     Authenticate with the blueapi using the OIDC (OpenID Connect) flow.
     """
     config: ApplicationConfig = obj["config"]
+    client = BlueapiClient.from_config(config)
     try:
-        auth: SessionManager = SessionManager.from_cache(config.auth_token_path)
-        access_token = auth.get_valid_access_token()
-        assert access_token
-        print("Logged in")
-    except Exception:
-        client = BlueapiClient.from_config(config)
-        try:
-            if oidc := client.oidc_config:
-                auth = SessionManager(
-                    oidc, cache_manager=SessionCacheManager(config.auth_token_path)
-                )
-                auth.start_device_flow()
-            else:
-                print("Server is not configured to use authentication!")
-        except NonJsonResponseError as e:
-            print(str(e))
+        client.login(config.auth_token_path)
+    except NonJsonResponseError as e:
+        print(str(e))
 
 
 @main.command(name="logout")
@@ -631,7 +625,13 @@ def logout(obj: dict) -> None:
     """
     config: ApplicationConfig = obj["config"]
     try:
-        auth: SessionManager = SessionManager.from_cache(config.auth_token_path)
+        if config.auth_token_path is None and (legacy := default_cache_path()).exists():
+            # Login cached by an older blueapi at the shared default path; log it
+            # out too so its refresh token is revoked rather than orphaned
+            SessionManager.from_cache(legacy).logout()
+        auth: SessionManager = SessionManager.from_cache(
+            config.auth_token_path, str(config.api.url)
+        )
         auth.logout()
     except FileNotFoundError:
         print("Logged out")

src/blueapi/client/client.py

@@ -13,7 +13,11 @@
 from observability_utils.tracing import get_tracer, start_as_current_span
 from pydantic import HttpUrl
 
-from blueapi.client.session import SessionCacheManager, SessionManager
+from blueapi.client.session import (
+    SessionCacheManager,
+    SessionManager,
+    migrate_legacy_cache,
+)
 from blueapi.config import (
     ApplicationConfig,
     ConfigLoader,
@@ -90,7 +94,9 @@ def from_config_file(cls, config_file: str) -> Self:
     def from_config(cls, config: ApplicationConfig) -> Self:
         session_manager: SessionManager | None = None
         try:
-            session_manager = SessionManager.from_cache(config.auth_token_path)
+            session_manager = SessionManager.from_cache(
+                config.auth_token_path, str(config.api.url)
+            )
         except Exception:
             ...  # Swallow exceptions
         rest = BlueapiRestClient(config.api, session_manager=session_manager)
@@ -657,20 +663,40 @@ def get_python_env(
         return self._rest.get_python_environment(name=name, source=source)
 
     def login(self, token_path: Path | None = None):
+        """
+        Log in to the server's OIDC provider, reusing a cached login if it is
+        still valid and was made with the same OIDC config the server uses now.
+
+        Raises:
+            AuthenticationServerUnavailableError: If the cached login needs
+                refreshing but the OIDC provider cannot be reached.
+        """
+        oidc = self.oidc_config
+        if oidc is None:
+            print("Server is not configured to use authentication!")
+            return
+        server_url = str(self._rest.url)
+        if token_path is None:
+            migrate_legacy_cache(server_url, oidc)
         try:
-            auth: SessionManager = SessionManager.from_cache(token_path)
-            access_token = auth.get_valid_access_token()
-            assert access_token
-            print("Logged in")
+            auth = SessionManager.from_cache(token_path, server_url)
         except Exception:
-            if oidc := self.oidc_config:
-                auth = SessionManager(
-                    oidc, cache_manager=SessionCacheManager(token_path)
+            auth = None
+        if auth is not None:
+            if not auth.is_for(oidc):
+                print(
+                    "The server's authentication config has changed since "
+                    "you logged in, logging in again"
                 )
-                auth.start_device_flow()
+            elif auth.get_valid_access_token():
+                print("Logged in")
                 self._rest.session_manager = auth
-            else:
-                print("Server is not configured to use authentication!")
+                return
+        auth = SessionManager(
+            oidc, cache_manager=SessionCacheManager(token_path, server_url)
+        )
+        auth.start_device_flow()
+        self._rest.session_manager = auth
 
     def logout(self):
         if sm := self._rest.session_manager:

src/blueapi/client/rest.py

@@ -10,14 +10,14 @@
     get_tracer,
     start_as_current_span,
 )
-from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic import BaseModel, HttpUrl, TypeAdapter, ValidationError
 from pydantic_core import PydanticSerializationError
 from websockets.exceptions import InvalidStatus
-from websockets.sync.client import connect
+from websockets.sync.client import ClientConnection, connect
 
 from blueapi import __version__
 from blueapi.client.protocols import ClientObjectRef
-from blueapi.client.session import JWTAuth, SessionManager
+from blueapi.client.session import SessionManager
 from blueapi.config import RestConfig
 from blueapi.core.bluesky_types import DataEvent
 from blueapi.service.model import (
@@ -201,6 +201,24 @@ def _response_json(response: requests.Response) -> Any:
         ) from exc
 
 
+def _access_token(session_manager: SessionManager | None) -> str:
+    return session_manager.get_valid_access_token() if session_manager else ""
+
+
+def _replacement_token(session_manager: SessionManager | None, rejected: str) -> str:
+    """A different token to retry with after `rejected` was refused, or """ ""
+    if not (session_manager and rejected):
+        return ""
+    token = session_manager.get_valid_access_token(rejected_token=rejected)
+    return token if token != rejected else ""
+
+
+def _with_bearer(headers: Mapping[str, str], token: str) -> dict[str, str]:
+    if not token:
+        return dict(headers)
+    return {**headers, "Authorization": f"Bearer {token}"}
+
+
 class BlueapiRestClient:
     _config: RestConfig
     session_manager: SessionManager | None
@@ -215,6 +233,10 @@ def __init__(
         self.session_manager = session_manager
         self._pool = requests.Session()
 
+    @property
+    def url(self) -> HttpUrl:
+        return self._config.url
+
     def get_plans(self) -> PlanResponse:
         return self._request_and_deserialize("/plans", PlanResponse)
 
@@ -298,7 +320,11 @@ def delete_environment(self) -> EnvironmentResponse:
 
     def get_oidc_config(self) -> OIDCConfig | None:
         try:
-            return self._request_and_deserialize("/config/oidc", OIDCConfig)
+            # Unauthenticated: a stale cached login (possibly for an OIDC provider
+            # the server no longer uses) must not stop us finding the new config
+            return self._request_and_deserialize(
+                "/config/oidc", OIDCConfig, authenticated=False
+            )
         except NoContentError:
             # Server is not using authentication
             return None
@@ -321,20 +347,31 @@ def _request_and_deserialize(
         method="GET",
         get_exception: Callable[[requests.Response], Exception | None] = _exception,
         params: Mapping[str, Any] | None = None,
+        authenticated: bool = True,
     ) -> T:
         url = self._config.url.unicode_string().removesuffix("/") + suffix
         # Get the trace context to propagate to the REST API
         headers = get_context_propagator()
         headers["User-Agent"] = USER_AGENT
-        try:
-            response = self._pool.request(
+        session_manager = self.session_manager if authenticated else None
+
+        def send(token: str) -> requests.Response:
+            return self._pool.request(
                 method,
                 url,
                 json=data,
                 params=params,
-                headers=headers,
-                auth=JWTAuth(self.session_manager),
+                headers=_with_bearer(headers, token),
             )
+
+        try:
+            token = _access_token(session_manager)
+            response = send(token)
+            if response.status_code == HTTPStatus.UNAUTHORIZED and (
+                new_token := _replacement_token(session_manager, token)
+            ):
+                # Token rejected (revoked, clock skew...): refresh and retry once
+                response = send(new_token)
         except requests.exceptions.ConnectionError as ce:
             raise ServiceUnavailableError() from ce
         exception = get_exception(response)
@@ -362,39 +399,55 @@ def run_blocking(
         self, req: TaskRequest
     ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]:
         url = self._config.ws_address.unicode_string().rstrip("/") + "/api/v2/run_plan"
+        with self._connect_websocket(url) as ws:
+            ws.send(Submit(task=req).model_dump_json(fallback=_task_model_fallback))
+            for message in ws:
+                event = ControlResponse.validate_json(message)
+                match event:
+                    case Update(data=data):
+                        yield data
+                    case InvalidArgs(errors=errors):
+                        raise InvalidParametersError(
+                            [
+                                ParameterError(
+                                    loc=e.loc, msg=e.msg, type=e.type, input=e.input
+                                )
+                                for e in errors
+                            ]
+                        )
+                    case PlanNotFound(plan_name=name):
+                        raise UnknownPlanError(message=name)
+                    case ServerBusy():
+                        raise BlueskyRemoteControlError(409, "Server is busy")
+                    case Unauthorized():
+                        raise UnauthorisedAccessError(
+                            403, "Not authorized to submit task"
+                        )
+
+    def _connect_websocket(self, url: str) -> ClientConnection:
+        """
+        Open the websocket, refreshing the token and retrying once if the
+        handshake is rejected as unauthenticated.
+        """
         headers = get_context_propagator()
-        if self.session_manager:
-            auth = self.session_manager.get_valid_access_token()
-            headers["Authorization"] = f"Bearer {auth}"
-        try:
-            with connect(
+
+        def open_with(token: str) -> ClientConnection:
+            return connect(
                 url,
-                additional_headers=headers,
+                additional_headers=_with_bearer(headers, token),
                 user_agent_header=USER_AGENT,
-            ) as ws:
-                ws.send(Submit(task=req).model_dump_json(fallback=_task_model_fallback))
-                for message in ws:
-                    event = ControlResponse.validate_json(message)
-                    match event:
-                        case Update(data=data):
-                            yield data
-                        case InvalidArgs(errors=errors):
-                            raise InvalidParametersError(
-                                [
-                                    ParameterError(
-                                        loc=e.loc, msg=e.msg, type=e.type, input=e.input
-                                    )
-                                    for e in errors
-                                ]
-                            )
-                        case PlanNotFound(plan_name=name):
-                            raise UnknownPlanError(message=name)
-                        case ServerBusy():
-                            raise BlueskyRemoteControlError(409, "Server is busy")
-                        case Unauthorized():
-                            raise UnauthorisedAccessError(
-                                403, "Not authorized to submit task"
-                            )
+            )
+
+        token = _access_token(self.session_manager)
+        try:
+            try:
+                return open_with(token)
+            except InvalidStatus as istat:
+                if istat.response.status_code == HTTPStatus.UNAUTHORIZED and (
+                    new_token := _replacement_token(self.session_manager, token)
+                ):
+                    return open_with(new_token)
+                raise
         except InvalidStatus as istat:
             match istat.response.status_code:
                 case 401 | 403: