From 8e1d330f8231d93d74a426fd450f7845783a24c6 Mon Sep 17 00:00:00 2001
From: Jacob Williamson <jacob.williamson@diamond.ac.uk>
Date: Wed, 10 Jun 2026 12:56:39 +0000
Subject: [PATCH] Create TokenRetriever protocol to support UDC clients

---
 src/blueapi/client/client.py          |  2 +-
 src/blueapi/client/rest.py            | 10 +++++-----
 src/blueapi/service/authentication.py |  8 ++++++--
 tests/unit_tests/client/test_rest.py  |  2 +-
 4 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 32d589089..edfa3bfe7 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -233,7 +233,7 @@ def from_config(cls, config: ApplicationConfig) -> Self:
             session_manager = SessionManager.from_cache(config.auth_token_path)
         except Exception:
             ...  # Swallow exceptions
-        rest = BlueapiRestClient(config.api, session_manager=session_manager)
+        rest = BlueapiRestClient(config.api, token_retreiver=session_manager)
         if config.stomp.enabled:
             assert config.stomp.url.host is not None, "Stomp URL missing host"
             assert config.stomp.url.port is not None, "Stomp URL missing port"
diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 6caa0add2..8386f7c2e 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -14,7 +14,7 @@
 
 from blueapi import __version__
 from blueapi.config import RestConfig
-from blueapi.service.authentication import JWTAuth, SessionManager
+from blueapi.service.authentication import JWTAuth, TokenRetriever
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
@@ -157,16 +157,16 @@ def _response_json(response: requests.Response) -> Any:
 
 class BlueapiRestClient:
     _config: RestConfig
-    _session_manager: SessionManager | None
+    _token_retreiver: TokenRetriever | None
     _pool: requests.Session
 
     def __init__(
         self,
         config: RestConfig | None = None,
-        session_manager: SessionManager | None = None,
+        token_retreiver: TokenRetriever | None = None,
     ) -> None:
         self._config = config or RestConfig()
-        self._session_manager = session_manager
+        self._token_retreiver = token_retreiver
         self._pool = requests.Session()
 
     def get_plans(self) -> PlanResponse:
@@ -283,7 +283,7 @@ def _request_and_deserialize(
                 json=data,
                 params=params,
                 headers=carr,
-                auth=JWTAuth(self._session_manager),
+                auth=JWTAuth(self._token_retreiver),
             )
         except requests.exceptions.ConnectionError as ce:
             raise ServiceUnavailableError() from ce
diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py
index b107f7b2b..b5c40c3d9 100644
--- a/src/blueapi/service/authentication.py
+++ b/src/blueapi/service/authentication.py
@@ -9,7 +9,7 @@
 from functools import cached_property
 from http import HTTPStatus
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, Protocol, cast
 
 import httpx
 import jwt
@@ -24,6 +24,10 @@
 SCOPES = "openid offline_access"
 
 
+class TokenRetriever(Protocol):
+    def get_valid_access_token(self) -> str: ...
+
+
 class CacheManager(ABC):
     @abstractmethod
     def can_access_cache(self) -> bool: ...
@@ -238,7 +242,7 @@ def start_device_flow(self):
 
 
 class JWTAuth(AuthBase):
-    def __init__(self, session_manager: SessionManager | None):
+    def __init__(self, session_manager: TokenRetriever | None):
         self.token: str = (
             session_manager.get_valid_access_token() if session_manager else ""
         )
diff --git a/tests/unit_tests/client/test_rest.py b/tests/unit_tests/client/test_rest.py
index 19d912694..ca50c1584 100644
--- a/tests/unit_tests/client/test_rest.py
+++ b/tests/unit_tests/client/test_rest.py
@@ -43,7 +43,7 @@ def rest() -> BlueapiRestClient:
 @pytest.fixture
 def rest_with_auth(oidc_config: OIDCConfig, tmp_path) -> BlueapiRestClient:
     return BlueapiRestClient(
-        session_manager=SessionManager(
+        token_retreiver=SessionManager(
             server_config=oidc_config,
             cache_manager=SessionCacheManager(tmp_path / "blueapi_cache"),
         )