From 08a32336e4f651e00d91986677667fcb00050572 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 25 Jun 2026 18:19:51 +0100
Subject: [PATCH 1/2] fix: Update rest session manager when logging in

If the user is not logged in when the client is created, the rest client
does not get a session manager. When the user calls login, a new session
manager is created to cache the new token but if the rest client is not
updated, any subsequent calls remain unauthenticated.

Setting the rest client's session manager to the new instance when a
login is successful means the tokens can be used without recreating the
client.

Now that the login method works as expected, the client created in the
`just repl` task can now automatically refresh its login if required.
---
 justfile                     | 1 +
 src/blueapi/client/client.py | 6 +++++-
 src/blueapi/client/rest.py   | 6 +++---
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/justfile b/justfile
index 2627d12e6..f6f31a1ac 100644
--- a/justfile
+++ b/justfile
@@ -40,5 +40,6 @@ repl:
     uv run --with ptpython ptpython -i <(cat << EOF
     from blueapi.client import BlueapiClient
     bc = BlueapiClient.from_config_file("tests/system_tests/config.yaml").with_instrument_session("cm12345-1")
+    bc.login()
     EOF
     )
diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 20de15892..ff68c574d 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -770,7 +770,11 @@ def login(self, token_path: Path | None = None):
                 auth = SessionManager(
                     oidc, cache_manager=SessionCacheManager(token_path)
                 )
-                auth.start_device_flow()
+                try:
+                    auth.start_device_flow()
+                    self._rest.session_manager = auth
+                except KeyboardInterrupt:
+                    pass
             else:
                 print("Server is not configured to use authentication!")
 
diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 0bddb5c87..980825708 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -186,7 +186,7 @@ def _response_json(response: requests.Response) -> Any:
 
 class BlueapiRestClient:
     _config: RestConfig
-    _session_manager: SessionManager | None
+    session_manager: SessionManager | None
     _pool: requests.Session
 
     def __init__(
@@ -195,7 +195,7 @@ def __init__(
         session_manager: SessionManager | None = None,
     ) -> None:
         self._config = config or RestConfig()
-        self._session_manager = session_manager
+        self.session_manager = session_manager
         self._pool = requests.Session()
 
     def get_plans(self) -> PlanResponse:
@@ -315,7 +315,7 @@ def _request_and_deserialize(
                 json=data,
                 params=params,
                 headers=carr,
-                auth=JWTAuth(self._session_manager),
+                auth=JWTAuth(self.session_manager),
             )
         except requests.exceptions.ConnectionError as ce:
             raise ServiceUnavailableError() from ce

From f451651f149ac063ec340ca59a161c2a01ea0213 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 26 Jun 2026 10:27:57 +0100
Subject: [PATCH 2/2] Move error handling to startup commands

---
 justfile                     | 12 +++++++++++-
 src/blueapi/client/client.py |  7 ++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/justfile b/justfile
index f6f31a1ac..0cc7222f5 100644
--- a/justfile
+++ b/justfile
@@ -39,7 +39,17 @@ repl:
     #!/usr/bin/env bash
     uv run --with ptpython ptpython -i <(cat << EOF
     from blueapi.client import BlueapiClient
+    from blueapi.client.rest import ServiceUnavailableError
     bc = BlueapiClient.from_config_file("tests/system_tests/config.yaml").with_instrument_session("cm12345-1")
-    bc.login()
+    try:
+        bc.login()
+    except KeyboardInterrupt:
+        print("Login cancelled")
+    except ServiceUnavailableError:
+        print("Couldn't access blueapi server to log in")
+    except Exception as e:
+        import traceback
+        print("Couldn't log in")
+        traceback.print_exception(e, chain=False)
     EOF
     )
diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index ff68c574d..ce5ec3001 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -770,11 +770,8 @@ def login(self, token_path: Path | None = None):
                 auth = SessionManager(
                     oidc, cache_manager=SessionCacheManager(token_path)
                 )
-                try:
-                    auth.start_device_flow()
-                    self._rest.session_manager = auth
-                except KeyboardInterrupt:
-                    pass
+                auth.start_device_flow()
+                self._rest.session_manager = auth
             else:
                 print("Server is not configured to use authentication!")