From 886bdcc0327a824461752e951d200a037fe1a6b5 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:25:11 +0100 Subject: [PATCH 01/30] Enhance SSE connection handling in `tellupdater` for compatibility and improved error handling --- src/aare/daq/tellupdater.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index e9815dec..449a0fcc 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -40,13 +40,17 @@ def listen_to_sse(): return sse_url = tell_client.url + "/events" try: - #response = requests.get(sse_url, stream=True) - client = sseclient.SSEClient(sse_url) + # Build a streaming HTTP response first to support common sseclient variants. + response = requests.get(sse_url, stream=True, timeout=(5, 60)) + response.raise_for_status() + client = sseclient.SSEClient(response) print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") handle_tell_change_event() - for event in client.events(): + # Compatibility: some SSEClient versions are iterable, others expose .events(). + events_iter = client.events() if hasattr(client, "events") else iter(client) + for event in events_iter: print(f"event = {event.event} with data: {event.data}") if event.event == "DewarContentUpdate": on_sse_event(event) @@ -164,4 +168,4 @@ def main(): time.sleep(5) if __name__ == "__main__": - main() \ No newline at end of file + main() From 68df2930b3c0ecd7ec33be5c97d8d30ec03c8498 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Fri, 13 Mar 2026 09:35:21 +0100 Subject: [PATCH 02/30] Streamline SSE client setup in `tellupdater` by removing redundant HTTP stream handling --- src/aare/daq/tellupdater.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 449a0fcc..defc8dd1 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -4,7 +4,6 @@ import threading import websocket import sseclient -import requests import time from aareDB.models import PuckWithTellPosition @@ -40,10 +39,8 @@ def listen_to_sse(): return sse_url = tell_client.url + "/events" try: - # Build a streaming HTTP response first to support common sseclient variants. - response = requests.get(sse_url, stream=True, timeout=(5, 60)) - response.raise_for_status() - client = sseclient.SSEClient(response) + # Use URL directly for sseclient variants that manage their own HTTP stream. + client = sseclient.SSEClient(sse_url) print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") handle_tell_change_event() From 6d0cf1fc4711e26c622c75963b0d7aef33eb462d Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 13 Mar 2026 09:38:11 +0100 Subject: [PATCH 03/30] Tell Client: temp fix of bug --- src/aare/devices/tell_client.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/aare/devices/tell_client.py b/src/aare/devices/tell_client.py index 9ff31f99..9ef0b5f5 100755 --- a/src/aare/devices/tell_client.py +++ b/src/aare/devices/tell_client.py @@ -633,6 +633,9 @@ class TellClientProxy: def url(self): return self._get_client().url + def __getattr__(self, name): + return getattr(self._get_client(), name) + # Delegate methods used by DAQ; add more as needed def get_mounted_sample(self) -> SampleDewarAddress | None: return self._get_client().get_mounted_sample() From 133a4a68766b14643bd947db7d057197220b4d87 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 13 Mar 2026 09:40:06 +0100 Subject: [PATCH 04/30] TellUpdater: remove unessacery print statement --- src/aare/daq/tellupdater.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index defc8dd1..5959fe45 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -48,7 +48,7 @@ def listen_to_sse(): # Compatibility: some SSEClient versions are iterable, others expose .events(). events_iter = client.events() if hasattr(client, "events") else iter(client) for event in events_iter: - print(f"event = {event.event} with data: {event.data}") + #print(f"event = {event.event} with data: {event.data}") if event.event == "DewarContentUpdate": on_sse_event(event) except Exception as exc: From 6a7e61c1a7abc98fc2c967766d8470593f32913f Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 13 Mar 2026 09:40:22 +0100 Subject: [PATCH 05/30] SmargonTrace: change location of csv to /sls/mx/applications/logs --- src/aare/daq/daq.py | 2 +- src/aare/gui/panels/smargon_trace_panel.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 9a72ec0e..db2a5d3b 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -62,7 +62,7 @@ class AareDAQ: self.__bl = bl.value.upper() self.__aare = AareWrapper(bl) self.__saved_box = None - self._smargon_trace_path = Path("logs") / "smargon_trace.csv" + self._smargon_trace_path = Path("/sls/mx/applications/logs") / "smargon_trace.csv" self._face_detection_progress_cb: Callable[[dict], None] | None = None self._last_sample_sync_ts = 0.0 self._sample_sync_min_interval_s = 2.0 diff --git a/src/aare/gui/panels/smargon_trace_panel.py b/src/aare/gui/panels/smargon_trace_panel.py index fe958702..7f031cc4 100644 --- a/src/aare/gui/panels/smargon_trace_panel.py +++ b/src/aare/gui/panels/smargon_trace_panel.py @@ -507,6 +507,7 @@ class SmargonTracePanel(QWidget): project_root / self._csv_path, project_root / "src" / "aare" / "daq" / "logs" / "smargon_trace.csv", project_root / "src" / "aare" / "gui" / "logs" / "smargon_trace.csv", + Path("/sls/mx/applications/logs/smargon_trace.csv"), ] out: list[Path] = [] From 3c769106bf8254edd0c9081561deca17f7c518a3 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 13 Mar 2026 09:40:44 +0100 Subject: [PATCH 06/30] Aerotech: think we are removing automation1 api, WIP --- src/aare/devices/aerotech.py | 290 +++++++++++++++++------------------ 1 file changed, 145 insertions(+), 145 deletions(-) diff --git a/src/aare/devices/aerotech.py b/src/aare/devices/aerotech.py index 5884e40b..bed65a09 100644 --- a/src/aare/devices/aerotech.py +++ b/src/aare/devices/aerotech.py @@ -5,7 +5,7 @@ import time from enum import Enum from typing import Union -import automation1 as a1 +#import automation1 as a1 from epics import PV, Motor, caput, caget from aare.common.beamline import MXBeamline, mx_beamline @@ -173,150 +173,150 @@ class AerotechControllerEpics: def set_global_variable(self, index: int, value: Union[int, float, str]): self.__set_global(index, value) - -class AerotechController: - def __init__(self, controller_ip: str): - if controller_ip is None: - self.controller = None - else: - self.controller = a1.Controller.connect(controller_ip) - - self.status_item_configuration = a1.StatusItemConfiguration() - self.start_controller() - - def start_controller(self): - self.controller.start() - - def disconnect(self): - self.controller.disconnect() - - def enable_motion(self, axis: str): - self.controller.runtime.commands.motion.enable(axis.upper()) - - def home_motor(self, axis: str): - self.__configure_axis_status(axis, a1.AxisStatusItem.AxisEnabled) - result = self.get_status_via_status_items(name=axis, status_item=a1.AxisStatusItem.AxisEnabled) - if not result == a1.AxisStatusItem.AxisEnabled: - self.enable_motion(axis) - self.controller.runtime.commands.motion.home(axis.upper()) - - def __configure_axis_status(self, axis_name: str, axis_status_item: a1.AxisStatusItem): - self.status_item_configuration.axis.add(axis_status_item=axis_status_item, axis=axis_name) - - def __configure_task_status(self, task_id:int, task_status_item: a1.TaskStatusItem): - self.status_item_configuration.task.add(task_status_item=task_status_item, task=f"Task {task_id}") - - def get_status_via_status_items(self, name: str | int, status_item: a1.TaskStatusItem | a1.AxisStatusItem): - result = self.controller.runtime.status.get_status_items(self.status_item_configuration) - if int(name): - return result.task.get(status_item, f"Task {name}").value - elif str(name): - return result.axis.get(status_item, name).value - else: - raise ValueError(f"Invalid status item {status_item} for task {name}") - - def set_global_variable(self, index:int, value: Union[int, float, str]): - if type(value) is int: - self.controller.runtime.variables.global_.set_integer(index, value) - elif type(value) is float: - self.controller.runtime.variables.global_.set_real(index, value) - elif type(value) is str: - self.controller.runtime.variables.global_.set_string(index, value) - else: - raise ValueError(f"Invalid type {type(value)} for global variable") - - def get_axis_status_via_status_items(self, axis_name:str, status_item: a1.AxisStatusItem): - result = self.controller.runtime.status.get_status_items(self.status_item_configuration) - return result.task.get(status_item, axis_name).value - - def wait_for_status_to_change(self, name: str | int, status_item:a1.TaskStatusItem | a1.AxisStatusItem, enum, timeout: float = 60.0): - start = time.perf_counter() - while self.get_status_via_status_items(name, status_item) == enum: - time.sleep(0.1) - if time.perf_counter() - start > timeout: - print(self.get_status_via_status_items(name, status_item)) - raise TimeoutError(f"Timeout waiting for task {name} to finish") - return - - def wait_program_finish(self, task_id:int=3, timeout:float=60.0): - start = time.perf_counter() - status_item = a1.TaskStatusItem.TaskState - while True: - status = self.get_status_via_status_items(task_id, status_item) - if status != a1.TaskState.ProgramRunning: - if status == a1.TaskState.Idle: - return - elif status == a1.TaskState.ProgramComplete: - print(f"Program completed on Task {task_id}") - return - elif status == a1.TaskState.Error: - raise RuntimeError(f"Task {task_id} failed to start") - elif status == a1.TaskState.ProgramPaused: - print(f"Program paused on Task {task_id}, not sure how") - else: - raise RuntimeError(f"Unknown status {status} for task {task_id}") - if time.perf_counter() - start > timeout: - print(self.get_status_via_status_items(task_id, status_item)) - raise TimeoutError(f"Timeout waiting for task {task_id} to finish") - time.sleep(0.05) - print(f"Task {task_id} finished with status {status}") - - def __run_program(self, script_name:str, task_id:int=3, timeout:float=60.0): - self.__configure_task_status(task_id, a1.TaskStatusItem.TaskState) - state = self.get_status_via_status_items(task_id, a1.TaskStatusItem.TaskState) - print(f"Task {task_id} is in state {state}: {a1.TaskState(state).name}") - if state != a1.TaskState.ProgramComplete and state != a1.TaskState.Idle: - if a1.TaskState.ProgramRunning == state: - self.wait_for_status_to_change(task_id, a1.TaskStatusItem.TaskState, state, timeout) - elif a1.TaskState.ProgramComplete == state: - print('can continue') - else: - print(f"Task {task_id} is not Idle: {a1.TaskState(state).name} , aborting") - return - try: - print(f"Running script {script_name} on task {task_id}") - self.controller.runtime.tasks[task_id].program.run(script_name) - self.wait_program_finish(task_id, timeout) - except Exception as e: - print(f"Error executing script {script_name}: {e}") - - def run_grid_scan(self, cell_height_mm:float, num_rows:int, - row_width_mm:float, time_per_row_s: float, task_id:int =3): - self.set_global_variable(0, cell_height_mm) - self.set_global_variable(1, row_width_mm) - self.set_global_variable(2, time_per_row_s) - self.set_global_variable(1, num_rows) - self.set_global_variable(0, 1) - #self.__run_program(script_name="grid_scan.a1exe", task_id=task_id, timeout=120.0) - def home_all(self, task_id:int = 3): - self.__run_program(script_name="home_all.a1exe", task_id=task_id, timeout=120.0) - - def rotation_scan(self, task_id = 3, start_angle:float = 0.0, end_angle:float = 360.0, step_size:float = 1.0, num_steps:int = 10): - self.__run_program(script_name="rotation_scan.a1exe", task_id=task_id, timeout=90.0) - - def move_motor_absolute(self, axis:str, position:float, speed:float=1.0): - self.controller.runtime.commands.motion.moveabsolute(axis.upper(), [position], [speed]) - - def move_motor_linear(self, axis: str, position: float, speed: float = 1.0): - self.controller.runtime.commands.motion.movelinear(axis.upper(), [position], speed) - -if __name__ == "__main__": - beamline = mx_beamline() - print(beamline) - ### test on 10S - aerotech = AerotechController(controller_ip="129.129.118.96") - #aerotech.enable_motion("X") - #aerotech.home_all() - rw = 0.320 - ch = 0.010 - nr = 10 - tpr_s = rw / ch * 0.02 - #aerotech.run_grid_scan(cell_height_mm=ch, num_rows=nr, - # row_width_mm=rw, time_per_row_s=tpr_s) - st = time.perf_counter() - aerotech.move_motor_absolute("Z", 0, 10000) - print(f"time to move: {time.perf_counter() - st}") - aerotech.disconnect() +# +# class AerotechController: +# def __init__(self, controller_ip: str): +# if controller_ip is None: +# self.controller = None +# else: +# self.controller = a1.Controller.connect(controller_ip) +# +# self.status_item_configuration = a1.StatusItemConfiguration() +# self.start_controller() +# +# def start_controller(self): +# self.controller.start() +# +# def disconnect(self): +# self.controller.disconnect() +# +# def enable_motion(self, axis: str): +# self.controller.runtime.commands.motion.enable(axis.upper()) +# +# def home_motor(self, axis: str): +# self.__configure_axis_status(axis, a1.AxisStatusItem.AxisEnabled) +# result = self.get_status_via_status_items(name=axis, status_item=a1.AxisStatusItem.AxisEnabled) +# if not result == a1.AxisStatusItem.AxisEnabled: +# self.enable_motion(axis) +# self.controller.runtime.commands.motion.home(axis.upper()) +# +# def __configure_axis_status(self, axis_name: str, axis_status_item: a1.AxisStatusItem): +# self.status_item_configuration.axis.add(axis_status_item=axis_status_item, axis=axis_name) +# +# def __configure_task_status(self, task_id:int, task_status_item: a1.TaskStatusItem): +# self.status_item_configuration.task.add(task_status_item=task_status_item, task=f"Task {task_id}") +# +# def get_status_via_status_items(self, name: str | int, status_item: a1.TaskStatusItem | a1.AxisStatusItem): +# result = self.controller.runtime.status.get_status_items(self.status_item_configuration) +# if int(name): +# return result.task.get(status_item, f"Task {name}").value +# elif str(name): +# return result.axis.get(status_item, name).value +# else: +# raise ValueError(f"Invalid status item {status_item} for task {name}") +# +# def set_global_variable(self, index:int, value: Union[int, float, str]): +# if type(value) is int: +# self.controller.runtime.variables.global_.set_integer(index, value) +# elif type(value) is float: +# self.controller.runtime.variables.global_.set_real(index, value) +# elif type(value) is str: +# self.controller.runtime.variables.global_.set_string(index, value) +# else: +# raise ValueError(f"Invalid type {type(value)} for global variable") +# +# def get_axis_status_via_status_items(self, axis_name:str, status_item: a1.AxisStatusItem): +# result = self.controller.runtime.status.get_status_items(self.status_item_configuration) +# return result.task.get(status_item, axis_name).value +# +# def wait_for_status_to_change(self, name: str | int, status_item:a1.TaskStatusItem | a1.AxisStatusItem, enum, timeout: float = 60.0): +# start = time.perf_counter() +# while self.get_status_via_status_items(name, status_item) == enum: +# time.sleep(0.1) +# if time.perf_counter() - start > timeout: +# print(self.get_status_via_status_items(name, status_item)) +# raise TimeoutError(f"Timeout waiting for task {name} to finish") +# return +# +# def wait_program_finish(self, task_id:int=3, timeout:float=60.0): +# start = time.perf_counter() +# status_item = a1.TaskStatusItem.TaskState +# while True: +# status = self.get_status_via_status_items(task_id, status_item) +# if status != a1.TaskState.ProgramRunning: +# if status == a1.TaskState.Idle: +# return +# elif status == a1.TaskState.ProgramComplete: +# print(f"Program completed on Task {task_id}") +# return +# elif status == a1.TaskState.Error: +# raise RuntimeError(f"Task {task_id} failed to start") +# elif status == a1.TaskState.ProgramPaused: +# print(f"Program paused on Task {task_id}, not sure how") +# else: +# raise RuntimeError(f"Unknown status {status} for task {task_id}") +# if time.perf_counter() - start > timeout: +# print(self.get_status_via_status_items(task_id, status_item)) +# raise TimeoutError(f"Timeout waiting for task {task_id} to finish") +# time.sleep(0.05) +# print(f"Task {task_id} finished with status {status}") +# +# def __run_program(self, script_name:str, task_id:int=3, timeout:float=60.0): +# self.__configure_task_status(task_id, a1.TaskStatusItem.TaskState) +# state = self.get_status_via_status_items(task_id, a1.TaskStatusItem.TaskState) +# print(f"Task {task_id} is in state {state}: {a1.TaskState(state).name}") +# if state != a1.TaskState.ProgramComplete and state != a1.TaskState.Idle: +# if a1.TaskState.ProgramRunning == state: +# self.wait_for_status_to_change(task_id, a1.TaskStatusItem.TaskState, state, timeout) +# elif a1.TaskState.ProgramComplete == state: +# print('can continue') +# else: +# print(f"Task {task_id} is not Idle: {a1.TaskState(state).name} , aborting") +# return +# try: +# print(f"Running script {script_name} on task {task_id}") +# self.controller.runtime.tasks[task_id].program.run(script_name) +# self.wait_program_finish(task_id, timeout) +# except Exception as e: +# print(f"Error executing script {script_name}: {e}") +# +# def run_grid_scan(self, cell_height_mm:float, num_rows:int, +# row_width_mm:float, time_per_row_s: float, task_id:int =3): +# self.set_global_variable(0, cell_height_mm) +# self.set_global_variable(1, row_width_mm) +# self.set_global_variable(2, time_per_row_s) +# self.set_global_variable(1, num_rows) +# self.set_global_variable(0, 1) +# #self.__run_program(script_name="grid_scan.a1exe", task_id=task_id, timeout=120.0) +# def home_all(self, task_id:int = 3): +# self.__run_program(script_name="home_all.a1exe", task_id=task_id, timeout=120.0) +# +# def rotation_scan(self, task_id = 3, start_angle:float = 0.0, end_angle:float = 360.0, step_size:float = 1.0, num_steps:int = 10): +# self.__run_program(script_name="rotation_scan.a1exe", task_id=task_id, timeout=90.0) +# +# def move_motor_absolute(self, axis:str, position:float, speed:float=1.0): +# self.controller.runtime.commands.motion.moveabsolute(axis.upper(), [position], [speed]) +# +# def move_motor_linear(self, axis: str, position: float, speed: float = 1.0): +# self.controller.runtime.commands.motion.movelinear(axis.upper(), [position], speed) +# +# if __name__ == "__main__": +# beamline = mx_beamline() +# print(beamline) +# ### test on 10S +# aerotech = AerotechController(controller_ip="129.129.118.96") +# #aerotech.enable_motion("X") +# #aerotech.home_all() +# rw = 0.320 +# ch = 0.010 +# nr = 10 +# tpr_s = rw / ch * 0.02 +# #aerotech.run_grid_scan(cell_height_mm=ch, num_rows=nr, +# # row_width_mm=rw, time_per_row_s=tpr_s) +# st = time.perf_counter() +# aerotech.move_motor_absolute("Z", 0, 10000) +# print(f"time to move: {time.perf_counter() - st}") +# aerotech.disconnect() #aerotech_epics = AerotechControllerEpics(beamline) #aerotech_epics.omega.speed = 80.0 From 6110a24d61a75e867359c40f9a92499edc254bc5 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 13 Mar 2026 09:40:52 +0100 Subject: [PATCH 07/30] Aerotech: think we are removing automation1 api, WIP --- src/aare/daq/devices.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aare/daq/devices.py b/src/aare/daq/devices.py index 0455beea..075aa57c 100644 --- a/src/aare/daq/devices.py +++ b/src/aare/daq/devices.py @@ -28,7 +28,7 @@ class BeamlineDevices: BEAMLINE = beamline.value.upper() self.tell = make_tell_client(beamline) self.__aerotech = aerotech.AerotechControllerEpics(beamline) - self.aerotech = aerotech.AerotechController(controller_ip="129.129.118.96") + #self.aerotech = aerotech.AerotechController(controller_ip="129.129.118.96") self.__smargon = smargon.Smargon(beamline) self.__ring_current_pv = PV(f"ARS07-DPCT-0100:CURR") From a3d4027b087865f950707459e82942ae78610d86 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 13 Mar 2026 09:41:27 +0100 Subject: [PATCH 08/30] GUI: Use Pserv, spark and x10sa-queue --- src/aare/gui/gui.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aare/gui/gui.py b/src/aare/gui/gui.py index 284ec734..de22821b 100644 --- a/src/aare/gui/gui.py +++ b/src/aare/gui/gui.py @@ -36,8 +36,8 @@ if __name__ == "__main__": default_gonio_cam_addr = "axis-accc8ed2972e.psi.ch" default_gonio_camera_id = 3 case MXBeamline.X10SA: - default_url = "http://127.0.0.1:5210" - default_zmq_addr = "tcp://x10sa-spark-01:9091" #"tcp://x10sa-pserv-01:9089" # + default_url = "http://mx-x10sa-queue-01.psi.ch:5210" #"http://127.0.0.1:5210" + default_zmq_addr = "tcp://x10sa-pserv-01:9089" # "tcp://x10sa-spark-01:9091" # default_pred_zmq_addr = "tcp://x10sa-spark-01:9091" #"tcp://sls-gpu-003:9089"#"" default_beamline_cam_addr = "axis-accc8eb02488.psi.ch" default_gonio_cam_addr = "axis-accc8ea5e463.psi.ch" From 12412c29d812200ebcad51be2db79d68067abf69 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Fri, 20 Mar 2026 11:55:38 +0100 Subject: [PATCH 09/30] Improve SSE client resilience in `tellupdater` by adding retry mechanism and enhanced error logging --- src/aare/daq/tellupdater.py | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 5959fe45..2616ff4e 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -38,21 +38,27 @@ def listen_to_sse(): print(f"[SSE][WARN] No TELL URL configured – SSE listener not started. (tell_client.url={tell_client.url})") return sse_url = tell_client.url + "/events" - try: - # Use URL directly for sseclient variants that manage their own HTTP stream. - client = sseclient.SSEClient(sse_url) + + while True: + try: + print(f"[SSE][INFO] Attempting to connect to {sse_url}...") + # Use URL directly for sseclient variants that manage their own HTTP stream. + client = sseclient.SSEClient(sse_url) - print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") - handle_tell_change_event() + print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") + handle_tell_change_event() - # Compatibility: some SSEClient versions are iterable, others expose .events(). - events_iter = client.events() if hasattr(client, "events") else iter(client) - for event in events_iter: - #print(f"event = {event.event} with data: {event.data}") - if event.event == "DewarContentUpdate": - on_sse_event(event) - except Exception as exc: - print(f"[SSE][listen_to_sse][ERROR] Failed to connect to {sse_url}: {exc}") + # Compatibility: some SSEClient versions are iterable, others expose .events(). + events_iter = client.events() if hasattr(client, "events") else iter(client) + for event in events_iter: + # print(f"event = {event.event} with data: {event.data}") + if event.event == "DewarContentUpdate": + on_sse_event(event) + except Exception as exc: + print(f"[SSE][listen_to_sse][ERROR] Connection lost or failed: {exc}") + + print("[SSE][INFO] Reconnecting to SSE in 5 seconds...") + time.sleep(5) def compare_and_report_change(old, new, key_func): """ From 2e4adb5eff521d862ae2a13822a7166ce7fe1d55 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 15:01:46 +0100 Subject: [PATCH 10/30] Enable mTLS support for WebSocket connections in `tellupdater` by configuring certificate and key files. --- src/aare/daq/tellupdater.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 2616ff4e..6addb2b5 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -147,6 +147,7 @@ def on_close(ws, close_status_code, close_msg): def on_open(ws): print("[WS][OPEN] WebSocket opened.") + def main(): # Start SSE listener in a separate background thread sse_thread = threading.Thread(target=listen_to_sse, daemon=True) @@ -163,7 +164,16 @@ def main(): on_close=on_close, on_open=on_open, ) - ws.run_forever(sslopt={"cert_reqs": 0}) + + # --- mTLS CERTIFICATES ADDED HERE -- + import ssl + ssl_settings = { + "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", + "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key", + "cert_reqs": ssl.CERT_REQUIRED + } + + ws.run_forever(sslopt=ssl_settings) except Exception as e: print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") From a9650876690bf2967a0f5a4dccdaec5f75a0c456 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 15:33:42 +0100 Subject: [PATCH 11/30] Add CA certificate configuration to SSL settings in `tellupdater` to complete mTLS setup --- src/aare/daq/tellupdater.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 6addb2b5..6c21f535 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -170,6 +170,7 @@ def main(): ssl_settings = { "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key", + "ca_certs": "/etc/ssl/certs/secrets/mx-db-01_DigiCert_Global_Root_G2.pem", "cert_reqs": ssl.CERT_REQUIRED } From 6a0761250980111e6dfbb7627cafd2772593830f Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 15:39:21 +0100 Subject: [PATCH 12/30] Enable mandatory mTLS for SSE connections in `tellupdater` by configuring client certificates and CA root files. --- src/aare/daq/tellupdater.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 6c21f535..13bf63af 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -38,12 +38,25 @@ def listen_to_sse(): print(f"[SSE][WARN] No TELL URL configured – SSE listener not started. (tell_client.url={tell_client.url})") return sse_url = tell_client.url + "/events" - + while True: try: print(f"[SSE][INFO] Attempting to connect to {sse_url}...") - # Use URL directly for sseclient variants that manage their own HTTP stream. - client = sseclient.SSEClient(sse_url) + + # --- MANDATORY mTLS FOR mx-db-01 --- + import requests + cert_pair = ( + "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", + "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key" + ) + # The DigiCert root you have on the machine + ca_root = "/etc/ssl/certs/secrets/mx-db-01_DigiCert_Global_Root_G2.pem" + + # Open the stream using the certificates required by Nginx + response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root) + response.raise_for_status() + + client = sseclient.SSEClient(response) print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") handle_tell_change_event() From 4e873fefc358c8346d8e2b0538d17dbb42d54c90 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 16:04:34 +0100 Subject: [PATCH 13/30] Update CA certificate path in `tellupdater` to use the full chain file for improved SSL configuration --- src/aare/daq/tellupdater.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 13bf63af..f4e796ca 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -50,7 +50,7 @@ def listen_to_sse(): "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key" ) # The DigiCert root you have on the machine - ca_root = "/etc/ssl/certs/secrets/mx-db-01_DigiCert_Global_Root_G2.pem" + ca_root = "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem" # Open the stream using the certificates required by Nginx response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root) @@ -183,7 +183,7 @@ def main(): ssl_settings = { "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key", - "ca_certs": "/etc/ssl/certs/secrets/mx-db-01_DigiCert_Global_Root_G2.pem", + "ca_certs": "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem", "cert_reqs": ssl.CERT_REQUIRED } From 87df42a179ded99705baa6dec4745abccba23c3f Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 16:17:45 +0100 Subject: [PATCH 14/30] Conditionally apply mTLS for HTTPS connections in `tellupdater` to avoid certificate usage on plain HTTP requests. --- src/aare/daq/tellupdater.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index f4e796ca..14f47bf6 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -45,15 +45,19 @@ def listen_to_sse(): # --- MANDATORY mTLS FOR mx-db-01 --- import requests - cert_pair = ( - "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", - "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key" - ) - # The DigiCert root you have on the machine - ca_root = "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem" + # --- FIX: Only use certs for HTTPS (DB) connections --- + if sse_url.startswith("https:"): + cert_pair = ( + "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", + "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key" + ) + # Use the combined bundle (Intermediate + Root) if you made it + ca_root = "/etc/ssl/certs/secrets/mx-db-01_DigiCert_Global_Root_G2.pem" + response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root) + else: + # Robot connection (PC17488): No certs allowed on plain HTTP + response = requests.get(sse_url, stream=True) - # Open the stream using the certificates required by Nginx - response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root) response.raise_for_status() client = sseclient.SSEClient(response) From a4b6c2575ee3296244a7ef1a4e7ea34cc3580589 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 17:18:42 +0100 Subject: [PATCH 15/30] Update certificate and key file paths in `tellupdater` to use database-specific mTLS credentials --- src/aare/daq/tellupdater.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 14f47bf6..b1394475 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -185,8 +185,8 @@ def main(): # --- mTLS CERTIFICATES ADDED HERE -- import ssl ssl_settings = { - "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", - "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key", + "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt", + "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key", "ca_certs": "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem", "cert_reqs": ssl.CERT_REQUIRED } From d6d20d598ff4133acea579de0d4fbf3ad6667222 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 17:28:11 +0100 Subject: [PATCH 16/30] Simplify mTLS handling in `tellupdater` by consolidating certificate paths and streamlining HTTP/HTTPS logic. --- src/aare/daq/tellupdater.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index b1394475..9ab979ba 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -42,25 +42,18 @@ def listen_to_sse(): while True: try: print(f"[SSE][INFO] Attempting to connect to {sse_url}...") - - # --- MANDATORY mTLS FOR mx-db-01 --- - import requests - # --- FIX: Only use certs for HTTPS (DB) connections --- - if sse_url.startswith("https:"): - cert_pair = ( - "/etc/ssl/certs/secrets/mx-x10sa-queue-01.crt", - "/etc/ssl/certs/secrets/mx-x10sa-queue-01.key" - ) - # Use the combined bundle (Intermediate + Root) if you made it - ca_root = "/etc/ssl/certs/secrets/mx-db-01_DigiCert_Global_Root_G2.pem" + if sse_url.startswith("https://mx-db-01"): + # mTLS path + import requests + cert_pair = ("/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt", "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key") + ca_root = "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem" response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root) + response.raise_for_status() + client = sseclient.SSEClient(response) else: - # Robot connection (PC17488): No certs allowed on plain HTTP - response = requests.get(sse_url, stream=True) - - response.raise_for_status() - - client = sseclient.SSEClient(response) + # Robot path (PC17488) - Pass URL STRING directly + # SSEClient will handle the simple HTTP GET itself + client = sseclient.SSEClient(sse_url) print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") handle_tell_change_event() From eaf3083eca6142e2c263fb1ffe8923d978486985 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 24 Mar 2026 20:35:58 +0100 Subject: [PATCH 17/30] Bump `aaredb` dependency version to `0.1.1a43` in `pyproject.toml`. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 044818a6..e4529c7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ "python-redis-lock==4.0.0", "fastapi==0.115.13", "uvicorn==0.34.2", - "aaredb==0.1.1a42", + "aaredb==0.1.1a43", "python_multipart==0.0.20", "websocket-client==1.8.0", "sseclient-py==1.8.0", From 100895f56bc8d27ba56a1fd6ad12b78b99943ac7 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 13:49:43 +0100 Subject: [PATCH 18/30] Added mTLS certificates authentication to aaredb --- src/aare/daq/aaredb.py | 12 ++++++++++-- src/aare/daq/spreadsheetupdater.py | 22 ++++++++++++++++++++-- src/aare/daq/tellupdater.py | 14 +++++++++----- 3 files changed, 39 insertions(+), 9 deletions(-) diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index e165f7b1..0ee85151 100644 --- a/src/aare/daq/aaredb.py +++ b/src/aare/daq/aaredb.py @@ -38,6 +38,7 @@ from jfjoch_client.models import ScanResult logger = setup_logger("aareDAQ") + class AareWrapper: def __init__( self, @@ -45,14 +46,21 @@ class AareWrapper: host: str = "https://mx-db-01.psi.ch/dispatcher", ): configuration = aareDB.Configuration(host=host) - configuration.verify_ssl = False # Disable SSL verification + + # --- mTLS & SSL CONFIGURATION --- + configuration.cert_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt" + configuration.key_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key" + + # Point this to the CA bundle that signed the server's certificate + configuration.verify_ssl = True + configuration.ssl_ca_cert = "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem" self.client = aareDB.ApiClient(configuration) self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD") self.__host = host self.__tell_api = aareDB.TellsRunnerApi(self.client) self.__sample_api = aareDB.SamplesRunnerApi(self.client) - self.__proc_api = aareDB.ProcessingsRunnerApi(self.client) + self.__proc_api = aareDB.ProcessingsRunnerApi(self.client) self.__raster_api = aareDB.GridscanRunnerApi(self.client) self.__bl = bl diff --git a/src/aare/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index 4f7c65f8..c1cd28af 100644 --- a/src/aare/daq/spreadsheetupdater.py +++ b/src/aare/daq/spreadsheetupdater.py @@ -9,7 +9,9 @@ from aare.common.beamline import MXBeamline, mx_beamline beamline = mx_beamline() SLOT_IDENTIFIER = beamline.value.upper() -WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/tell_runner/ws/samples-spreadsheet/{SLOT_IDENTIFIER}" +#WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/tell_runner/ws/samples-spreadsheet/{SLOT_IDENTIFIER}" +WS_URL = f"wss://mx-aaredb-dmz-01.psi.ch/dispatcher/protected_router/tell_runner/ws/samples-spreadsheet/{SLOT_IDENTIFIER}" + # Ensure the environment variable for the shared password is set password = os.getenv("AAREDB_SHARED_PASSWORD") @@ -142,7 +144,23 @@ def main(): on_close=on_close, on_open=on_open, ) - ws.run_forever(sslopt={"cert_reqs": 0}) + + # --- mTLS CONFIGURATION --- + import ssl + ssl_settings = { + # This is the "Machine Certificate" NGINX will verify + "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_aaredb-dmz-01.crt", + "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_aaredb-dmz-01.key", + + # This allows Python to trust the Server (NGINX/PSI) + "ca_certs": "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem", + + # Standard TLS requirements + "cert_reqs": ssl.CERT_REQUIRED, + "check_hostname": True, + } + + ws.run_forever(sslopt=ssl_settings) except Exception as e: print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 9ab979ba..1023ad76 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -18,7 +18,8 @@ logger = setup_logger("aareDAQ") # Configuration beamline = mx_beamline() SLOT_IDENTIFIER = beamline.value.upper() -WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" +#WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" +WS_URL = f"wss://mx-aaredb-dmz-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" #WS_URL = f"wss://localhost:8001/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" WS_HEADERS = [f"X-Shared-Password: {os.getenv('AAREDB_SHARED_PASSWORD')}"] print(WS_HEADERS) @@ -42,18 +43,21 @@ def listen_to_sse(): while True: try: print(f"[SSE][INFO] Attempting to connect to {sse_url}...") - if sse_url.startswith("https://mx-db-01"): + if sse_url.startswith("https://mx-aaredb-dmz-01"): #"https://mx-db-01" # mTLS path import requests - cert_pair = ("/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt", "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key") - ca_root = "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem" + #cert_pair = ("/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt", "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key") + cert_pair = ("/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt", + "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key") + #ca_root = "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem" + ca_root = "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem" response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root) response.raise_for_status() client = sseclient.SSEClient(response) else: # Robot path (PC17488) - Pass URL STRING directly # SSEClient will handle the simple HTTP GET itself - client = sseclient.SSEClient(sse_url) + client = sseclient.SSEClient(sse_url) print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") handle_tell_change_event() From b24e3aa12aba6517aa4ced2b8c59ae9b2d3056d6 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Wed, 25 Mar 2026 13:51:39 +0100 Subject: [PATCH 19/30] ml_box: changed to new model yolo26l-seg-overlap-false_2026-03-16.engine --- src/aare/daq/mlbox.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py index 094233c2..84d00fc7 100644 --- a/src/aare/daq/mlbox.py +++ b/src/aare/daq/mlbox.py @@ -30,7 +30,7 @@ class MlBox: elif bl == MXBeamline.X06DA: self.__url = "http://mx-aare-test.psi.ch:8002/predict/?model=best_v8_20102025.pt" elif bl == MXBeamline.X10SA: - self.__url = "http://x10sa-spark-01.psi.ch:8002/predict/?model=best_v12_22092025.engine" + self.__url = "http://x10sa-spark-01.psi.ch:8002/predict/?model=best_yolo26l-seg-overlap-false_2026-03-16.engine"#v12_22092025.engine" elif bl == MXBeamline.X06SA: self.__url = "" raise NotImplemented(f"MLBox not implemente for {bl}") From 6e3e14a837367a7b9eb5b9a5a1b0dd646f804324 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Wed, 25 Mar 2026 13:57:10 +0100 Subject: [PATCH 20/30] aaredb: changed certification, host and logging statements --- src/aare/daq/aaredb.py | 42 ++++++++++++++++++++---------------------- 1 file changed, 20 insertions(+), 22 deletions(-) diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index 0ee85151..1e70683b 100644 --- a/src/aare/daq/aaredb.py +++ b/src/aare/daq/aaredb.py @@ -38,22 +38,20 @@ from jfjoch_client.models import ScanResult logger = setup_logger("aareDAQ") - class AareWrapper: def __init__( self, bl: MXBeamline, - host: str = "https://mx-db-01.psi.ch/dispatcher", + host: str = "https://mx-aaredb-dmz-01.psi.ch/dispatcher", ): configuration = aareDB.Configuration(host=host) - # --- mTLS & SSL CONFIGURATION --- - configuration.cert_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt" - configuration.key_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key" # Point this to the CA bundle that signed the server's certificate configuration.verify_ssl = True configuration.ssl_ca_cert = "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem" + configuration.cert_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt" + configuration.key_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" self.client = aareDB.ApiClient(configuration) self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD") @@ -78,7 +76,7 @@ class AareWrapper: ret = self.__tell_api.set_tell_positions( set_tell_position_request=payload, ) - print(ret) + logger.debug(ret) def create_manual_sample(self, s: SampleShortInfo): from aareDB.models import ManualSampleCreate @@ -92,7 +90,7 @@ class AareWrapper: try: s.db_id = self.__sample_api.insert_sample(manual_sample).id except Exception as e: - print(f"Error inserting sample: {e}") + logger.error(f"Error inserting sample: {e}") def sample_mounted(self, s: Optional[SampleShortInfo]): if s is not None: @@ -102,7 +100,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("Mounted")), ) except Exception as e: - print(e) + logger.error(e) def sample_unmounted(self, s: Optional[SampleShortInfo]): if s is not None: @@ -112,7 +110,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("Unmounted")), ) except Exception as e: - print(e) + logger.error(e) def sample_centered(self, s: Optional[SampleShortInfo]): if s is not None: @@ -122,7 +120,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("Centered")), ) except Exception as e: - print(e) + logger.error(e) def sample_collected(self, s: Optional[SampleShortInfo]): if s is None: @@ -133,7 +131,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("Collected")), ) except Exception as e: - print(e) + logger.error(e) def sample_failed(self, s: Optional[SampleShortInfo], failed_comment: Optional[str] = None): if s is None: @@ -144,7 +142,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("Failed"), comment=failed_comment), ) except Exception as e: - print(e) + logger.error(e) def axc_failed(self, s: Optional[SampleShortInfo]): if s is None: @@ -155,7 +153,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("AXCFailed")), ) except Exception as e: - print(e) + logger.error(e) def alc_failed(self, s: Optional[SampleShortInfo], alc_comment: Optional[str] = None): if s is None: @@ -166,7 +164,7 @@ class AareWrapper: sample_event_create=SampleEventCreate(event_type=SampleEventType("ALCFailed"), comment=alc_comment), ) except Exception as e: - print(e) + logger.error(e) def sample_lost(self, s: Optional[SampleShortInfo]): if s is None: @@ -296,9 +294,9 @@ class AareWrapper: sample_id=s.db_id, experiment_parameters_create=experiment_params_payload ) - print("Experiment parameters created:", response) + logger.debug("Experiment parameters created:", response) except Exception as e: - print(e) + logger.error(e) def create_gridscan_run(self, s: Optional[SampleShortInfo], r:RasterGridRequest, d:DAQStatusModel): if s is None: @@ -360,9 +358,9 @@ class AareWrapper: sample_id=s.db_id, experiment_parameters_create=experiment_params_payload ) - print("Experiment parameters created:", response) + logger.info("Experiment parameters created:", response) except Exception as e: - print(e) + logger.debug(e) def ingest_gridscan(self, sample: Optional[SampleShortInfo], raster_result: ScanResult, raster_request: RasterGridRequest, geom: SampleGeometryModel, @@ -386,7 +384,7 @@ class AareWrapper: headers=headers, data=json.dumps(payload), timeout=30, verify=False) response.raise_for_status() - print(f"Response status code: {response.status_code}") + logger.info(f"Response status code: {response.status_code}") def format_gridscan_payload(self, sample: Optional[SampleShortInfo], raster_result:ScanResult, @@ -429,7 +427,7 @@ class AareWrapper: return payload except Exception as e: - print(e) + logger.error(e) raise e def ingest_scan(self, sample: Optional[SampleShortInfo], result: ScanResult, @@ -453,7 +451,7 @@ class AareWrapper: headers=headers, data=json.dumps(payload), timeout=30, verify=False) response.raise_for_status() - print(f"Response status code: {response.status_code}") + logger.info(f"Response status code: {response.status_code}") def format_scan_payload(self, sample: Optional[SampleShortInfo], result:ScanResult, geom:SampleGeometryModel, @@ -471,5 +469,5 @@ class AareWrapper: return payload except Exception as e: - print(e) + logger.error(e) raise e From f1a9147f107727abe117e1212cbd1b59f10e0742 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 13:59:02 +0100 Subject: [PATCH 21/30] modify tellupdater mTLS certificates --- src/aare/daq/tellupdater.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 1023ad76..c87e1059 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -170,6 +170,19 @@ def main(): # Main thread runs websocket client loop while True: try: + import ssl + # 1. Create a proper SSL Context + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + + # 2. Load the CA to trust the server + context.load_verify_locations(cafile="/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem") + + # 3. Load the Client Certificate (mTLS) + context.load_cert_chain( + certfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt", + keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" + ) + ws = websocket.WebSocketApp( WS_URL, header=WS_HEADERS, @@ -179,16 +192,9 @@ def main(): on_open=on_open, ) - # --- mTLS CERTIFICATES ADDED HERE -- - import ssl - ssl_settings = { - "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt", - "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key", - "ca_certs": "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem", - "cert_reqs": ssl.CERT_REQUIRED - } + # 4. Pass the context object to run_forever + ws.run_forever(sslopt={"context": context}) - ws.run_forever(sslopt=ssl_settings) except Exception as e: print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") From fd20ae50d805a0065c0778701d83d5f087f12131 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:08:00 +0100 Subject: [PATCH 22/30] modify tellupdater mTLS certificates --- src/aare/daq/aaredb.py | 13 +++++++++---- src/aare/daq/tellupdater.py | 21 ++++++++++++++------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index 1e70683b..2d1b9785 100644 --- a/src/aare/daq/aaredb.py +++ b/src/aare/daq/aaredb.py @@ -46,14 +46,19 @@ class AareWrapper: ): configuration = aareDB.Configuration(host=host) - - # Point this to the CA bundle that signed the server's certificate + # --- mTLS & SSL CONFIGURATION --- + # 1. Trust the Server (CA that signed mx-aaredb-dmz-01) configuration.verify_ssl = True configuration.ssl_ca_cert = "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem" - configuration.cert_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt" - configuration.key_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" + # 2. Present Machine Identity (The certs that worked in curl) + configuration.cert_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt" + configuration.key_file = "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" + + # 3. Initialize the Client with this config self.client = aareDB.ApiClient(configuration) + + # Identity Forwarding (Optional now that mTLS is active, but safe to keep) self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD") self.__host = host self.__tell_api = aareDB.TellsRunnerApi(self.client) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index c87e1059..9792741d 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -167,22 +167,28 @@ def main(): sse_thread = threading.Thread(target=listen_to_sse, daemon=True) sse_thread.start() - # Main thread runs websocket client loop + """ + Main function to initiate WebSocket connection with mTLS. + """ while True: try: import ssl - # 1. Create a proper SSL Context + # 1. Create a modern SSL Context for a TLS Client context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - # 2. Load the CA to trust the server + # 2. Load the CA to verify the NGINX server's identity context.load_verify_locations(cafile="/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem") - # 3. Load the Client Certificate (mTLS) + # 3. Load the Client Certificate and Key (mTLS) + # Using the 'dmz-01' paths that worked in your curl context.load_cert_chain( certfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt", keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" ) + # Optional: Ensure hostname matching is active (recommended) + context.check_hostname = True + ws = websocket.WebSocketApp( WS_URL, header=WS_HEADERS, @@ -192,13 +198,14 @@ def main(): on_open=on_open, ) - # 4. Pass the context object to run_forever + # 4. Pass the context directly via sslopt + print(f"[WS][INFO] Connecting to {WS_URL} using mTLS...") ws.run_forever(sslopt={"context": context}) except Exception as e: - print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") + print(f"[MAIN][ERROR] WebSocket connection failed: {e}") - print("[WS][INFO] WebSocket connection lost. Reconnecting in 5 seconds...") + print("[WS][INFO] Reconnecting in 5 seconds...") time.sleep(5) if __name__ == "__main__": From 25163a5b4dfc17c27d185247af510f5de600ddd7 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:10:17 +0100 Subject: [PATCH 23/30] modify spreadsheetupdater.py mTLS certificates --- src/aare/daq/spreadsheetupdater.py | 41 +++++++++++++++++------------- 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/src/aare/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index c1cd28af..204a8102 100644 --- a/src/aare/daq/spreadsheetupdater.py +++ b/src/aare/daq/spreadsheetupdater.py @@ -2,6 +2,7 @@ import os import json import websocket import time +import ssl from aareDB.models import PuckWithTellPosition from aare.common.models import SampleShortInfoList, SampleShortInfo, DewarAddress from config import BeamlineConfig @@ -136,6 +137,25 @@ def main(): """ while True: try: + # 1. Create the Context + # Use PROTOCOL_TLS_CLIENT as it's the modern standard + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + + # 2. Load the Server Trust (CA) + context.load_verify_locations(cafile="/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem") + + # 3. Load the Client Identity (The key/cert that worked in curl) + context.load_cert_chain( + certfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt", + keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" + ) + + # 4. Critical: If NGINX is using a self-signed or internal CA, + # and the hostname doesn't perfectly match the Cert SN/SAN, + # you might need to toggle this, but keep it True for now. + context.check_hostname = True + context.verify_mode = ssl.CERT_REQUIRED + ws = websocket.WebSocketApp( WS_URL, header=WS_HEADERS, @@ -145,26 +165,13 @@ def main(): on_open=on_open, ) - # --- mTLS CONFIGURATION --- - import ssl - ssl_settings = { - # This is the "Machine Certificate" NGINX will verify - "certfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_aaredb-dmz-01.crt", - "keyfile": "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_aaredb-dmz-01.key", + # 5. Connect using ONLY the context + # We remove other keys to ensure websocket-client uses our pre-configured context + ws.run_forever(sslopt={"context": context}) - # This allows Python to trust the Server (NGINX/PSI) - "ca_certs": "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem", - - # Standard TLS requirements - "cert_reqs": ssl.CERT_REQUIRED, - "check_hostname": True, - } - - ws.run_forever(sslopt=ssl_settings) except Exception as e: - print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") + print(f"[MAIN][ERROR] WebSocket setup failed: {e}") - print("[WS][INFO] WebSocket connection lost. Reconnecting in 5 seconds...") time.sleep(5) if __name__ == "__main__": From 97d7e561251579b4af5901937348db238223ba66 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:21:07 +0100 Subject: [PATCH 24/30] modify spreadsheetupdater.py mTLS certificates --- src/aare/daq/spreadsheetupdater.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/src/aare/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index 204a8102..fca2ff7a 100644 --- a/src/aare/daq/spreadsheetupdater.py +++ b/src/aare/daq/spreadsheetupdater.py @@ -137,25 +137,18 @@ def main(): """ while True: try: - # 1. Create the Context - # Use PROTOCOL_TLS_CLIENT as it's the modern standard + # Initialize with the modern client protocol context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - # 2. Load the Server Trust (CA) + # Trust the PSI DMZ CA context.load_verify_locations(cafile="/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem") - # 3. Load the Client Identity (The key/cert that worked in curl) + # Load your machine identity context.load_cert_chain( certfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt", keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" ) - # 4. Critical: If NGINX is using a self-signed or internal CA, - # and the hostname doesn't perfectly match the Cert SN/SAN, - # you might need to toggle this, but keep it True for now. - context.check_hostname = True - context.verify_mode = ssl.CERT_REQUIRED - ws = websocket.WebSocketApp( WS_URL, header=WS_HEADERS, @@ -165,8 +158,7 @@ def main(): on_open=on_open, ) - # 5. Connect using ONLY the context - # We remove other keys to ensure websocket-client uses our pre-configured context + print(f"[WS][INFO] Connecting to {WS_URL}...") ws.run_forever(sslopt={"context": context}) except Exception as e: From 12b749c4e63d9ff5f951b5269045afd7036d8ca2 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:24:48 +0100 Subject: [PATCH 25/30] modify spreadsheetupdater.py mTLS certificates --- src/aare/daq/spreadsheetupdater.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/aare/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index fca2ff7a..e7dc1ee3 100644 --- a/src/aare/daq/spreadsheetupdater.py +++ b/src/aare/daq/spreadsheetupdater.py @@ -149,6 +149,8 @@ def main(): keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" ) + context.check_hostname = True + ws = websocket.WebSocketApp( WS_URL, header=WS_HEADERS, @@ -158,7 +160,7 @@ def main(): on_open=on_open, ) - print(f"[WS][INFO] Connecting to {WS_URL}...") + print(f"[WS][INFO] Connecting to {WS_URL} using mTLS...") ws.run_forever(sslopt={"context": context}) except Exception as e: From 0de5cbfc01d652a8d899a07a215b641045415c97 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:28:16 +0100 Subject: [PATCH 26/30] modify spreadsheetupdater.py mTLS certificates --- src/aare/daq/spreadsheetupdater.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aare/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index e7dc1ee3..8f4b1f06 100644 --- a/src/aare/daq/spreadsheetupdater.py +++ b/src/aare/daq/spreadsheetupdater.py @@ -2,7 +2,6 @@ import os import json import websocket import time -import ssl from aareDB.models import PuckWithTellPosition from aare.common.models import SampleShortInfoList, SampleShortInfo, DewarAddress from config import BeamlineConfig @@ -137,6 +136,7 @@ def main(): """ while True: try: + import ssl # Initialize with the modern client protocol context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) From d98769d67ca34b5d976b15ac8cc32524484c9ddf Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:32:43 +0100 Subject: [PATCH 27/30] modify spreadsheetupdater.py mTLS certificates --- src/aare/daq/spreadsheetupdater.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/aare/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index 8f4b1f06..5cad17fb 100644 --- a/src/aare/daq/spreadsheetupdater.py +++ b/src/aare/daq/spreadsheetupdater.py @@ -137,19 +137,23 @@ def main(): while True: try: import ssl - # Initialize with the modern client protocol - context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + import websocket - # Trust the PSI DMZ CA + # FORCE a clean context + context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) context.load_verify_locations(cafile="/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem") - - # Load your machine identity context.load_cert_chain( certfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt", keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key" ) - context.check_hostname = True + # Explicitly set the SNI hostname to match NGINX server_name + # This is often what's missing when NGINX says "No cert sent" + ssl_opt = { + "context": context, + "server_hostname": "mx-aaredb-dmz-01.psi.ch", + "check_hostname": True + } ws = websocket.WebSocketApp( WS_URL, @@ -160,11 +164,11 @@ def main(): on_open=on_open, ) - print(f"[WS][INFO] Connecting to {WS_URL} using mTLS...") - ws.run_forever(sslopt={"context": context}) + print(f"[WS][INFO] Connecting to {WS_URL}...") + ws.run_forever(sslopt=ssl_opt) except Exception as e: - print(f"[MAIN][ERROR] WebSocket setup failed: {e}") + print(f"[MAIN][ERROR] WebSocket connection failed: {e}") time.sleep(5) From 3c1cf09417e0d3ace3f42b4600baf4de748ec540 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Thu, 26 Mar 2026 11:20:31 +0100 Subject: [PATCH 28/30] aaredaq: continuing work on baton exchange procedure --- src/aare/common/auth_models.py | 51 ++++ src/aare/daq/auth.py | 264 +++++++++++++++++++ src/aare/daq/config.py | 129 ++++++++- src/aare/daq/server.py | 137 +++++++++- src/aare/gui/main_window.py | 143 ++++++++++ src/aare/gui/threads/daq_worker.py | 160 ++++++++++- src/aare/gui/widgets/alert_banner.py | 75 ++++++ src/aare/gui/widgets/baton_request_dialog.py | 220 ++++++++++++++++ src/aare/gui/widgets/status_bar.py | 261 +++++++++++++++--- 9 files changed, 1402 insertions(+), 38 deletions(-) create mode 100644 src/aare/common/auth_models.py create mode 100644 src/aare/gui/widgets/baton_request_dialog.py diff --git a/src/aare/common/auth_models.py b/src/aare/common/auth_models.py new file mode 100644 index 00000000..6897751c --- /dev/null +++ b/src/aare/common/auth_models.py @@ -0,0 +1,51 @@ +from enum import Enum +from pydantic import BaseModel +from datetime import datetime + +class BatonRequestStatus(Enum): + PENDING = "pending" + ACCEPTED = "accepted" + REFUSED = "refused" + TIMEOUT = "timeout" + CANCELLED = "cancelled" + + +class BatonHolderInfo(BaseModel): + """Information about the current baton holder.""" + username: str + session: int + is_staff: bool + pgroup: str | None = None + + +class BatonRequest(BaseModel): + """A request from one user to take the baton from another.""" + request_id: str + requester_username: str + requester_session: int + requester_is_staff: bool + holder_username: str | None = None + holder_session: int | None = None + created_at: float # Unix timestamp + timeout_seconds: int = 30 + status: BatonRequestStatus = BatonRequestStatus.PENDING + + +class BatonTransferQueue(BaseModel): + """Queued baton transfer waiting for beamline to be available.""" + target_session: int + target_username: str + target_is_staff: bool + target_pgroup: str | None = None + queued_at: float # Unix timestamp + reason: str = "beamline_busy" + + +class BatonStatus(BaseModel): + """Full baton status for GUI display.""" + holder: BatonHolderInfo | None = None + pending_request: BatonRequest | None = None + queued_transfer: BatonTransferQueue | None = None + you_are_holder: bool = False + you_have_pending_request: bool = False + incoming_request: bool = False # True if YOU are being asked to give up baton \ No newline at end of file diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index 3af7f247..06d24829 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -1,14 +1,18 @@ import grp import os import pwd +import uuid from datetime import datetime, timedelta, UTC from typing import List +import time import jwt from fastapi import HTTPException, status from fastapi.security import OAuth2PasswordRequestForm from pydantic import BaseModel +from aare.common.auth_models import BatonRequestStatus, BatonTransferQueue, BatonRequest, BatonStatus +from aare.common.models import SessionsStateEnum from aare.daq.config import BeamlineConfig from aare.common.exception_handler import AuthenticationException, UserRightsException, AuthErrorCode @@ -21,6 +25,7 @@ SECRET_KEY = os.environ.get("JWT_AAREDAQ_KEY") ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 24 * 60 * 7 # 1 week SESSION_EXPIRE_SECONDS = 60 * 10 +BATON_REQUEST_TIMEOUT_SECONDS = 30 STAFF_GROUP = "unx-MXgroup" SUPER_USERS = ["e10019", "e11206", "e18147"] @@ -113,3 +118,262 @@ def check_jwt_staff(cfg: BeamlineConfig, data: TokenData) -> None: def force_current_sesion(cfg: BeamlineConfig, data: TokenData) -> None: cfg.force_set_active_session(data.session, SESSION_EXPIRE_SECONDS) + +def _finalize_expired_baton_request(cfg: BeamlineConfig, pending: BatonRequest) -> BatonStatus: + """ + Resolve an expired baton request in one place. + + Rules: + - staff can override immediately if beamline is free + - otherwise the transfer is queued + - pending request is cleared once it is no longer pending + """ + requester_is_staff = bool(pending.requester_is_staff) + + if cfg.can_transfer_baton_now(): + cfg.execute_baton_transfer( + to_session=pending.requester_session, + to_username=pending.requester_username, + to_is_staff=requester_is_staff, + to_pgroup=cfg.pgroup, + expiry_sec=SESSION_EXPIRE_SECONDS + ) + else: + cfg.queued_baton_transfer = BatonTransferQueue( + target_session=pending.requester_session, + target_username=pending.requester_username, + target_is_staff=requester_is_staff, + target_pgroup=cfg.pgroup, + queued_at=time.time(), + reason="timeout_beamline_busy" + ) + + cfg.clear_pending_baton_request() + return get_baton_status(cfg, TokenData( + sub=pending.requester_username, + pgroups=[], + session=pending.requester_session, + staff=requester_is_staff + )) + +def resolve_baton_timeout_if_needed(cfg: BeamlineConfig) -> BatonStatus | None: + """ + Check the current pending request and resolve it if expired. + Returns the updated BatonStatus when a timeout was processed, else None. + """ + pending = cfg.pending_baton_request + if pending is None or pending.status != BatonRequestStatus.PENDING: + return None + + elapsed = time.time() - pending.created_at + if elapsed < pending.timeout_seconds: + return None + + return _finalize_expired_baton_request(cfg, pending) + +def get_baton_status(cfg: BeamlineConfig, data: TokenData) -> BatonStatus: + """Get full baton status for a user.""" + pending = cfg.pending_baton_request + + if pending is not None and pending.status == BatonRequestStatus.PENDING: + elapsed = time.time() - pending.created_at + if elapsed >= pending.timeout_seconds: + resolve_baton_timeout_if_needed(cfg) + pending = cfg.pending_baton_request + + holder = cfg.baton_holder + queued = cfg.queued_baton_transfer + + you_are_holder = holder is not None and holder.session == data.session + you_have_pending_request = ( + pending is not None and + pending.requester_session == data.session and + pending.status == BatonRequestStatus.PENDING + ) + incoming_request = ( + pending is not None and + holder is not None and + holder.session == data.session and + pending.status == BatonRequestStatus.PENDING + ) + + return BatonStatus( + holder=holder, + pending_request=pending, + queued_transfer=queued, + you_are_holder=you_are_holder, + you_have_pending_request=you_have_pending_request, + incoming_request=incoming_request + ) + + +def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: + """ + Request the baton (control) of the beamline. + """ + resolve_baton_timeout_if_needed(cfg) + + session_state = cfg.session_state(data.session) + + if session_state == SessionsStateEnum.Vacant: + cfg.execute_baton_transfer( + to_session=data.session, + to_username=data.sub, + to_is_staff=data.staff, + to_pgroup=cfg.pgroup, + expiry_sec=SESSION_EXPIRE_SECONDS + ) + return {"granted": True, "message": "Baton acquired (beamline was vacant)"} + + if session_state == SessionsStateEnum.OwnedByYou: + cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS) + return {"already_holder": True, "message": "You already hold the baton"} + + holder = cfg.baton_holder + + if holder and holder.is_staff and not data.staff: + return { + "error": True, + "message": "Cannot request baton from staff. Please ask them directly." + } + + if data.staff: + if not cfg.can_transfer_baton_now(): + cfg.queued_baton_transfer = BatonTransferQueue( + target_session=data.session, + target_username=data.sub, + target_is_staff=data.staff, + target_pgroup=cfg.pgroup, + queued_at=time.time(), + reason="beamline_busy_staff_override" + ) + return { + "queued": True, + "message": "Staff override queued - will transfer when beamline is available" + } + + cfg.execute_baton_transfer( + to_session=data.session, + to_username=data.sub, + to_is_staff=data.staff, + to_pgroup=cfg.pgroup, + expiry_sec=SESSION_EXPIRE_SECONDS + ) + return {"granted": True, "override": True, "message": "Staff override - baton acquired"} + + existing_request = cfg.pending_baton_request + if existing_request and existing_request.status == BatonRequestStatus.PENDING: + if existing_request.requester_session == data.session: + elapsed = time.time() - existing_request.created_at + if elapsed >= existing_request.timeout_seconds: + return {"timeout": True, "message": "Request timed out"} + remaining = existing_request.timeout_seconds - elapsed + return { + "pending": True, + "existing": True, + "remaining_seconds": max(0, remaining), + "message": f"Request already pending ({remaining:.0f}s remaining)" + } + return { + "error": True, + "message": "Another user already has a pending request" + } + + request = BatonRequest( + request_id=str(uuid.uuid4()), + requester_username=data.sub, + requester_session=data.session, + requester_is_staff=data.staff, + holder_username=holder.username if holder else None, + holder_session=holder.session if holder else None, + created_at=time.time(), + timeout_seconds=BATON_REQUEST_TIMEOUT_SECONDS, + status=BatonRequestStatus.PENDING + ) + cfg.set_pending_baton_request(request, timeout_sec=BATON_REQUEST_TIMEOUT_SECONDS) + + return { + "pending": True, + "request_id": request.request_id, + "timeout_seconds": BATON_REQUEST_TIMEOUT_SECONDS, + "message": f"Request sent to {holder.username if holder else 'current user'}" + } + +def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) -> dict: + """ + Current baton holder responds to a pending request. + """ + resolve_baton_timeout_if_needed(cfg) + + holder = cfg.baton_holder + if holder is None or holder.session != data.session: + return {"error": True, "message": "You are not the current baton holder"} + + pending = cfg.pending_baton_request + if pending is None or pending.status != BatonRequestStatus.PENDING: + return {"error": True, "message": "No pending request to respond to"} + + if accept: + if cfg.can_transfer_baton_now(): + cfg.execute_baton_transfer( + to_session=pending.requester_session, + to_username=pending.requester_username, + to_is_staff=pending.requester_is_staff, + to_pgroup=cfg.pgroup, + expiry_sec=SESSION_EXPIRE_SECONDS + ) + return {"accepted": True, "transferred": True, "message": "Baton transferred"} + else: + cfg.queued_baton_transfer = BatonTransferQueue( + target_session=pending.requester_session, + target_username=pending.requester_username, + target_is_staff=pending.requester_is_staff, + target_pgroup=cfg.pgroup, + queued_at=time.time(), + reason="accepted_beamline_busy" + ) + cfg.clear_pending_baton_request() + return { + "accepted": True, + "queued": True, + "message": "Request accepted - will transfer when beamline is available" + } + else: + pending.status = BatonRequestStatus.REFUSED + cfg.set_pending_baton_request(pending, timeout_sec=5) + return {"refused": True, "message": "Request refused"} + +def release_baton(cfg: BeamlineConfig, data: TokenData) -> dict: + """ + Voluntarily release the baton (set session to free). + """ + resolve_baton_timeout_if_needed(cfg) + + holder = cfg.baton_holder + if holder is None: + return {"released": True, "message": "Baton was already vacant"} + + if holder.session != data.session: + return {"info": True, "message": "You don't hold the baton"} + + cfg.end_active_session(data.session) + cfg.baton_holder = None + cfg.clear_pending_baton_request() + + return {"released": True, "message": "Baton released - beamline is now vacant"} + +def cancel_baton_request(cfg: BeamlineConfig, data: TokenData) -> dict: + """ + Cancel your own pending baton request. + """ + resolve_baton_timeout_if_needed(cfg) + + pending = cfg.pending_baton_request + if pending is None: + return {"error": True, "message": "No pending request to cancel"} + + if pending.requester_session != data.session: + return {"error": True, "message": "You can only cancel your own request"} + + cfg.clear_pending_baton_request() + return {"cancelled": True, "message": "Request cancelled"} \ No newline at end of file diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index 61196563..ba8f8328 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -18,6 +18,13 @@ from aare.common.models import ( FluorescenceSpectrumOutputModel, CrystalSize, SimpleStrategyInputModel, SimpleScanParameters ) +from aare.common.auth_models import ( + BatonStatus, + BatonRequest, + BatonHolderInfo, + BatonRequestStatus, + BatonTransferQueue, +) from aare.common.beamline import MXBeamline from aare.common.logger_config import setup_logger @@ -152,6 +159,7 @@ class BeamlineConfig: return if active == session: self.__client.delete(f"{self.__bl}:active_session") + self.__client.delete(f"{self.__bl}:baton_holder") def force_set_active_session(self, session: int, expiry_sec: int) -> None: # Ensure that there is no active try-set for active session @@ -161,6 +169,125 @@ class BeamlineConfig: self.__client.set(f"{self.__bl}:active_session", session) self.__client.expire(f"{self.__bl}:active_session", expiry_sec) + # ========== BATON SYSTEM ========== + + @property + def baton_holder(self) -> BatonHolderInfo | None: + """Get information about the current baton holder.""" + tmp = self.__client.get(f"{self.__bl}:baton_holder") + if tmp is None: + return None + try: + return BatonHolderInfo(**json.loads(tmp)) + except Exception: + return None + + @baton_holder.setter + def baton_holder(self, info: BatonHolderInfo | None) -> None: + if info is None: + self.__client.delete(f"{self.__bl}:baton_holder") + else: + self.__client.set(f"{self.__bl}:baton_holder", info.model_dump_json()) + + @property + def pending_baton_request(self) -> BatonRequest | None: + """Get the current pending baton request, if any.""" + tmp = self.__client.get(f"{self.__bl}:baton_request") + if tmp is None: + return None + try: + return BatonRequest(**json.loads(tmp)) + except Exception: + return None + + def set_pending_baton_request(self, request: BatonRequest | None, timeout_sec: int = 30) -> None: + """Set a pending baton request with auto-expiry for timeout.""" + if request is None: + self.__client.delete(f"{self.__bl}:baton_request") + else: + self.__client.set(f"{self.__bl}:baton_request", request.model_dump_json()) + # Add a few seconds buffer so we can detect timeout vs expiry + self.__client.expire(f"{self.__bl}:baton_request", timeout_sec + 5) + + def clear_pending_baton_request(self) -> None: + self.__client.delete(f"{self.__bl}:baton_request") + + @property + def queued_baton_transfer(self) -> BatonTransferQueue | None: + """Get queued transfer waiting for beamline to be available.""" + tmp = self.__client.get(f"{self.__bl}:baton_transfer_queue") + if tmp is None: + return None + try: + return BatonTransferQueue(**json.loads(tmp)) + except Exception: + return None + + @queued_baton_transfer.setter + def queued_baton_transfer(self, transfer: BatonTransferQueue | None) -> None: + if transfer is None: + self.__client.delete(f"{self.__bl}:baton_transfer_queue") + else: + self.__client.set(f"{self.__bl}:baton_transfer_queue", transfer.model_dump_json()) + + def can_transfer_baton_now(self) -> bool: + """Check if baton can be transferred (beamline not mid-operation).""" + # Can't transfer while beamline is busy + if self.state_busy: + return False + # Add automation queue check here when you implement it + # if self.automation_queue_running: + # return False + return True + + def execute_baton_transfer( + self, + to_session: int, + to_username: str, + to_is_staff: bool, + to_pgroup: str | None, + expiry_sec: int + ) -> None: + """ + Atomically transfer the baton to a new holder. + Use existing active_session_lock for consistency. + """ + with redis_lock.Lock( + self.__client, f"{self.__bl}:active_session_lock", expire=10 + ): + self.__client.set(f"{self.__bl}:active_session", to_session) + self.__client.expire(f"{self.__bl}:active_session", expiry_sec) + self.baton_holder = BatonHolderInfo( + username=to_username, + session=to_session, + is_staff=to_is_staff, + pgroup=to_pgroup + ) + # Clear any pending request or queued transfer + self.clear_pending_baton_request() + self.queued_baton_transfer = None + + def process_queued_transfer_if_ready(self, expiry_sec: int) -> bool: + """ + Check if there's a queued transfer and beamline is now available. + Returns True if transfer was executed. + """ + queued = self.queued_baton_transfer + if queued is None: + return False + + if not self.can_transfer_baton_now(): + return False + + self.execute_baton_transfer( + to_session=queued.target_session, + to_username=queued.target_username, + to_is_staff=queued.target_is_staff, + to_pgroup=queued.target_pgroup, + expiry_sec=expiry_sec + ) + return True + @property def pgroup(self) -> str | None: tmp = self.__client.get(f"{self.__bl}:pgroup") @@ -624,4 +751,4 @@ class BeamlineConfig: self.__client.set(f"{self.__bl}:failed_mount_count", count) def increment_failed_mount_count(self) -> int: - return int(self.__client.incr(f"{self.__bl}:failed_mount_count")) + return int(self.__client.incr(f"{self.__bl}:failed_mount_count")) \ No newline at end of file diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 45f57d4f..94bcfe09 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -7,6 +7,8 @@ import json import cv2 import urllib3 import uvicorn + +from aare.common.auth_models import BatonStatus from aare.common.coordinate import SmargonCoordinate, Coordinate from aare.common.error_codes import export_error_codes, export_error_codes_grouped from aare.common.logger_config import setup_logger @@ -659,11 +661,19 @@ async def pgroup(token: str = Depends(oauth2_scheme)) -> str: @app.put("/access/pgroup") async def set_pgroup(val: str, token: str = Depends(oauth2_scheme)) -> str: - auth.check_jwt_ro(cfg, auth.parse_token(token)) + data = auth.parse_token(token) + + holder = cfg.baton_holder + is_current_holder = holder is not None and holder.session == data.session + + # Staff or current baton holder may change p-group even if it is not currently active. + # Everyone else must still belong to the active p-group. + if not (data.staff or is_current_holder): + auth.check_jwt_ro(cfg, data) + cfg.pgroup = val return "OK" - @app.delete("/access/pgroup") async def del_pgroup(token: str = Depends(oauth2_scheme)) -> str: auth.check_jwt_ro(cfg, auth.parse_token(token)) @@ -696,6 +706,129 @@ async def force_current_session(token: str = Depends(oauth2_scheme)) -> str: auth.force_current_sesion(cfg, data) return "OK" +# ========== BATON CONTROL ENDPOINTS ========== + +@app.get("/baton/status") +async def baton_status(token: str = Depends(oauth2_scheme)) -> BatonStatus: + """Get the current baton status for the requesting user.""" + data = auth.parse_token(token) + auth.resolve_baton_timeout_if_needed(cfg) + return auth.get_baton_status(cfg, data) + + +@app.post("/baton/request") +async def baton_request(token: str = Depends(oauth2_scheme)) -> dict: + """ + Request control (baton) of the beamline. + + - If vacant: granted immediately + - If staff requesting: granted immediately (or queued if busy) + - If same level: creates pending request with timeout + - Non-staff cannot request from staff + """ + data = auth.parse_token(token) + auth.check_jwt_ro(cfg, data) + + # holder = cfg.baton_holder + # if holder is not None and holder.is_staff and not data.staff: + # return { + # "error": True, + # "message": "Cannot request baton from staff. Please ask them directly." + # } + + return auth.request_baton(cfg, data) + +@app.post("/baton/respond") +async def baton_respond(accept: bool, token: str = Depends(oauth2_scheme)) -> dict: + """ + Current baton holder responds to a pending request. + + - accept=true: transfers baton (or queues if busy) + - accept=false: refuses the request + """ + data = auth.parse_token(token) + auth.check_jwt_ro(cfg, data) + return auth.respond_to_baton_request(cfg, data, accept) + + +@app.post("/baton/release") +async def baton_release(token: str = Depends(oauth2_scheme)) -> dict: + """Voluntarily release the baton, making the beamline vacant.""" + data = auth.parse_token(token) + return auth.release_baton(cfg, data) + + +@app.post("/baton/cancel") +async def baton_cancel(token: str = Depends(oauth2_scheme)) -> dict: + """Cancel your own pending baton request.""" + data = auth.parse_token(token) + return auth.cancel_baton_request(cfg, data) + + +@app.get("/baton/check_timeout") +async def baton_check_timeout(token: str = Depends(oauth2_scheme)) -> dict: + """ + Check if a pending request has timed out and process it. + Called by GUI to poll for timeout completion. + """ + data = auth.parse_token(token) + + resolved = auth.resolve_baton_timeout_if_needed(cfg) + if resolved is not None: + return resolved.model_dump() + + pending = cfg.pending_baton_request + if pending is None: + return {"no_pending": True} + + if pending.requester_session != data.session: + return {"not_your_request": True} + + elapsed = time.time() - pending.created_at + if elapsed < pending.timeout_seconds: + return { + "pending": True, + "remaining_seconds": pending.timeout_seconds - elapsed + } + + return auth.request_baton(cfg, data) + + +async def baton_status_event_stream(data: TokenData) -> AsyncGenerator[str, None]: + """SSE stream for baton status updates.""" + last_status = None + try: + while True: + auth.resolve_baton_timeout_if_needed(cfg) + + new_baton_status = auth.get_baton_status(cfg, data) + status_json = new_baton_status.model_dump_json() + + if status_json != last_status: + last_status = status_json + yield f"data: {status_json}\n\n" + + cfg.process_queued_transfer_if_ready(auth.SESSION_EXPIRE_SECONDS) + + await asyncio.sleep(0.5) + except asyncio.CancelledError: + return + + +@app.get("/sse/baton") +async def sse_baton(token: str = Depends(oauth2_scheme)): + """SSE endpoint for real-time baton status updates.""" + data = auth.parse_token(token) + return StreamingResponse( + baton_status_event_stream(data), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Headers": "Cache-Control" + } + ) @app.get("/beamline/settings") async def get_settings(token: str = Depends(oauth2_scheme)) -> BeamlineSettingsModel: diff --git a/src/aare/gui/main_window.py b/src/aare/gui/main_window.py index dcc40ce0..3e2b8522 100644 --- a/src/aare/gui/main_window.py +++ b/src/aare/gui/main_window.py @@ -1,4 +1,5 @@ import time +import requests import jwt from PySide6.QtCore import Qt, Slot, Signal, QTimer, QSettings @@ -12,6 +13,7 @@ from PySide6.QtWidgets import ( QDockWidget, QTabWidget, QFrame, QSizePolicy, QLabel) +from aare.common.auth_models import BatonStatus, BatonRequestStatus from aare.common.coordinate import Coordinate, SmargonCoordinate from aare.common.diffraction_geometry import DiffractionGeometry from aare.common.logger_config import setup_logger @@ -42,6 +44,7 @@ from aare.gui.threads.daq_worker import DAQWorker from aare.gui.threads.jfjoch_viewer import JFJochDBusClient from aare.gui.tutorials.tutorial_registration import register_tutorials from aare.gui.widgets.alert_banner import AlertBanner +from aare.gui.widgets.baton_request_dialog import BatonRequestDialog from aare.gui.widgets.camera_image import SampleCameraImageLabel from aare.gui.widgets.no_wheel_scroll_area import NoWheelScrollArea from aare.gui.widgets.status_bar import StatusBar @@ -71,6 +74,8 @@ class MainWindow(QMainWindow): self._controls_help_dialog = None self._cleanup_done = False + self._waiting_for_baton_response: bool = False + # Tutorial manager (define tutorials after widgets exist) self.tutorial_manager = TutorialManager(self) @@ -283,6 +288,12 @@ class MainWindow(QMainWindow): self.setStatusBar(self.status_bar) self.daq = DAQWorker(base_url=self.__base_url, token=self.__token) + + self.daq.baton_status_changed.connect(self._on_baton_status_changed) + self.daq.baton_request_result.connect(self._on_baton_request_result) + self.daq.baton_response_result.connect(self._on_baton_response_result) + self.daq.baton_timeout_checked.connect(self._on_baton_timeout_checked) + self.daq.spreadsheet.connect(self.tell_samples.new_sample_list) if self.__decoded_token.staff: self.daq.reference_tools.connect(self.ref_tools_panel.new_list) @@ -394,9 +405,18 @@ class MainWindow(QMainWindow): self.data_collection.simple.parameters_changed.connect(self.daq.smart_params) self.raster.grid_scan_size_changed.connect(self.data_collection.raster.grid_scan_size_change) + self.status_bar.set_pgroup.connect(self.daq.set_pgroup) self.status_bar.end_session.connect(self.daq.end_session) self.status_bar.force_session.connect(self.daq.force_session) + + self.status_bar.request_baton.connect(self.daq.request_baton) + self.status_bar.cancel_baton_request.connect(self.daq.cancel_baton_request) + self.status_bar.release_baton.connect(self.daq.release_baton) + + self.daq.baton_status_changed.connect(self.status_bar.update_baton_status) + self.daq.baton_status_changed.connect(self.status_bar.update_baton_status) + self.status_bar.dewar_exchange.connect(self.daq.dewar_exchange) self.status_bar.sample_exchange.connect(self.daq.sample_exchange) self.status_bar.sample_alignment.connect(self.daq.sample_alignment) @@ -404,6 +424,7 @@ class MainWindow(QMainWindow): self.status_bar.close_shutter.connect(self.daq.close_shutter) self.status_bar.open_shutter.connect(self.daq.open_shutter) + self.rotation.file_ready.connect(self.viewer.load_image) self.raster.image_selected.connect(self.viewer.load_image) @@ -657,6 +678,122 @@ class MainWindow(QMainWindow): self.__mounting = False self.video_tab.setCurrentIndex(0) + + # ========== BATON DIALOG HANDLING ========== + + @Slot(dict) + def _on_baton_incoming_request(self, info: dict): + """Show dialog when someone requests our baton.""" + requester = info.get("requester", "Another user") + timeout = info.get("timeout", 30) + + # Don't show multiple dialogs + if self._baton_request_dialog is not None and self._baton_request_dialog.isVisible(): + return + + self._baton_request_dialog = BatonRequestDialog( + requester=requester, + timeout_seconds=timeout, + parent=self + ) + self._baton_request_dialog.accepted_signal.connect(self._on_baton_dialog_accepted) + self._baton_request_dialog.refused_signal.connect(self._on_baton_dialog_refused) + self._baton_request_dialog.show() + + @Slot() + def _on_baton_dialog_accepted(self): + """User clicked Accept in baton dialog.""" + self.daq.respond_to_baton_request(True) + self._baton_request_dialog = None + + @Slot() + def _on_baton_dialog_refused(self): + """User clicked Refuse in baton dialog.""" + self.daq.respond_to_baton_request(False) + self._baton_request_dialog = None + + @Slot(dict) + def _on_baton_request_result(self, result: dict): + """Handle result of our baton request - show waiting banner with countdown.""" + if result.get("granted"): + self._waiting_for_baton_response = False + self.alert_banner.show_message("Baton acquired!", False) + logger.info("Baton acquired") + elif result.get("pending"): + self._waiting_for_baton_response = True + timeout = result.get("timeout_seconds", 30) + holder = result.get("message", "Waiting for response...") + self.alert_banner.show_waiting(f"Requesting control - {holder}", timeout) + logger.info(f"Baton request pending - {timeout}s timeout") + elif result.get("queued"): + self._waiting_for_baton_response = True + self.alert_banner.show_waiting("Control transfer queued - waiting for beamline") + logger.info("Baton transfer queued") + elif result.get("already_holder"): + self._waiting_for_baton_response = False + # Don't show anything - user already has control + logger.debug("Already baton holder") + elif result.get("error"): + self._waiting_for_baton_response = False + self.alert_banner.show_message(result.get("message", "Request failed"), True) + logger.warning(f"Baton request failed: {result.get('message')}") + + @Slot(dict) + def _on_baton_response_result(self, result: dict): + """Handle result after we responded to someone else's request.""" + if result.get("accepted"): + self.alert_banner.show_message("Control transferred", False) + elif result.get("refused"): + self.alert_banner.show_message("Request declined", False) + + @Slot(dict) + def _on_baton_request_received(self, info: dict): + """Incoming baton requests are now owned by StatusBar.""" + requester = info.get("requester", "Another user") + timeout = info.get("timeout", 30) + logger.info(f"Incoming baton request from {requester} ({timeout}s)") + + @Slot(BatonStatus) + def _on_baton_status_changed(self, status: BatonStatus): + """Handle baton status updates from SSE stream.""" + if self._waiting_for_baton_response: + pending = status.pending_request + if status.you_are_holder: + self._waiting_for_baton_response = False + self.alert_banner.show_message("Baton acquired!", False) + elif pending and pending.status == BatonRequestStatus.REFUSED: + self._waiting_for_baton_response = False + holder_name = pending.holder_username or "Current user" + self.alert_banner.show_message(f"Request declined by {holder_name}", True) + elif pending is None: + self._waiting_for_baton_response = False + self.alert_banner.clear_message() + elif status.you_have_pending_request: + import time + elapsed = time.time() - pending.created_at + remaining = max(0, int(pending.timeout_seconds - elapsed)) + if remaining > 0: + self.alert_banner.show_waiting( + f"Requesting control from {pending.holder_username or 'current user'}", + remaining + ) + else: + self.alert_banner.show_waiting("Processing timeout...", 0) + + @Slot(dict) + def _on_baton_timeout_checked(self, result: dict): + """Refresh waiting UI when the backend confirms timeout state.""" + if result.get("pending"): + remaining = int(result.get("remaining_seconds", 0)) + if self._waiting_for_baton_response: + self.alert_banner.show_waiting("Requesting control", remaining) + elif result.get("granted"): + self._waiting_for_baton_response = False + self.alert_banner.show_message("Baton acquired!", False) + elif result.get("queued"): + self._waiting_for_baton_response = True + self.alert_banner.show_waiting("Control transfer queued - waiting for beamline") + def _restore_window_state(self) -> None: settings = QSettings() geometry = settings.value("main_window/geometry") @@ -675,6 +812,12 @@ class MainWindow(QMainWindow): except Exception as e: logger.warning(f"Failed to save main window state: {e}") + # Release baton before closing + try: + self.daq.release_baton_on_close() + except Exception as e: + logger.warning(f"Failed to release baton on close: {e}") + try: self.cleanup() except Exception as e: diff --git a/src/aare/gui/threads/daq_worker.py b/src/aare/gui/threads/daq_worker.py index 16f77b8d..53c4c212 100644 --- a/src/aare/gui/threads/daq_worker.py +++ b/src/aare/gui/threads/daq_worker.py @@ -8,6 +8,7 @@ from PySide6.QtCore import Signal, QUrl, Slot, QTimer, QObject, QByteArray from PySide6.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkReply from jfjoch_client import ScanResult, ScanResultImagesInner +from aare.common.auth_models import BatonStatus from aare.common.coordinate import SmargonCoordinate, Coordinate from aare.common.error_codes import export_error_codes from aare.common.models import DAQStatusModel, SampleShortInfoList, SampleShortInfo, SampleCameraSettings, \ @@ -45,6 +46,12 @@ class DAQWorker(QObject): last_error_payload_changed = Signal(dict) last_error_payloads_changed = Signal(list) + baton_status_changed = Signal(BatonStatus) + baton_request_result = Signal(dict) + baton_response_result = Signal(dict) + baton_incoming_request = Signal(dict) + baton_timeout_checked = Signal(dict) + def __init__(self, base_url: str | None, token: str, parent=None): super().__init__(parent) self.__token = token @@ -81,10 +88,19 @@ class DAQWorker(QObject): self._server_connected: bool | None = None self._last_server_error: str | None = None + self._baton_stream_reply: QNetworkReply | None = None + self._last_baton_status: BatonStatus | None = None + + self._baton_timeout_timer = QTimer(self) + self._baton_timeout_timer.setInterval(1000) + self._baton_timeout_timer.timeout.connect(self.check_baton_timeout) + self._face_detection_stream_reply: QNetworkReply | None = None if self.__base_url is not None: self.start_face_detection_stream() + self.start_baton_stream() + self._baton_timeout_timer.start() def get_last_error_payload(self) -> dict: return dict(self._last_error_payload or {}) @@ -547,6 +563,7 @@ class DAQWorker(QObject): self.generic_delete("access/pgroup") else: self.generic_put(f"access/pgroup?val={val}") + self.send_status_request() @Slot(QNetworkReply) def _handle_all_pgroups_response(self, reply: QNetworkReply): @@ -1022,7 +1039,6 @@ class DAQWorker(QObject): out[str(k)] = str(v) return out - @Slot() @Slot() def get_error_codes(self) -> None: """ @@ -1089,4 +1105,144 @@ class DAQWorker(QObject): query.append(f"message={quote(message)}") suffix = f"?{'&'.join(query)}" if query else "" - self.generic_post(f"samcam/send_screenshot_db{suffix}") \ No newline at end of file + self.generic_post(f"samcam/send_screenshot_db{suffix}") + + def start_baton_stream(self): + """Start SSE stream for baton status updates.""" + if self.__base_url is None: + return + + if self._baton_stream_reply is not None: + return + + request = QNetworkRequest(QUrl(f"{self.__base_url}/sse/baton")) + request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8")) + reply = self.__net_manager.get(request) + reply.readyRead.connect(lambda: self._read_baton_stream(reply)) + reply.finished.connect(self._restart_baton_stream) + self._baton_stream_reply = reply + + def _restart_baton_stream(self): + self._baton_stream_reply = None + if self.__base_url is not None: + QTimer.singleShot(1000, self.start_baton_stream) + + def _read_baton_stream(self, reply: QNetworkReply): + try: + chunk = reply.readAll().data().decode("utf-8") + for line in chunk.splitlines(): + if line.startswith("data:"): + payload = line[5:].strip() + if payload: + status = BatonStatus.model_validate_json(payload) + + # Detect incoming request + if (status.incoming_request and + (self._last_baton_status is None or + not self._last_baton_status.incoming_request)): + self.baton_incoming_request.emit({ + "requester": status.pending_request.requester_username if status.pending_request else "Unknown", + "timeout": status.pending_request.timeout_seconds if status.pending_request else 30 + }) + + self._last_baton_status = status + self.baton_status_changed.emit(status) + except Exception as e: + logger.error(f"Baton stream parse error: {e}") + + @Slot() + def request_baton(self): + """Request the baton.""" + if self.__base_url is None: + logger.info("POST /baton/request") + return + + request = QNetworkRequest(QUrl(f"{self.__base_url}/baton/request")) + request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8")) + request.setRawHeader(b"Content-Type", b"application/json") + reply = self.__net_manager.post(request, QByteArray(b"")) + reply.finished.connect(lambda: self._handle_baton_request_response(reply)) + + def _handle_baton_request_response(self, reply: QNetworkReply): + try: + response_data = self.handle_response(reply) + result = json.loads(response_data) if response_data else {} + self.baton_request_result.emit(result) + + if result.get("granted") or result.get("queued") or result.get("pending"): + self.check_baton_timeout() + except Exception as e: + logger.error(f"Baton request failed: {e}") + self.http_error.emit(str(e)) + + @Slot(bool) + def respond_to_baton_request(self, accept: bool): + """Respond to an incoming baton request.""" + if self.__base_url is None: + logger.info(f"POST /baton/respond?accept={accept}") + return + + request = QNetworkRequest(QUrl(f"{self.__base_url}/baton/respond?accept={str(accept).lower()}")) + request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8")) + request.setRawHeader(b"Content-Type", b"application/json") + reply = self.__net_manager.post(request, QByteArray(b"")) + reply.finished.connect(lambda: self._handle_baton_response_result(reply)) + + def _handle_baton_response_result(self, reply: QNetworkReply): + try: + response_data = self.handle_response(reply) + result = json.loads(response_data) if response_data else {} + self.baton_response_result.emit(result) + + if result.get("accepted") or result.get("refused"): + self.send_status_request() + except Exception as e: + logger.error(f"Baton response failed: {e}") + self.http_error.emit(str(e)) + + @Slot() + def release_baton(self): + """Release the baton voluntarily.""" + self.generic_post("baton/release") + + @Slot() + def cancel_baton_request(self): + """Cancel your pending baton request.""" + self.generic_post("baton/cancel") + + @Slot() + def check_baton_timeout(self): + """Poll to check if timeout has been reached.""" + if self.__base_url is None: + return + + request = QNetworkRequest(QUrl(f"{self.__base_url}/baton/check_timeout")) + request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8")) + reply = self.__net_manager.get(request) + reply.finished.connect(lambda: self._handle_baton_timeout_response(reply)) + + def _handle_baton_timeout_response(self, reply: QNetworkReply): + try: + response_data = self.handle_response(reply) + result = json.loads(response_data) if response_data else {} + self.baton_timeout_checked.emit(result) + + if result.get("granted") or result.get("queued"): + self.send_status_request() + except Exception as e: + logger.error(f"Baton timeout check failed: {e}") + self.http_error.emit(str(e)) + + def release_baton_on_close(self): + """Release baton when GUI is closed to free the beamline.""" + try: + if hasattr(self, "_baton_timeout_timer") and self._baton_timeout_timer is not None: + self._baton_timeout_timer.stop() + self.release_baton() + from PySide6.QtCore import QEventLoop, QTimer + loop = QEventLoop() + QTimer.singleShot(500, loop.quit) + loop.exec() + logger.info("Baton release requested on GUI close") + except Exception as e: + logger.warning(f"Error releasing baton on close: {e}") \ No newline at end of file diff --git a/src/aare/gui/widgets/alert_banner.py b/src/aare/gui/widgets/alert_banner.py index fe10f4c7..c1ffe20b 100644 --- a/src/aare/gui/widgets/alert_banner.py +++ b/src/aare/gui/widgets/alert_banner.py @@ -15,6 +15,13 @@ class AlertBanner(QFrame): self._clear_timer.setSingleShot(True) self._clear_timer.timeout.connect(self.clear_message) + # Countdown timer for "waiting" state + self._countdown_timer = QTimer(self) + self._countdown_timer.setInterval(1000) + self._countdown_timer.timeout.connect(self._tick_countdown) + self._countdown_remaining = 0 + self._countdown_base_message = "" + self._label = QLabel("", self) self._label.setWordWrap(True) self._label.setAlignment(Qt.AlignmentFlag.AlignCenter) @@ -34,6 +41,8 @@ class AlertBanner(QFrame): @Slot(str, bool) def show_message(self, msg: str, is_error: bool = True): + """Show error (red) or success (green) message.""" + self._stop_countdown() self._clear_timer.stop() if not msg: @@ -77,6 +86,72 @@ class AlertBanner(QFrame): self._label.setText(decorated) self.setVisible(True) + @Slot(str, int) + def show_waiting(self, msg: str, countdown_seconds: int = 0): + """ + Show a waiting/pending message (yellow) with optional countdown. + + Args: + msg: Base message to display + countdown_seconds: If > 0, append countdown and auto-update + """ + self._clear_timer.stop() + self._stop_countdown() + + if not msg: + self.clear_message() + return + + self._countdown_base_message = msg + self._countdown_remaining = countdown_seconds + + self._apply_waiting_style() + self._update_waiting_text() + + if countdown_seconds > 0: + self._countdown_timer.start() + + self.setVisible(True) + + def _apply_waiting_style(self): + """Apply yellow/waiting style.""" + self.setStyleSheet( + "QFrame {" + " background-color: #fff8e1;" + " border: 2px solid #ffb300;" + " border-radius: 12px;" + " margin: 8px 12px 8px 12px;" + "}" + "QLabel {" + " color: #e65100;" + " font-weight: 700;" + " font-size: 20px;" + " padding: 2px 6px 2px 6px;" + "}" + ) + + def _update_waiting_text(self): + """Update the waiting message text, including countdown if active.""" + if self._countdown_remaining > 0: + decorated = f"⏳ {self._countdown_base_message} ({self._countdown_remaining}s) ⏳" + else: + decorated = f"⏳ {self._countdown_base_message} ⏳" + self._label.setText(decorated) + + def _tick_countdown(self): + """Called every second during countdown.""" + self._countdown_remaining -= 1 + if self._countdown_remaining <= 0: + self._stop_countdown() + # Don't auto-clear - let the baton status update handle that + self._update_waiting_text() + + def _stop_countdown(self): + """Stop the countdown timer.""" + self._countdown_timer.stop() + self._countdown_remaining = 0 + self._countdown_base_message = "" + @Slot() def clear_message(self): self._clear_timer.stop() diff --git a/src/aare/gui/widgets/baton_request_dialog.py b/src/aare/gui/widgets/baton_request_dialog.py new file mode 100644 index 00000000..90aab3b9 --- /dev/null +++ b/src/aare/gui/widgets/baton_request_dialog.py @@ -0,0 +1,220 @@ +from PySide6.QtCore import Qt, Signal, QTimer +from PySide6.QtWidgets import ( + QDialog, QVBoxLayout, QHBoxLayout, QLabel, + QPushButton, QProgressBar, QFrame +) +from PySide6.QtGui import QFont + + +class BatonRequestDialog(QDialog): + """ + Dialog shown to current baton holder when someone requests control. + + Based on the workflow diagram: + - User can Accept (transfer immediately or queue if busy) + - User can Refuse (deny the request) + - If user ignores/closes, timeout causes auto-transfer + """ + + accepted_signal = Signal() + refused_signal = Signal() + + def __init__(self, requester: str, timeout_seconds: int = 30, parent=None): + super().__init__(parent) + self.setWindowTitle("⚡ Baton Request") + self.setModal(False) # Non-modal so user can see beamline status + self.setMinimumWidth(400) + self.setWindowFlags( + self.windowFlags() | + Qt.WindowType.WindowStaysOnTopHint + ) + + self._timeout = timeout_seconds + self._remaining = timeout_seconds + self._requester = requester + + self._setup_ui() + self._start_timer() + + def _setup_ui(self): + layout = QVBoxLayout(self) + layout.setSpacing(15) + + # Header + header = QLabel("🔔 Control Request") + header_font = QFont() + header_font.setPointSize(14) + header_font.setBold(True) + header.setFont(header_font) + header.setAlignment(Qt.AlignmentFlag.AlignCenter) + layout.addWidget(header) + + # Separator + line = QFrame() + line.setFrameShape(QFrame.Shape.HLine) + line.setFrameShadow(QFrame.Shadow.Sunken) + layout.addWidget(line) + + # Message + self.message_label = QLabel( + f"{self._requester} is requesting control of the beamline." + ) + self.message_label.setWordWrap(True) + self.message_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + layout.addWidget(self.message_label) + + # Timeout progress + progress_layout = QVBoxLayout() + + self.progress = QProgressBar() + self.progress.setRange(0, self._timeout) + self.progress.setValue(self._timeout) + self.progress.setTextVisible(False) + self.progress.setFixedHeight(8) + self.progress.setStyleSheet(""" + QProgressBar { + border: 1px solid #ccc; + border-radius: 4px; + background-color: #f0f0f0; + } + QProgressBar::chunk { + background-color: #4CAF50; + border-radius: 3px; + } + """) + progress_layout.addWidget(self.progress) + + self.time_label = QLabel(f"{self._timeout} seconds remaining") + self.time_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.time_label.setStyleSheet("color: #666;") + progress_layout.addWidget(self.time_label) + + layout.addLayout(progress_layout) + + # Warning about auto-transfer + self.warning_label = QLabel( + "⚠️ If you don't respond, control will transfer automatically." + ) + self.warning_label.setWordWrap(True) + self.warning_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.warning_label.setStyleSheet("color: #ff9800; font-style: italic;") + layout.addWidget(self.warning_label) + + # Buttons + button_layout = QHBoxLayout() + button_layout.setSpacing(20) + + self.accept_btn = QPushButton("✓ Accept") + self.accept_btn.setMinimumHeight(40) + self.accept_btn.setStyleSheet(""" + QPushButton { + background-color: #4CAF50; + color: white; + border: none; + border-radius: 5px; + font-weight: bold; + font-size: 13px; + } + QPushButton:hover { + background-color: #45a049; + } + QPushButton:pressed { + background-color: #3d8b40; + } + """) + self.accept_btn.clicked.connect(self._on_accept) + button_layout.addWidget(self.accept_btn) + + self.refuse_btn = QPushButton("✗ Refuse") + self.refuse_btn.setMinimumHeight(40) + self.refuse_btn.setStyleSheet(""" + QPushButton { + background-color: #f44336; + color: white; + border: none; + border-radius: 5px; + font-weight: bold; + font-size: 13px; + } + QPushButton:hover { + background-color: #da190b; + } + QPushButton:pressed { + background-color: #c41000; + } + """) + self.refuse_btn.clicked.connect(self._on_refuse) + button_layout.addWidget(self.refuse_btn) + + layout.addLayout(button_layout) + + # Info text + info_label = QLabel( + "If the beamline is busy, transfer will occur after " + "the current operation completes." + ) + info_label.setWordWrap(True) + info_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + info_label.setStyleSheet("color: #999;") + layout.addWidget(info_label) + + def _start_timer(self): + self._timer = QTimer(self) + self._timer.setInterval(1000) + self._timer.timeout.connect(self._tick) + self._timer.start() + + def _tick(self): + self._remaining -= 1 + self.progress.setValue(self._remaining) + self.time_label.setText(f"{self._remaining} seconds remaining") + + # Change progress bar color as time runs out + if self._remaining <= 10: + self.progress.setStyleSheet(""" + QProgressBar { + border: 1px solid #ccc; + border-radius: 4px; + background-color: #f0f0f0; + } + QProgressBar::chunk { + background-color: #ff9800; + border-radius: 3px; + } + """) + + if self._remaining <= 5: + self.progress.setStyleSheet(""" + QProgressBar { + border: 1px solid #ccc; + border-radius: 4px; + background-color: #f0f0f0; + } + QProgressBar::chunk { + background-color: #f44336; + border-radius: 3px; + } + """) + self.time_label.setStyleSheet("color: #f44336; font-weight: bold;") + + if self._remaining <= 0: + self._timer.stop() + # Timeout = auto-accept (as per your diagram: "Ignores request" → auto transfer) + self._on_accept() + + def _on_accept(self): + self._timer.stop() + self.accepted_signal.emit() + self.accept() + + def _on_refuse(self): + self._timer.stop() + self.refused_signal.emit() + self.reject() + + def closeEvent(self, event): + """Closing the dialog counts as ignoring = auto-accept on timeout.""" + # Don't emit anything here - let the timeout handle it + # or the SSE stream will close the dialog when resolved + self._timer.stop() + super().closeEvent(event) \ No newline at end of file diff --git a/src/aare/gui/widgets/status_bar.py b/src/aare/gui/widgets/status_bar.py index b59ae947..7dec639a 100644 --- a/src/aare/gui/widgets/status_bar.py +++ b/src/aare/gui/widgets/status_bar.py @@ -5,8 +5,10 @@ from PySide6.QtGui import QFont from PySide6.QtWidgets import QStatusBar, QDialog, QMenu, QMessageBox, QLabel, QSizePolicy from aare.common.models import TokenData, BeamlineStateEnum, DAQStatusModel, SessionsStateEnum +from aare.gui.widgets.baton_request_dialog import BatonRequestDialog from aare.gui.widgets.clickable_label import ClickableLabel from aare.gui.widgets.pgroup_dialog import PGroupDialog +from aare.common.auth_models import BatonStatus, BatonRequestStatus from aare.gui.widgets.value_label import ValueLabel from aare.common.logger_config import setup_logger @@ -22,10 +24,17 @@ class StatusBar(QStatusBar): force_session = Signal() end_session = Signal() - close_shutter = Signal() - open_shutter = Signal() + request_baton = Signal() # New signal for baton request + cancel_baton_request = Signal() # Cancel pending request + release_baton = Signal() # Voluntarily release + get_all_pgroups = Signal() staff_pgroups_loaded = Signal(list) + baton_request_received = Signal(dict) + + close_shutter = Signal() + open_shutter = Signal() + def __init__(self, token: TokenData, parent=None): super().__init__(parent) @@ -38,6 +47,11 @@ class StatusBar(QStatusBar): self._message_clear_timer.setSingleShot(True) self._message_clear_timer.timeout.connect(self.clear_connection_message) + self._baton_status: BatonStatus | None = None + self._has_pending_request: bool = False + self._pgroup_dialog_for_baton: PGroupDialog | None = None + self._baton_request_dialog: BatonRequestDialog | None = None + self.message_label = QLabel("", self) self.message_label.setVisible(False) self.message_label.setSizePolicy(QSizePolicy.Policy.Maximum, QSizePolicy.Policy.Preferred) @@ -200,22 +214,133 @@ class StatusBar(QStatusBar): html_content_session = f"""Session: {session_flag}""" self.session_label.setText(html_content_session) + @Slot(BatonStatus) + def update_baton_status(self, status: BatonStatus): + """Update baton status from SSE stream.""" + prev_incoming = bool(self._baton_status and self._baton_status.incoming_request) + self._baton_status = status + self._has_pending_request = status.you_have_pending_request if status else False + self._update_session_display() + + incoming = bool(status and status.incoming_request) + if incoming and not prev_incoming: + self._show_incoming_baton_request_dialog(status) + + if self._baton_request_dialog is not None and self._baton_request_dialog.isVisible(): + if not incoming: + self._baton_request_dialog.close() + self._baton_request_dialog = None + + def _show_incoming_baton_request_dialog(self, status: BatonStatus) -> None: + if self._baton_request_dialog is not None and self._baton_request_dialog.isVisible(): + return + + requester = "Another user" + timeout = 30 + if status.pending_request is not None: + requester = status.pending_request.requester_username or requester + timeout = int(status.pending_request.timeout_seconds or timeout) + + self.baton_request_received.emit({ + "requester": requester, + "timeout": timeout, + }) + + self._baton_request_dialog = BatonRequestDialog( + requester=requester, + timeout_seconds=timeout, + parent=self + ) + self._baton_request_dialog.accepted_signal.connect(self._on_baton_dialog_accepted) + self._baton_request_dialog.refused_signal.connect(self._on_baton_dialog_refused) + self._baton_request_dialog.show() + + def _update_session_display(self): + """Update session label based on current status.""" + if self.__status is None: + return + + session_state = self.__status.session.session + + # Base text + if session_state == SessionsStateEnum.OwnedByYou: + text = "Session: You" + if self._baton_status and self._baton_status.incoming_request: + text = "Session: You (⚡ Request)" + elif session_state == SessionsStateEnum.OwnedByElse: + holder_name = "" + if self._baton_status and self._baton_status.holder: + holder_name = self._baton_status.holder.username + text = f"Session: {holder_name or 'Other'}" + if self._has_pending_request: + text += " (⏳ Waiting)" + else: + text = "Session: Vacant" + + self.session_label.setText(text) + def show_session_menu(self): menu = QMenu(self) is_busy = self.__status and self.__status.busy is_vacant = self.__status and self.__status.session.session == SessionsStateEnum.Vacant - action_1 = menu.addAction("Grab") - action_1.setEnabled(bool(not is_busy or self.__is_staff or is_vacant)) - action_1.triggered.connect(self._on_grab_clicked) - action_2 = menu.addAction("End") - action_2.setEnabled(bool(not is_busy or self.__is_staff)) - action_2.triggered.connect(self.end_session_clicked) + is_yours = self.__status and self.__status.session.session == SessionsStateEnum.OwnedByYou + is_other = self.__status and self.__status.session.session == SessionsStateEnum.OwnedByElse + + # Determine holder info from baton status + holder_is_staff = ( + self._baton_status and + self._baton_status.holder and + self._baton_status.holder.is_staff + ) + + # --- GRAB / REQUEST --- + if is_vacant: + # Vacant - simple grab + action_grab = menu.addAction("Grab") + action_grab.setEnabled(True) + action_grab.triggered.connect(self._on_grab_clicked) + elif is_other: + # Someone else has it + if self._has_pending_request: + # Already have a pending request - show cancel option + action_cancel = menu.addAction("Cancel Request") + action_cancel.triggered.connect(self._on_cancel_request_clicked) + elif self.__is_staff: + # Staff can always grab (override) + action_grab = menu.addAction("Grab (Override)") + action_grab.setEnabled(not is_busy) # Still respect busy for safety + action_grab.triggered.connect(self._on_grab_clicked) + elif holder_is_staff: + # Non-staff cannot request from staff + action_grab = menu.addAction("Request (Staff has control)") + action_grab.setEnabled(False) + else: + # Same level - request with timeout + action_request = menu.addAction("Request Control") + action_request.setEnabled(True) + action_request.triggered.connect(self._on_grab_clicked) + elif is_yours: + # You have it - show release option + action_release = menu.addAction("Release") + action_release.setEnabled(not is_busy) + action_release.triggered.connect(self._on_release_clicked) + + menu.addSeparator() + + # --- END SESSION (cleanup) --- + action_end = menu.addAction("End Session") + action_end.setEnabled(bool((is_yours and not is_busy) or self.__is_staff)) + action_end.triggered.connect(self.end_session_clicked) + + # --- STAFF: FORCE GRAB (emergency) --- + if self.__is_staff and is_other: + menu.addSeparator() + action_force = menu.addAction("⚠️ Force Take Over") + action_force.triggered.connect(self._on_force_session_clicked) label_geometry = self.session_label.geometry() menu.move(self.mapToGlobal(label_geometry.topLeft()) - QPoint(0, menu.sizeHint().height())) - menu.setFixedWidth(label_geometry.width()) - menu.exec() def show_pgroup_menu(self): @@ -302,39 +427,109 @@ class StatusBar(QStatusBar): menu.exec() def _on_grab_clicked(self): - self.grab_session_clicked() + """Handle grab/request click - shows p-group dialog first if needed.""" + available_pgroups = self.__allowed_pgroups or [] - def after_grab(): - - if not self.__status: - logger.info("status is None when session is grabbed") - return - - check_state = self.__status.state in ( - BeamlineStateEnum.SampleAlignment, - BeamlineStateEnum.SampleExchange, - BeamlineStateEnum.DewarTransfer, - BeamlineStateEnum.Maintenance, + if not self.__is_staff and not available_pgroups: + QMessageBox.warning( + self, + "No P-Groups", + "You don't have any p-groups available. Cannot request control." ) + return - session_ownership = self.__status.session.session == SessionsStateEnum.OwnedByYou - have_pgroup = (self.__status.session.current_pgroup in (self.__allowed_pgroups or [])) - allowed = (not self.__status.busy and check_state and session_ownership and have_pgroup) or self.__is_staff + if self.__is_staff and not available_pgroups: + self.request_baton.emit() + return - if allowed: - self.show_change_dialog() + if len(available_pgroups) == 1: + self.set_pgroup.emit(available_pgroups[0]) + self.request_baton.emit() + return + self._show_pgroup_selection_for_baton(available_pgroups) + + def _show_pgroup_selection_for_baton(self, available_pgroups: list[str]): + """Show p-group selection dialog before requesting baton.""" + if self._pgroup_dialog_for_baton is not None and self._pgroup_dialog_for_baton.isVisible(): + return + + # Get current p-group if available + curr_pgroup = None + if self.__status and self.__status.session: + curr_pgroup = self.__status.session.current_pgroup + + # If staff, allow all p-groups + if self.__is_staff: + def _on_staff_loaded(lst: list): + try: + merged = {str(p).strip() for p in (lst or []) if p is not None and str(p).strip()} + # Add user's own p-groups too + for pg in available_pgroups: + if pg and str(pg).strip(): + merged.add(str(pg).strip()) + if not merged: + merged = set(available_pgroups) + self._show_pgroup_dialog_and_request_baton(curr_pgroup, sorted(merged)) + finally: + try: + self.staff_pgroups_loaded.disconnect(_on_staff_loaded) + except Exception as e: + logger.debug(f"Error disconnecting staff_pgroups_loaded: {e}") + + self.staff_pgroups_loaded.connect(_on_staff_loaded) + self.__list_staff_pgroups() + else: + # Non-staff: only show their own p-groups + self._show_pgroup_dialog_and_request_baton(curr_pgroup, available_pgroups) + + def _show_pgroup_dialog_and_request_baton(self, curr_pgroup: str | None, pgroups: list[str]): + """Show the dialog and request baton if accepted.""" + self._pgroup_dialog_for_baton = PGroupDialog( + curr_pgroup=curr_pgroup, + pgroups=pgroups, + parent=self + ) + + if self._pgroup_dialog_for_baton.exec() == QDialog.DialogCode.Accepted: + selected_pgroup = self._pgroup_dialog_for_baton.get_input() + if selected_pgroup and selected_pgroup.strip(): + # Validate the selection + if not self.__is_staff and selected_pgroup not in (self.__allowed_pgroups or []): + QMessageBox.warning( + self, + "Invalid P-Group", + f"P-group '{selected_pgroup}' is not in your allowed list.\n" + f"Please select from: {', '.join(self.__allowed_pgroups or [])}" + ) + else: + # Set the p-group and request baton + self.set_pgroup.emit(selected_pgroup) + self.request_baton.emit() else: - logger.debug( - "not allowed to change pgroup due to; " - f"session ownership: {session_ownership}, allowed_pgroup: {have_pgroup}, " - f"beamline busy: {self.__status.busy}, beamline state: {check_state}" + QMessageBox.warning( + self, + "No P-Group Selected", + "You must select a p-group to request control." ) - QTimer.singleShot(300, after_grab) + self._pgroup_dialog_for_baton = None + + def _on_release_clicked(self): + """Handle release click.""" + self.release_baton.emit() + + def _on_cancel_request_clicked(self): + """Handle cancel request click.""" + self.cancel_baton_request.emit() + + def _on_force_session_clicked(self): + """Staff emergency force take over (bypasses baton protocol).""" + self.force_session.emit() def grab_session_clicked(self): - self.force_session.emit() + """Legacy method - now routes to baton request.""" + self.request_baton.emit() def end_session_clicked(self): self.end_session.emit() From 5a47411727acd72f949955c3f58e8761f44a21d9 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Thu, 26 Mar 2026 14:16:25 +0100 Subject: [PATCH 29/30] aaredaq: continuing work on baton exchange procedure --- src/aare/daq/auth.py | 3 + src/aare/daq/server.py | 10 -- src/aare/gui/main_window.py | 124 +++++++++++--------- src/aare/gui/threads/daq_worker.py | 18 ++- src/aare/gui/widgets/alert_banner.py | 5 +- src/aare/gui/widgets/pgroup_dialog.py | 100 +++++++++++++--- src/aare/gui/widgets/status_bar.py | 161 ++++++++++++-------------- 7 files changed, 243 insertions(+), 178 deletions(-) diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index 06d24829..5e946fd0 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -210,6 +210,9 @@ def get_baton_status(cfg: BeamlineConfig, data: TokenData) -> BatonStatus: def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: """ Request the baton (control) of the beamline. + + This should not depend on active p-group membership at the server route level. + Baton ownership rules are handled here. """ resolve_baton_timeout_if_needed(cfg) diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 94bcfe09..c0877fdf 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -727,15 +727,6 @@ async def baton_request(token: str = Depends(oauth2_scheme)) -> dict: - Non-staff cannot request from staff """ data = auth.parse_token(token) - auth.check_jwt_ro(cfg, data) - - # holder = cfg.baton_holder - # if holder is not None and holder.is_staff and not data.staff: - # return { - # "error": True, - # "message": "Cannot request baton from staff. Please ask them directly." - # } - return auth.request_baton(cfg, data) @app.post("/baton/respond") @@ -747,7 +738,6 @@ async def baton_respond(accept: bool, token: str = Depends(oauth2_scheme)) -> di - accept=false: refuses the request """ data = auth.parse_token(token) - auth.check_jwt_ro(cfg, data) return auth.respond_to_baton_request(cfg, data, accept) diff --git a/src/aare/gui/main_window.py b/src/aare/gui/main_window.py index 3e2b8522..a237f3f7 100644 --- a/src/aare/gui/main_window.py +++ b/src/aare/gui/main_window.py @@ -75,6 +75,7 @@ class MainWindow(QMainWindow): self._cleanup_done = False self._waiting_for_baton_response: bool = False + self._baton_request_dialog: BatonRequestDialog | None = None # Tutorial manager (define tutorials after widgets exist) self.tutorial_manager = TutorialManager(self) @@ -289,11 +290,15 @@ class MainWindow(QMainWindow): self.daq = DAQWorker(base_url=self.__base_url, token=self.__token) - self.daq.baton_status_changed.connect(self._on_baton_status_changed) + self.daq.baton_status_changed.connect(self.status_bar.update_baton_status) self.daq.baton_request_result.connect(self._on_baton_request_result) self.daq.baton_response_result.connect(self._on_baton_response_result) self.daq.baton_timeout_checked.connect(self._on_baton_timeout_checked) + self.status_bar.baton_request_received.connect(self._show_baton_request_dialog) + self.status_bar.baton_request_accepted.connect(self._accept_baton_request) + self.status_bar.baton_request_refused.connect(self._refuse_baton_request) + self.daq.spreadsheet.connect(self.tell_samples.new_sample_list) if self.__decoded_token.staff: self.daq.reference_tools.connect(self.ref_tools_panel.new_list) @@ -413,9 +418,12 @@ class MainWindow(QMainWindow): self.status_bar.request_baton.connect(self.daq.request_baton) self.status_bar.cancel_baton_request.connect(self.daq.cancel_baton_request) self.status_bar.release_baton.connect(self.daq.release_baton) - - self.daq.baton_status_changed.connect(self.status_bar.update_baton_status) - self.daq.baton_status_changed.connect(self.status_bar.update_baton_status) + self.status_bar.baton_request_accepted.connect( + lambda: self.daq.respond_to_baton_request(True) + ) + self.status_bar.baton_request_refused.connect( + lambda: self.daq.respond_to_baton_request(False) + ) self.status_bar.dewar_exchange.connect(self.daq.dewar_exchange) self.status_bar.sample_exchange.connect(self.daq.sample_exchange) @@ -682,57 +690,63 @@ class MainWindow(QMainWindow): # ========== BATON DIALOG HANDLING ========== @Slot(dict) - def _on_baton_incoming_request(self, info: dict): - """Show dialog when someone requests our baton.""" - requester = info.get("requester", "Another user") - timeout = info.get("timeout", 30) + def _show_baton_request_dialog(self, payload: dict): + requester = str(payload.get("requester") or "Another user") + timeout = int(payload.get("timeout") or 30) - # Don't show multiple dialogs if self._baton_request_dialog is not None and self._baton_request_dialog.isVisible(): return self._baton_request_dialog = BatonRequestDialog( requester=requester, timeout_seconds=timeout, - parent=self + parent=self, ) - self._baton_request_dialog.accepted_signal.connect(self._on_baton_dialog_accepted) - self._baton_request_dialog.refused_signal.connect(self._on_baton_dialog_refused) + self._baton_request_dialog.accepted_signal.connect(self.status_bar._on_baton_dialog_accepted) + self._baton_request_dialog.refused_signal.connect(self.status_bar._on_baton_dialog_refused) self._baton_request_dialog.show() + self._baton_request_dialog.raise_() + self._baton_request_dialog.activateWindow() @Slot() - def _on_baton_dialog_accepted(self): - """User clicked Accept in baton dialog.""" + def _accept_baton_request(self): self.daq.respond_to_baton_request(True) - self._baton_request_dialog = None @Slot() - def _on_baton_dialog_refused(self): - """User clicked Refuse in baton dialog.""" + def _refuse_baton_request(self): self.daq.respond_to_baton_request(False) - self._baton_request_dialog = None @Slot(dict) def _on_baton_request_result(self, result: dict): """Handle result of our baton request - show waiting banner with countdown.""" if result.get("granted"): self._waiting_for_baton_response = False - self.alert_banner.show_message("Baton acquired!", False) + self.alert_banner.show_message("Baton acquired!", False, auto_clear_ms=10000) logger.info("Baton acquired") + self._close_baton_dialog() + + available_pgroups = [str(p).strip() for p in (self.__decoded_token.pgroups or []) if p is not None and str(p).strip()] + if len(available_pgroups) == 1: + self.status_bar.set_pgroup.emit(available_pgroups[0]) + else: + self.status_bar._after_baton_granted_select_pgroup() + elif result.get("pending"): self._waiting_for_baton_response = True timeout = result.get("timeout_seconds", 30) holder = result.get("message", "Waiting for response...") self.alert_banner.show_waiting(f"Requesting control - {holder}", timeout) logger.info(f"Baton request pending - {timeout}s timeout") + elif result.get("queued"): self._waiting_for_baton_response = True self.alert_banner.show_waiting("Control transfer queued - waiting for beamline") logger.info("Baton transfer queued") + elif result.get("already_holder"): self._waiting_for_baton_response = False - # Don't show anything - user already has control logger.debug("Already baton holder") + elif result.get("error"): self._waiting_for_baton_response = False self.alert_banner.show_message(result.get("message", "Request failed"), True) @@ -741,58 +755,54 @@ class MainWindow(QMainWindow): @Slot(dict) def _on_baton_response_result(self, result: dict): """Handle result after we responded to someone else's request.""" + logger.debug(f"Baton response result: {result}") if result.get("accepted"): - self.alert_banner.show_message("Control transferred", False) + self._waiting_for_baton_response = False + self.alert_banner.show_message("Control transferred", False, auto_clear_ms=10000) + self._close_baton_dialog() + self.status_bar.update_baton_status(self.status_bar._baton_status) # refresh label state elif result.get("refused"): - self.alert_banner.show_message("Request declined", False) - - @Slot(dict) - def _on_baton_request_received(self, info: dict): - """Incoming baton requests are now owned by StatusBar.""" - requester = info.get("requester", "Another user") - timeout = info.get("timeout", 30) - logger.info(f"Incoming baton request from {requester} ({timeout}s)") - - @Slot(BatonStatus) - def _on_baton_status_changed(self, status: BatonStatus): - """Handle baton status updates from SSE stream.""" - if self._waiting_for_baton_response: - pending = status.pending_request - if status.you_are_holder: - self._waiting_for_baton_response = False - self.alert_banner.show_message("Baton acquired!", False) - elif pending and pending.status == BatonRequestStatus.REFUSED: - self._waiting_for_baton_response = False - holder_name = pending.holder_username or "Current user" - self.alert_banner.show_message(f"Request declined by {holder_name}", True) - elif pending is None: - self._waiting_for_baton_response = False - self.alert_banner.clear_message() - elif status.you_have_pending_request: - import time - elapsed = time.time() - pending.created_at - remaining = max(0, int(pending.timeout_seconds - elapsed)) - if remaining > 0: - self.alert_banner.show_waiting( - f"Requesting control from {pending.holder_username or 'current user'}", - remaining - ) - else: - self.alert_banner.show_waiting("Processing timeout...", 0) + self._waiting_for_baton_response = False + self.alert_banner.show_message("Request declined", False, auto_clear_ms=10000) + self._close_baton_dialog() + self.status_bar.update_baton_status(self.status_bar._baton_status) # refresh label state + else: + logger.debug(f"replied with {result}") @Slot(dict) def _on_baton_timeout_checked(self, result: dict): """Refresh waiting UI when the backend confirms timeout state.""" + logger.debug(f"Baton timeout checked: {result}") if result.get("pending"): remaining = int(result.get("remaining_seconds", 0)) if self._waiting_for_baton_response: self.alert_banner.show_waiting("Requesting control", remaining) + elif result.get("granted"): self._waiting_for_baton_response = False - self.alert_banner.show_message("Baton acquired!", False) + self.alert_banner.show_message("Baton acquired!", False, auto_clear_ms=10000) + self._close_baton_dialog() + elif result.get("queued"): self._waiting_for_baton_response = True self.alert_banner.show_waiting("Control transfer queued - waiting for beamline") + self._close_baton_dialog() + + elif result.get("refused"): + self._waiting_for_baton_response = False + self.alert_banner.show_message("Request declined", False, auto_clear_ms=10000) + self._close_baton_dialog() + + else: + logger.debug(f"replied with {result}") + self.alert_banner.clear_message() + + def _close_baton_dialog(self) -> None: + if self._baton_request_dialog is not None: + try: + self._baton_request_dialog.close() + finally: + self._baton_request_dialog = None def _restore_window_state(self) -> None: settings = QSettings() diff --git a/src/aare/gui/threads/daq_worker.py b/src/aare/gui/threads/daq_worker.py index 53c4c212..7643aaf9 100644 --- a/src/aare/gui/threads/daq_worker.py +++ b/src/aare/gui/threads/daq_worker.py @@ -1169,8 +1169,16 @@ class DAQWorker(QObject): result = json.loads(response_data) if response_data else {} self.baton_request_result.emit(result) - if result.get("granted") or result.get("queued") or result.get("pending"): - self.check_baton_timeout() + if result.get("granted"): + self.status_message.emit("Baton acquired", False) + self.send_status_request() + elif result.get("pending"): + self.status_message.emit( + f"Request sent - waiting for response ({result.get('timeout_seconds', 30)}s timeout)", + False + ) + elif result.get("error"): + self.status_message.emit(result.get("message", "Request failed"), True) except Exception as e: logger.error(f"Baton request failed: {e}") self.http_error.emit(str(e)) @@ -1196,10 +1204,13 @@ class DAQWorker(QObject): if result.get("accepted") or result.get("refused"): self.send_status_request() + self.check_baton_timeout() + self.start_baton_stream() except Exception as e: logger.error(f"Baton response failed: {e}") self.http_error.emit(str(e)) + @Slot() def release_baton(self): """Release the baton voluntarily.""" @@ -1227,8 +1238,9 @@ class DAQWorker(QObject): result = json.loads(response_data) if response_data else {} self.baton_timeout_checked.emit(result) - if result.get("granted") or result.get("queued"): + if result.get("granted") or result.get("queued") or result.get("refused"): self.send_status_request() + self.start_baton_stream() except Exception as e: logger.error(f"Baton timeout check failed: {e}") self.http_error.emit(str(e)) diff --git a/src/aare/gui/widgets/alert_banner.py b/src/aare/gui/widgets/alert_banner.py index c1ffe20b..6389a4e3 100644 --- a/src/aare/gui/widgets/alert_banner.py +++ b/src/aare/gui/widgets/alert_banner.py @@ -40,7 +40,7 @@ class AlertBanner(QFrame): self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed) @Slot(str, bool) - def show_message(self, msg: str, is_error: bool = True): + def show_message(self, msg: str, is_error: bool = True, auto_clear_ms: int | None = None): """Show error (red) or success (green) message.""" self._stop_countdown() self._clear_timer.stop() @@ -81,7 +81,8 @@ class AlertBanner(QFrame): " padding: 2px 6px 2px 6px;" "}" ) - self._clear_timer.start(5000) + timeout = 5000 if auto_clear_ms is None else int(auto_clear_ms) + self._clear_timer.start(timeout) self._label.setText(decorated) self.setVisible(True) diff --git a/src/aare/gui/widgets/pgroup_dialog.py b/src/aare/gui/widgets/pgroup_dialog.py index 156ddb3a..88387dcb 100644 --- a/src/aare/gui/widgets/pgroup_dialog.py +++ b/src/aare/gui/widgets/pgroup_dialog.py @@ -1,5 +1,8 @@ from PySide6.QtCore import Qt -from PySide6.QtWidgets import QDialog, QLineEdit, QVBoxLayout, QPushButton, QLabel, QComboBox, QCompleter +from PySide6.QtWidgets import ( + QDialog, QVBoxLayout, QPushButton, QLabel, + QComboBox, QCompleter, QMessageBox +) class PGroupDialog(QDialog): @@ -8,42 +11,107 @@ class PGroupDialog(QDialog): self.setWindowTitle("Change current p-group") self.setMinimumWidth(300) - # Create a layout + self._pgroups = [str(p).strip() for p in (pgroups or []) if p is not None and str(p).strip()] + layout = QVBoxLayout(self) - # Add a label self.label = QLabel("Set p-group:", self) layout.addWidget(self.label) self.combo = QComboBox(self) self.combo.setEditable(True) - items = [str(p) for p in (pgroups or []) if p is not None and str(p).strip()] - self.combo.addItems(items) + self.combo.setInsertPolicy(QComboBox.InsertPolicy.NoInsert) + self.combo.addItems(self._pgroups) - completer = QCompleter(items, self) + completer = QCompleter(self._pgroups, self) completer.setCaseSensitivity(Qt.CaseInsensitive) - completer.setFilterMode(Qt.MatchFlag.MatchContains) # requires Qt import; fallback to default if not desired + completer.setFilterMode(Qt.MatchFlag.MatchContains) + completer.setCompletionMode(QCompleter.CompletionMode.PopupCompletion) self.combo.setCompleter(completer) - if curr_pgroup and curr_pgroup in items: - self.combo.setCurrentText(curr_pgroup) - elif curr_pgroup: + default_pgroup = self._latest_pgroup(self._pgroups) + if curr_pgroup and curr_pgroup in self._pgroups: self.combo.setCurrentText(curr_pgroup) + elif default_pgroup is not None: + self.combo.setCurrentText(default_pgroup) + layout.addWidget(self.combo) - # Create buttons self.ok_button = QPushButton("OK", self) self.cancel_button = QPushButton("Cancel", self) - # Add buttons to the layout layout.addWidget(self.ok_button) layout.addWidget(self.cancel_button) - # Connect button signals - self.ok_button.clicked.connect(self.accept) + self.ok_button.clicked.connect(self._validate_and_accept) self.cancel_button.clicked.connect(self.reject) + if self.combo.lineEdit() is not None: + self.combo.lineEdit().textEdited.connect(self._live_validate) + self._live_validate(self.combo.currentText()) + + @staticmethod + def _latest_pgroup(pgroups: list[str]) -> str | None: + """ + Return the numerically largest p-group, e.g. p16371 over p01234. + Falls back to lexicographic max if parsing fails. + """ + if not pgroups: + return None + + def _key(pg: str): + s = str(pg).strip() + if s.startswith("p") and s[1:].isdigit(): + return (1, int(s[1:]), s) + return (0, -1, s) + + return max(pgroups, key=_key) + + def _set_error_state(self, is_error: bool, message: str | None = None) -> None: + if is_error: + self.combo.setStyleSheet("border: 2px solid #d9534f;") + if message: + self.label.setText(f"Set p-group: {message}") + else: + self.label.setText("Set p-group:") + else: + self.combo.setStyleSheet("") + self.label.setText("Set p-group:") + + def _live_validate(self, text: str) -> None: + text = (text or "").strip() + if not text: + self._set_error_state(True, "Select a p-group") + return + if self._pgroups and text not in self._pgroups: + self._set_error_state(True, "Not in allowed list") + return + self._set_error_state(False) + + def _validate_and_accept(self) -> None: + entered_text = (self.combo.currentText() or "").strip() + + if not entered_text: + QMessageBox.warning(self, "Invalid P-Group", "You must select a p-group.") + self.combo.setFocus() + return + + if self._pgroups and entered_text not in self._pgroups: + QMessageBox.warning( + self, + "Invalid P-Group", + f"P-group '{entered_text}' is not in your allowed list.\n" + f"Please select from: {', '.join(self._pgroups)}" + ) + self.combo.setFocus() + if self.combo.lineEdit() is not None: + self.combo.lineEdit().selectAll() + self._set_error_state(True, "Not in allowed list") + return + + self._set_error_state(False) + self.accept() + def get_input(self): """Return the input text when the dialog is accepted.""" - #return self.text_entry.text() - return self.combo.currentText() \ No newline at end of file + return (self.combo.currentText() or "").strip() \ No newline at end of file diff --git a/src/aare/gui/widgets/status_bar.py b/src/aare/gui/widgets/status_bar.py index 7dec639a..6bf7d833 100644 --- a/src/aare/gui/widgets/status_bar.py +++ b/src/aare/gui/widgets/status_bar.py @@ -24,9 +24,11 @@ class StatusBar(QStatusBar): force_session = Signal() end_session = Signal() - request_baton = Signal() # New signal for baton request - cancel_baton_request = Signal() # Cancel pending request - release_baton = Signal() # Voluntarily release + request_baton = Signal() + cancel_baton_request = Signal() + release_baton = Signal() + baton_request_accepted = Signal() + baton_request_refused = Signal() get_all_pgroups = Signal() staff_pgroups_loaded = Signal(list) @@ -35,7 +37,6 @@ class StatusBar(QStatusBar): close_shutter = Signal() open_shutter = Signal() - def __init__(self, token: TokenData, parent=None): super().__init__(parent) self.__decoded_token = token @@ -224,17 +225,14 @@ class StatusBar(QStatusBar): incoming = bool(status and status.incoming_request) if incoming and not prev_incoming: - self._show_incoming_baton_request_dialog(status) + self._emit_incoming_baton_request(status) if self._baton_request_dialog is not None and self._baton_request_dialog.isVisible(): if not incoming: self._baton_request_dialog.close() self._baton_request_dialog = None - def _show_incoming_baton_request_dialog(self, status: BatonStatus) -> None: - if self._baton_request_dialog is not None and self._baton_request_dialog.isVisible(): - return - + def _emit_incoming_baton_request(self, status: BatonStatus) -> None: requester = "Another user" timeout = 30 if status.pending_request is not None: @@ -246,14 +244,15 @@ class StatusBar(QStatusBar): "timeout": timeout, }) - self._baton_request_dialog = BatonRequestDialog( - requester=requester, - timeout_seconds=timeout, - parent=self - ) - self._baton_request_dialog.accepted_signal.connect(self._on_baton_dialog_accepted) - self._baton_request_dialog.refused_signal.connect(self._on_baton_dialog_refused) - self._baton_request_dialog.show() + @Slot() + def _on_baton_dialog_accepted(self): + self.baton_request_accepted.emit() + self._baton_request_dialog = None + + @Slot() + def _on_baton_dialog_refused(self): + self.baton_request_refused.emit() + self._baton_request_dialog = None def _update_session_display(self): """Update session label based on current status.""" @@ -426,95 +425,75 @@ class StatusBar(QStatusBar): menu.exec() - def _on_grab_clicked(self): - """Handle grab/request click - shows p-group dialog first if needed.""" - available_pgroups = self.__allowed_pgroups or [] + def _latest_pgroup(self, pgroups: list[str]) -> str | None: + if not pgroups: + return None - if not self.__is_staff and not available_pgroups: - QMessageBox.warning( - self, - "No P-Groups", - "You don't have any p-groups available. Cannot request control." - ) + def _key(pg: str): + s = str(pg).strip() + if s.startswith("p") and s[1:].isdigit(): + return (1, int(s[1:]), s) + return (0, -1, s) + + return max(pgroups, key=_key) + + def _after_baton_granted_select_pgroup(self) -> None: + """ + After baton grant: + - if exactly one allowed p-group, apply it automatically + - otherwise prompt user to choose from their allowed list + """ + pgroups = [str(p).strip() for p in (self.__allowed_pgroups or []) if p is not None and str(p).strip()] + if not pgroups: return - if self.__is_staff and not available_pgroups: - self.request_baton.emit() + if len(pgroups) == 1: + self.set_pgroup.emit(pgroups[0]) return - if len(available_pgroups) == 1: - self.set_pgroup.emit(available_pgroups[0]) - self.request_baton.emit() - return - - self._show_pgroup_selection_for_baton(available_pgroups) - - def _show_pgroup_selection_for_baton(self, available_pgroups: list[str]): - """Show p-group selection dialog before requesting baton.""" - if self._pgroup_dialog_for_baton is not None and self._pgroup_dialog_for_baton.isVisible(): - return - - # Get current p-group if available - curr_pgroup = None + default_pgroup = self._latest_pgroup(pgroups) + curr = None if self.__status and self.__status.session: - curr_pgroup = self.__status.session.current_pgroup + curr = self.__status.session.current_pgroup or default_pgroup - # If staff, allow all p-groups - if self.__is_staff: - def _on_staff_loaded(lst: list): - try: - merged = {str(p).strip() for p in (lst or []) if p is not None and str(p).strip()} - # Add user's own p-groups too - for pg in available_pgroups: - if pg and str(pg).strip(): - merged.add(str(pg).strip()) - if not merged: - merged = set(available_pgroups) - self._show_pgroup_dialog_and_request_baton(curr_pgroup, sorted(merged)) - finally: - try: - self.staff_pgroups_loaded.disconnect(_on_staff_loaded) - except Exception as e: - logger.debug(f"Error disconnecting staff_pgroups_loaded: {e}") - - self.staff_pgroups_loaded.connect(_on_staff_loaded) - self.__list_staff_pgroups() - else: - # Non-staff: only show their own p-groups - self._show_pgroup_dialog_and_request_baton(curr_pgroup, available_pgroups) - - def _show_pgroup_dialog_and_request_baton(self, curr_pgroup: str | None, pgroups: list[str]): - """Show the dialog and request baton if accepted.""" self._pgroup_dialog_for_baton = PGroupDialog( - curr_pgroup=curr_pgroup, + curr_pgroup=curr, pgroups=pgroups, - parent=self + parent=self, ) if self._pgroup_dialog_for_baton.exec() == QDialog.DialogCode.Accepted: selected_pgroup = self._pgroup_dialog_for_baton.get_input() - if selected_pgroup and selected_pgroup.strip(): - # Validate the selection - if not self.__is_staff and selected_pgroup not in (self.__allowed_pgroups or []): - QMessageBox.warning( - self, - "Invalid P-Group", - f"P-group '{selected_pgroup}' is not in your allowed list.\n" - f"Please select from: {', '.join(self.__allowed_pgroups or [])}" - ) - else: - # Set the p-group and request baton - self.set_pgroup.emit(selected_pgroup) - self.request_baton.emit() - else: - QMessageBox.warning( - self, - "No P-Group Selected", - "You must select a p-group to request control." - ) + if selected_pgroup: + self.set_pgroup.emit(selected_pgroup) self._pgroup_dialog_for_baton = None + def _show_post_grant_pgroup_dialog(self, available_pgroups: list[str]) -> None: + curr_pgroup = None + if self.__status and self.__status.session: + curr_pgroup = self.__status.session.current_pgroup + + if self._pgroup_dialog_for_baton is not None and self._pgroup_dialog_for_baton.isVisible(): + return + + self._pgroup_dialog_for_baton = PGroupDialog( + curr_pgroup=curr_pgroup, + pgroups=available_pgroups, + parent=self + ) + + if self._pgroup_dialog_for_baton.exec() == QDialog.DialogCode.Accepted: + selected_pgroup = (self._pgroup_dialog_for_baton.get_input() or "").strip() + if selected_pgroup: + self.set_pgroup.emit(selected_pgroup) + + self._pgroup_dialog_for_baton = None + + def _on_grab_clicked(self): + """Handle grab/request click - baton first, p-group after grant.""" + self.request_baton.emit() + def _on_release_clicked(self): """Handle release click.""" self.release_baton.emit() @@ -554,6 +533,7 @@ class StatusBar(QStatusBar): return def _generate_pgroup_dialogue(self, curr: str | None = None, pgroups: list | None = None): + logger.info(pgroups) dialog = PGroupDialog(curr_pgroup=curr, pgroups=pgroups) if dialog.exec() == QDialog.DialogCode.Accepted: entered_text = dialog.get_input() @@ -565,6 +545,7 @@ class StatusBar(QStatusBar): f"P-group '{entered_text}' is not in your allowed list.\n" f"Please select from: {', '.join(pgroups)}" ) + self._generate_pgroup_dialogue(curr=curr, pgroups=pgroups) return self.set_pgroup.emit(entered_text) From 5949f14599fcf261e7cab42d741927fb3ece5192 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Thu, 26 Mar 2026 17:31:13 +0100 Subject: [PATCH 30/30] aaredaq: continuing work on baton exchange procedure --- src/aare/common/auth_models.py | 3 +- src/aare/daq/auth.py | 89 ++++++------- src/aare/daq/config.py | 21 ++- src/aare/daq/server.py | 22 +++- src/aare/gui/main_window.py | 79 +++++++++++- src/aare/gui/threads/daq_worker.py | 9 +- src/aare/gui/widgets/alert_banner.py | 4 +- src/aare/gui/widgets/baton_request_dialog.py | 128 +++++++++++++++++++ src/aare/gui/widgets/status_bar.py | 18 ++- 9 files changed, 311 insertions(+), 62 deletions(-) diff --git a/src/aare/common/auth_models.py b/src/aare/common/auth_models.py index 6897751c..9930e7cd 100644 --- a/src/aare/common/auth_models.py +++ b/src/aare/common/auth_models.py @@ -48,4 +48,5 @@ class BatonStatus(BaseModel): queued_transfer: BatonTransferQueue | None = None you_are_holder: bool = False you_have_pending_request: bool = False - incoming_request: bool = False # True if YOU are being asked to give up baton \ No newline at end of file + incoming_request: bool = False + allow_non_staff_request: bool = False \ No newline at end of file diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index 5e946fd0..ac9b6601 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -117,7 +117,13 @@ def check_jwt_staff(cfg: BeamlineConfig, data: TokenData) -> None: ) from e def force_current_sesion(cfg: BeamlineConfig, data: TokenData) -> None: - cfg.force_set_active_session(data.session, SESSION_EXPIRE_SECONDS) + cfg.execute_baton_transfer( + to_session=data.session, + to_username=data.sub, + to_is_staff=data.staff, + to_pgroup=cfg.pgroup, + expiry_sec=SESSION_EXPIRE_SECONDS + ) def _finalize_expired_baton_request(cfg: BeamlineConfig, pending: BatonRequest) -> BatonStatus: """ @@ -172,48 +178,43 @@ def resolve_baton_timeout_if_needed(cfg: BeamlineConfig) -> BatonStatus | None: return _finalize_expired_baton_request(cfg, pending) def get_baton_status(cfg: BeamlineConfig, data: TokenData) -> BatonStatus: - """Get full baton status for a user.""" + """ + Build baton status scoped to the requesting session. + + Important: + - requester sees you_have_pending_request + - holder sees incoming_request + - nobody else sees the request as actionable + """ + holder = cfg.baton_holder pending = cfg.pending_baton_request - if pending is not None and pending.status == BatonRequestStatus.PENDING: - elapsed = time.time() - pending.created_at - if elapsed >= pending.timeout_seconds: - resolve_baton_timeout_if_needed(cfg) - pending = cfg.pending_baton_request - - holder = cfg.baton_holder - queued = cfg.queued_baton_transfer - - you_are_holder = holder is not None and holder.session == data.session - you_have_pending_request = ( - pending is not None and - pending.requester_session == data.session and - pending.status == BatonRequestStatus.PENDING + is_requester = bool( + pending + and pending.status in (BatonRequestStatus.PENDING, BatonRequestStatus.REFUSED) + and pending.requester_session == data.session ) - incoming_request = ( - pending is not None and - holder is not None and - holder.session == data.session and - pending.status == BatonRequestStatus.PENDING + + is_holder = bool( + pending + and pending.status == BatonRequestStatus.PENDING + and pending.holder_session == data.session ) + # Only expose the pending request object to the two relevant sessions + scoped_pending = pending if (is_requester or is_holder) else None + return BatonStatus( holder=holder, - pending_request=pending, - queued_transfer=queued, - you_are_holder=you_are_holder, - you_have_pending_request=you_have_pending_request, - incoming_request=incoming_request + pending_request=scoped_pending, + queued_transfer=cfg.queued_baton_transfer, + you_are_holder=bool(holder and holder.session == data.session), + you_have_pending_request=is_requester, + incoming_request=is_holder, + allow_non_staff_request=cfg.allow_non_staff_request_from_staff, ) - def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: - """ - Request the baton (control) of the beamline. - - This should not depend on active p-group membership at the server route level. - Baton ownership rules are handled here. - """ resolve_baton_timeout_if_needed(cfg) session_state = cfg.session_state(data.session) @@ -224,7 +225,7 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: to_username=data.sub, to_is_staff=data.staff, to_pgroup=cfg.pgroup, - expiry_sec=SESSION_EXPIRE_SECONDS + expiry_sec=SESSION_EXPIRE_SECONDS, ) return {"granted": True, "message": "Baton acquired (beamline was vacant)"} @@ -233,11 +234,11 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: return {"already_holder": True, "message": "You already hold the baton"} holder = cfg.baton_holder - - if holder and holder.is_staff and not data.staff: + print(cfg.allow_non_staff_request_from_staff) + if holder and holder.is_staff and not data.staff and not cfg.allow_non_staff_request_from_staff: return { "error": True, - "message": "Cannot request baton from staff. Please ask them directly." + "message": "Requesting baton from staff is disabled by backend policy.", } if data.staff: @@ -248,11 +249,11 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: target_is_staff=data.staff, target_pgroup=cfg.pgroup, queued_at=time.time(), - reason="beamline_busy_staff_override" + reason="beamline_busy_staff_override", ) return { "queued": True, - "message": "Staff override queued - will transfer when beamline is available" + "message": "Staff override queued - will transfer when beamline is available", } cfg.execute_baton_transfer( @@ -260,7 +261,7 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: to_username=data.sub, to_is_staff=data.staff, to_pgroup=cfg.pgroup, - expiry_sec=SESSION_EXPIRE_SECONDS + expiry_sec=SESSION_EXPIRE_SECONDS, ) return {"granted": True, "override": True, "message": "Staff override - baton acquired"} @@ -275,11 +276,11 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: "pending": True, "existing": True, "remaining_seconds": max(0, remaining), - "message": f"Request already pending ({remaining:.0f}s remaining)" + "message": f"Request already pending ({remaining:.0f}s remaining)", } return { "error": True, - "message": "Another user already has a pending request" + "message": "Another user already has a pending request", } request = BatonRequest( @@ -291,7 +292,7 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: holder_session=holder.session if holder else None, created_at=time.time(), timeout_seconds=BATON_REQUEST_TIMEOUT_SECONDS, - status=BatonRequestStatus.PENDING + status=BatonRequestStatus.PENDING, ) cfg.set_pending_baton_request(request, timeout_sec=BATON_REQUEST_TIMEOUT_SECONDS) @@ -299,7 +300,7 @@ def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict: "pending": True, "request_id": request.request_id, "timeout_seconds": BATON_REQUEST_TIMEOUT_SECONDS, - "message": f"Request sent to {holder.username if holder else 'current user'}" + "message": f"Request sent to {holder.username if holder else 'current holder'}", } def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) -> dict: diff --git a/src/aare/daq/config.py b/src/aare/daq/config.py index ba8f8328..827915d3 100644 --- a/src/aare/daq/config.py +++ b/src/aare/daq/config.py @@ -84,6 +84,20 @@ class BeamlineConfig: # Session and authentication management + @property + def allow_non_staff_request_from_staff(self) -> bool: + raw = self.__client.get(f"{self.__bl}:allow_non_staff_request_from_staff") + if raw is None: + return False + return str(raw).strip().lower() in {"1", "true", "yes", "on"} + + @allow_non_staff_request_from_staff.setter + def allow_non_staff_request_from_staff(self, enabled: bool) -> None: + if enabled: + self.__client.set(f"{self.__bl}:allow_non_staff_request_from_staff", "1") + else: + self.__client.delete(f"{self.__bl}:allow_non_staff_request_from_staff") + def generate_session(self) -> int: return int(self.__client.incr(f"{self.__bl}:session")) @@ -751,4 +765,9 @@ class BeamlineConfig: self.__client.set(f"{self.__bl}:failed_mount_count", count) def increment_failed_mount_count(self) -> int: - return int(self.__client.incr(f"{self.__bl}:failed_mount_count")) \ No newline at end of file + return int(self.__client.incr(f"{self.__bl}:failed_mount_count")) + +if __name__ == "__main__": + from aare.common.beamline import mx_beamline + cfg = BeamlineConfig(bl=mx_beamline()) + cfg.allow_non_staff_request_from_staff = True \ No newline at end of file diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index c0877fdf..9393c8c5 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -8,7 +8,7 @@ import cv2 import urllib3 import uvicorn -from aare.common.auth_models import BatonStatus +from aare.common.auth_models import BatonStatus, BatonRequestStatus from aare.common.coordinate import SmargonCoordinate, Coordinate from aare.common.error_codes import export_error_codes, export_error_codes_grouped from aare.common.logger_config import setup_logger @@ -726,6 +726,7 @@ async def baton_request(token: str = Depends(oauth2_scheme)) -> dict: - If same level: creates pending request with timeout - Non-staff cannot request from staff """ + logger.debug(cfg.allow_non_staff_request_from_staff) data = auth.parse_token(token) return auth.request_baton(cfg, data) @@ -763,17 +764,24 @@ async def baton_check_timeout(token: str = Depends(oauth2_scheme)) -> dict: """ data = auth.parse_token(token) - resolved = auth.resolve_baton_timeout_if_needed(cfg) - if resolved is not None: - return resolved.model_dump() + auth.resolve_baton_timeout_if_needed(cfg) pending = cfg.pending_baton_request if pending is None: + if cfg.baton_holder and cfg.baton_holder.session == data.session: + return {"granted": True, "message": "Baton acquired!"} + queued = cfg.queued_baton_transfer + if queued and queued.target_session == data.session: + return {"queued": True, "message": "Transfer queued"} return {"no_pending": True} if pending.requester_session != data.session: return {"not_your_request": True} + if pending.status == BatonRequestStatus.REFUSED: + cfg.clear_pending_baton_request() + return {"refused": True, "message": "Request refused"} + elapsed = time.time() - pending.created_at if elapsed < pending.timeout_seconds: return { @@ -783,6 +791,12 @@ async def baton_check_timeout(token: str = Depends(oauth2_scheme)) -> dict: return auth.request_baton(cfg, data) +@app.put("/access/allow_non_staff_request_from_staff") +async def set_allow_non_staff_request_from_staff(val: bool, token: str = Depends(oauth2_scheme)) -> str: + data = auth.parse_token(token) + auth.check_jwt_staff_only(data) + cfg.allow_non_staff_request_from_staff = val + return "OK" async def baton_status_event_stream(data: TokenData) -> AsyncGenerator[str, None]: """SSE stream for baton status updates.""" diff --git a/src/aare/gui/main_window.py b/src/aare/gui/main_window.py index a237f3f7..dcdfeca6 100644 --- a/src/aare/gui/main_window.py +++ b/src/aare/gui/main_window.py @@ -44,7 +44,7 @@ from aare.gui.threads.daq_worker import DAQWorker from aare.gui.threads.jfjoch_viewer import JFJochDBusClient from aare.gui.tutorials.tutorial_registration import register_tutorials from aare.gui.widgets.alert_banner import AlertBanner -from aare.gui.widgets.baton_request_dialog import BatonRequestDialog +from aare.gui.widgets.baton_request_dialog import BatonRequestDialog, BatonPendingDialog from aare.gui.widgets.camera_image import SampleCameraImageLabel from aare.gui.widgets.no_wheel_scroll_area import NoWheelScrollArea from aare.gui.widgets.status_bar import StatusBar @@ -76,6 +76,7 @@ class MainWindow(QMainWindow): self._waiting_for_baton_response: bool = False self._baton_request_dialog: BatonRequestDialog | None = None + self._baton_pending_dialog: BatonPendingDialog | None = None # Tutorial manager (define tutorials after widgets exist) self.tutorial_manager = TutorialManager(self) @@ -291,6 +292,7 @@ class MainWindow(QMainWindow): self.daq = DAQWorker(base_url=self.__base_url, token=self.__token) self.daq.baton_status_changed.connect(self.status_bar.update_baton_status) + self.daq.baton_status_changed.connect(self._on_baton_status_changed) self.daq.baton_request_result.connect(self._on_baton_request_result) self.daq.baton_response_result.connect(self._on_baton_response_result) self.daq.baton_timeout_checked.connect(self._on_baton_timeout_checked) @@ -716,6 +718,26 @@ class MainWindow(QMainWindow): def _refuse_baton_request(self): self.daq.respond_to_baton_request(False) + @Slot() + def _accept_baton_request(self): + self.daq.respond_to_baton_request(True) + + @Slot() + def _refuse_baton_request(self): + self.daq.respond_to_baton_request(False) + + @Slot(BatonStatus) + def _on_baton_status_changed(self, status: BatonStatus): + """Close the pending dialog immediately if the baton request has been resolved via SSE.""" + if self._waiting_for_baton_response and not status.you_have_pending_request: + self._waiting_for_baton_response = False + self._close_baton_pending_dialog() + + if status.you_are_holder: + self.alert_banner.show_message("Baton acquired!", False, auto_clear_ms=10000) + else: + self.alert_banner.show_message("Request declined or cancelled", False, auto_clear_ms=10000) + @Slot(dict) def _on_baton_request_result(self, result: dict): """Handle result of our baton request - show waiting banner with countdown.""" @@ -723,7 +745,9 @@ class MainWindow(QMainWindow): self._waiting_for_baton_response = False self.alert_banner.show_message("Baton acquired!", False, auto_clear_ms=10000) logger.info("Baton acquired") - self._close_baton_dialog() + + # Close the pending dialog immediately before showing p-group prompt + self._close_baton_pending_dialog() available_pgroups = [str(p).strip() for p in (self.__decoded_token.pgroups or []) if p is not None and str(p).strip()] if len(available_pgroups) == 1: @@ -735,6 +759,16 @@ class MainWindow(QMainWindow): self._waiting_for_baton_response = True timeout = result.get("timeout_seconds", 30) holder = result.get("message", "Waiting for response...") + + if getattr(self, "_baton_pending_dialog", None) is None: + target_user = holder.replace("Request sent to ", "") + self._baton_pending_dialog = BatonPendingDialog(target_user=target_user, timeout_seconds=timeout, + parent=self) + self._baton_pending_dialog.cancelled_signal.connect(self.daq.cancel_baton_request) + self._baton_pending_dialog.show() + else: + self._baton_pending_dialog.update_remaining(timeout) + self.alert_banner.show_waiting(f"Requesting control - {holder}", timeout) logger.info(f"Baton request pending - {timeout}s timeout") @@ -743,14 +777,23 @@ class MainWindow(QMainWindow): self.alert_banner.show_waiting("Control transfer queued - waiting for beamline") logger.info("Baton transfer queued") + if getattr(self, "_baton_pending_dialog", None) is None: + self._baton_pending_dialog = BatonPendingDialog(target_user="Current Holder", timeout_seconds=0, + parent=self) + self._baton_pending_dialog.cancelled_signal.connect(self.daq.cancel_baton_request) + self._baton_pending_dialog.show() + self._baton_pending_dialog.set_queued_state() + elif result.get("already_holder"): self._waiting_for_baton_response = False logger.debug("Already baton holder") + self._close_baton_pending_dialog() elif result.get("error"): self._waiting_for_baton_response = False self.alert_banner.show_message(result.get("message", "Request failed"), True) logger.warning(f"Baton request failed: {result.get('message')}") + self._close_baton_pending_dialog() @Slot(dict) def _on_baton_response_result(self, result: dict): @@ -777,33 +820,57 @@ class MainWindow(QMainWindow): remaining = int(result.get("remaining_seconds", 0)) if self._waiting_for_baton_response: self.alert_banner.show_waiting("Requesting control", remaining) + if getattr(self, "_baton_pending_dialog", None) is not None: + self._baton_pending_dialog.update_remaining(remaining) elif result.get("granted"): self._waiting_for_baton_response = False self.alert_banner.show_message("Baton acquired!", False, auto_clear_ms=10000) - self._close_baton_dialog() + + # Close the pending dialog immediately before showing p-group prompt + self._close_baton_pending_dialog() + # P-group logic will be handled automatically by the status_bar stream update elif result.get("queued"): self._waiting_for_baton_response = True self.alert_banner.show_waiting("Control transfer queued - waiting for beamline") - self._close_baton_dialog() + + if getattr(self, "_baton_pending_dialog", None) is not None: + self._baton_pending_dialog.set_queued_state() + else: + self._baton_pending_dialog = BatonPendingDialog(target_user="Current Holder", timeout_seconds=0, + parent=self) + self._baton_pending_dialog.cancelled_signal.connect(self.daq.cancel_baton_request) + self._baton_pending_dialog.show() + self._baton_pending_dialog.set_queued_state() elif result.get("refused"): self._waiting_for_baton_response = False self.alert_banner.show_message("Request declined", False, auto_clear_ms=10000) - self._close_baton_dialog() + self._close_baton_pending_dialog() else: logger.debug(f"replied with {result}") self.alert_banner.clear_message() + self._close_baton_pending_dialog() def _close_baton_dialog(self) -> None: - if self._baton_request_dialog is not None: + if getattr(self, "_baton_request_dialog", None) is not None: try: self._baton_request_dialog.close() finally: self._baton_request_dialog = None + def _close_baton_pending_dialog(self) -> None: + if getattr(self, "_baton_pending_dialog", None) is not None: + try: + if hasattr(self._baton_pending_dialog, '_timer'): + self._baton_pending_dialog._timer.stop() + self._baton_pending_dialog.close() + finally: + self._baton_pending_dialog = None + + def _restore_window_state(self) -> None: settings = QSettings() geometry = settings.value("main_window/geometry") diff --git a/src/aare/gui/threads/daq_worker.py b/src/aare/gui/threads/daq_worker.py index 7643aaf9..9c3a1de9 100644 --- a/src/aare/gui/threads/daq_worker.py +++ b/src/aare/gui/threads/daq_worker.py @@ -100,7 +100,6 @@ class DAQWorker(QObject): if self.__base_url is not None: self.start_face_detection_stream() self.start_baton_stream() - self._baton_timeout_timer.start() def get_last_error_payload(self) -> dict: return dict(self._last_error_payload or {}) @@ -1136,7 +1135,13 @@ class DAQWorker(QObject): if payload: status = BatonStatus.model_validate_json(payload) - # Detect incoming request + if status.you_have_pending_request: + if not self._baton_timeout_timer.isActive(): + self._baton_timeout_timer.start() + else: + if self._baton_timeout_timer.isActive(): + self._baton_timeout_timer.stop() + if (status.incoming_request and (self._last_baton_status is None or not self._last_baton_status.incoming_request)): diff --git a/src/aare/gui/widgets/alert_banner.py b/src/aare/gui/widgets/alert_banner.py index 6389a4e3..3267a7f2 100644 --- a/src/aare/gui/widgets/alert_banner.py +++ b/src/aare/gui/widgets/alert_banner.py @@ -144,7 +144,9 @@ class AlertBanner(QFrame): self._countdown_remaining -= 1 if self._countdown_remaining <= 0: self._stop_countdown() - # Don't auto-clear - let the baton status update handle that + self.clear_message() + return + self._update_waiting_text() def _stop_countdown(self): diff --git a/src/aare/gui/widgets/baton_request_dialog.py b/src/aare/gui/widgets/baton_request_dialog.py index 90aab3b9..b23bca2e 100644 --- a/src/aare/gui/widgets/baton_request_dialog.py +++ b/src/aare/gui/widgets/baton_request_dialog.py @@ -216,5 +216,133 @@ class BatonRequestDialog(QDialog): """Closing the dialog counts as ignoring = auto-accept on timeout.""" # Don't emit anything here - let the timeout handle it # or the SSE stream will close the dialog when resolved + self._timer.stop() + super().closeEvent(event) + +class BatonPendingDialog(QDialog): + """ + Dialog shown to the user who requested the baton while they wait for a response + or for the beamline queue to clear. + """ + cancelled_signal = Signal() + + def __init__(self, target_user: str, timeout_seconds: int = 30, parent=None): + super().__init__(parent) + self.setWindowTitle("⏳ Baton Request Pending") + self.setModal(False) + self.setMinimumWidth(400) + self.setWindowFlags(self.windowFlags() | Qt.WindowType.WindowStaysOnTopHint) + + self._timeout = timeout_seconds + self._remaining = timeout_seconds + self._target_user = target_user + + self._setup_ui() + self._start_timer() + + def _setup_ui(self): + layout = QVBoxLayout(self) + layout.setSpacing(15) + + self.header = QLabel("⏳ Requesting Control") + header_font = QFont() + header_font.setPointSize(14) + header_font.setBold(True) + self.header.setFont(header_font) + self.header.setAlignment(Qt.AlignmentFlag.AlignCenter) + layout.addWidget(self.header) + + line = QFrame() + line.setFrameShape(QFrame.Shape.HLine) + line.setFrameShadow(QFrame.Shadow.Sunken) + layout.addWidget(line) + + self.message_label = QLabel( + f"Waiting for {self._target_user} to respond..." + ) + self.message_label.setWordWrap(True) + self.message_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + layout.addWidget(self.message_label) + + self.progress_layout = QVBoxLayout() + self.progress = QProgressBar() + self.progress.setRange(0, max(1, self._timeout)) + self.progress.setValue(self._timeout) + self.progress.setTextVisible(False) + self.progress.setFixedHeight(8) + self.progress.setStyleSheet(""" + QProgressBar { + border: 1px solid #ccc; + border-radius: 4px; + background-color: #f0f0f0; + } + QProgressBar::chunk { + background-color: #2196F3; + border-radius: 3px; + } + """) + self.progress_layout.addWidget(self.progress) + + self.time_label = QLabel(f"{self._timeout} seconds remaining") + self.time_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.time_label.setStyleSheet("color: #666;") + self.progress_layout.addWidget(self.time_label) + + layout.addLayout(self.progress_layout) + + button_layout = QHBoxLayout() + self.cancel_btn = QPushButton("✗ Cancel Request") + self.cancel_btn.setMinimumHeight(40) + self.cancel_btn.setStyleSheet(""" + QPushButton { + background-color: #f44336; + color: white; + border: none; + border-radius: 5px; + font-weight: bold; + font-size: 13px; + } + QPushButton:hover { background-color: #da190b; } + QPushButton:pressed { background-color: #c41000; } + """) + self.cancel_btn.clicked.connect(self._on_cancel) + button_layout.addWidget(self.cancel_btn) + layout.addLayout(button_layout) + + def _start_timer(self): + self._timer = QTimer(self) + self._timer.setInterval(1000) + self._timer.timeout.connect(self._tick) + self._timer.start() + + def _tick(self): + self._remaining -= 1 + if self._remaining < 0: + self._remaining = 0 + + self.progress.setValue(self._remaining) + self.time_label.setText(f"{self._remaining} seconds remaining") + if self._remaining <= 0: + self._timer.stop() + + def update_remaining(self, remaining: int): + self._remaining = remaining + self.progress.setValue(self._remaining) + self.time_label.setText(f"{self._remaining} seconds remaining") + + def set_queued_state(self): + self._timer.stop() + self.header.setText("⏳ Transfer Queued") + self.message_label.setText("Waiting for current action to finish before receiving baton...") + self.progress.hide() + self.time_label.hide() + # Keep cancel button so they can abort the wait if they change their mind + + def _on_cancel(self): + self._timer.stop() + self.cancelled_signal.emit() + self.reject() + + def closeEvent(self, event): self._timer.stop() super().closeEvent(event) \ No newline at end of file diff --git a/src/aare/gui/widgets/status_bar.py b/src/aare/gui/widgets/status_bar.py index 6bf7d833..a6b546a9 100644 --- a/src/aare/gui/widgets/status_bar.py +++ b/src/aare/gui/widgets/status_bar.py @@ -219,10 +219,19 @@ class StatusBar(QStatusBar): def update_baton_status(self, status: BatonStatus): """Update baton status from SSE stream.""" prev_incoming = bool(self._baton_status and self._baton_status.incoming_request) + + # Detect if we just became the holder (e.g., from a queue resolving) + was_holder = bool(self._baton_status and self._baton_status.you_are_holder) + now_holder = bool(status and status.you_are_holder) + self._baton_status = status self._has_pending_request = status.you_have_pending_request if status else False self._update_session_display() + # If we just received the baton (and weren't the holder a moment ago) + if now_holder and not was_holder: + self._after_baton_granted_select_pgroup() + incoming = bool(status and status.incoming_request) if incoming and not prev_incoming: self._emit_incoming_baton_request(status) @@ -311,8 +320,11 @@ class StatusBar(QStatusBar): action_grab.triggered.connect(self._on_grab_clicked) elif holder_is_staff: # Non-staff cannot request from staff - action_grab = menu.addAction("Request (Staff has control)") - action_grab.setEnabled(False) + allowed = self._baton_status and getattr(self._baton_status, "allow_non_staff_request", False) + action_grab = menu.addAction("Request from Staff") + action_grab.setEnabled(allowed) + if allowed: + action_grab.triggered.connect(self._on_grab_clicked) else: # Same level - request with timeout action_request = menu.addAction("Request Control") @@ -343,7 +355,7 @@ class StatusBar(QStatusBar): menu.exec() def show_pgroup_menu(self): - in_curr = self.__status and self.__status.session.current_pgroup in (self.__allowed_pgroups or []) + in_curr = self.__status logger.info(f"in_curr is {in_curr}")