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", diff --git a/src/aare/common/auth_models.py b/src/aare/common/auth_models.py new file mode 100644 index 00000000..9930e7cd --- /dev/null +++ b/src/aare/common/auth_models.py @@ -0,0 +1,52 @@ +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 + allow_non_staff_request: bool = False \ No newline at end of file diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index e165f7b1..2d1b9785 100644 --- a/src/aare/daq/aaredb.py +++ b/src/aare/daq/aaredb.py @@ -42,17 +42,28 @@ 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) - configuration.verify_ssl = False # Disable SSL verification + # --- 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" + + # 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) 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 @@ -70,7 +81,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 @@ -84,7 +95,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: @@ -94,7 +105,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: @@ -104,7 +115,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: @@ -114,7 +125,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: @@ -125,7 +136,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: @@ -136,7 +147,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: @@ -147,7 +158,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: @@ -158,7 +169,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: @@ -288,9 +299,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: @@ -352,9 +363,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, @@ -378,7 +389,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, @@ -421,7 +432,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, @@ -445,7 +456,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, @@ -463,5 +474,5 @@ class AareWrapper: return payload except Exception as e: - print(e) + logger.error(e) raise e diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index d93edd06..22704334 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 Depends, HTTPException, status from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer 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"] @@ -114,4 +119,267 @@ 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: + """ + 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: + """ + 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 + + is_requester = bool( + pending + and pending.status in (BatonRequestStatus.PENDING, BatonRequestStatus.REFUSED) + and pending.requester_session == data.session + ) + + 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=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: + 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 + 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": "Requesting baton from staff is disabled by backend policy.", + } + + 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 holder'}", + } + +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 4b353ae0..3011466f 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 @@ -83,6 +90,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")) @@ -158,6 +179,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 @@ -167,6 +189,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") @@ -631,3 +772,8 @@ class BeamlineConfig: def increment_failed_mount_count(self) -> int: 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/daq.py b/src/aare/daq/daq.py index 57968015..606d812e 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -67,7 +67,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/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}") diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index cf9a02ec..92e27d6e 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -7,7 +7,10 @@ import json import cv2 import urllib3 import uvicorn + from aare.common.coordinate import SmargonCoordinate, Coordinate, AerotechCoordinate +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 from aare.common.models import SampleShortInfo, DAQStatusModel, BeamlineStateEnum, BeamlineSettingsModel, \ @@ -691,11 +694,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)) @@ -728,6 +739,133 @@ 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 + """ + logger.debug(cfg.allow_non_staff_request_from_staff) + data = auth.parse_token(token) + 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) + 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) + + 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 { + "pending": True, + "remaining_seconds": pending.timeout_seconds - elapsed + } + + 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.""" + 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/daq/spreadsheetupdater.py b/src/aare/daq/spreadsheetupdater.py index 4f7c65f8..5cad17fb 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") @@ -134,6 +136,25 @@ def main(): """ while True: try: + import ssl + import websocket + + # 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") + 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" + ) + + # 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, header=WS_HEADERS, @@ -142,11 +163,13 @@ def main(): on_close=on_close, on_open=on_open, ) - ws.run_forever(sslopt={"cert_reqs": 0}) - except Exception as e: - print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") - print("[WS][INFO] WebSocket connection lost. Reconnecting in 5 seconds...") + print(f"[WS][INFO] Connecting to {WS_URL}...") + ws.run_forever(sslopt=ssl_opt) + + except Exception as e: + print(f"[MAIN][ERROR] WebSocket connection failed: {e}") + time.sleep(5) if __name__ == "__main__": diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index e9815dec..9792741d 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 @@ -19,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) @@ -39,19 +39,40 @@ 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: - #response = requests.get(sse_url, stream=True) - client = sseclient.SSEClient(sse_url) - print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") - handle_tell_change_event() + while True: + try: + print(f"[SSE][INFO] Attempting to connect to {sse_url}...") + 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") + 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) - for event in client.events(): - 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}") + 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] 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): """ @@ -140,14 +161,34 @@ 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) 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 modern SSL Context for a TLS Client + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + + # 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 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, @@ -156,12 +197,16 @@ def main(): on_close=on_close, on_open=on_open, ) - ws.run_forever(sslopt={"cert_reqs": 0}) - except Exception as e: - print(f"[MAIN][ERROR] ws.run_forever() crashed with: {e}") - print("[WS][INFO] WebSocket connection lost. Reconnecting in 5 seconds...") + # 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] WebSocket connection failed: {e}") + + print("[WS][INFO] Reconnecting in 5 seconds...") time.sleep(5) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/aare/gui/gui.py b/src/aare/gui/gui.py index d5957486..9b4ffb63 100644 --- a/src/aare/gui/gui.py +++ b/src/aare/gui/gui.py @@ -36,9 +36,9 @@ 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://sls-gpu-003:9089"#"tcp://x10sa-spark-01:9091" #"tcp://x10sa-spark-01:9091" #"tcp://x10sa-pserv-01:9089" # - default_pred_zmq_addr = "tcp://sls-gpu-003:9089"#"""tcp://x10sa-spark-01:9091" # + default_url = "http://mx-x10sa-queue-01.psi.ch:5210" #"http://127.0.0.1:5210" + default_zmq_addr = "tcp://x10sa-spark-01:9091" # "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" default_gonio_camera_id = 1 diff --git a/src/aare/gui/main_window.py b/src/aare/gui/main_window.py index be28070b..5c1e63bc 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 @@ -49,6 +51,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, 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 @@ -83,6 +86,10 @@ class MainWindow(QMainWindow): self._cleanup_done = False self._default_window_state = None + 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_event_bus = TutorialEventBus(self) self._tutorial_text_resolver = DictionaryTextResolver(MANUAL_MOUNT_TUTORIAL) @@ -326,6 +333,17 @@ 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.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) + + 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) @@ -436,9 +454,21 @@ 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.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) self.status_bar.sample_alignment.connect(self.daq.sample_alignment) @@ -446,6 +476,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) @@ -793,6 +824,189 @@ class MainWindow(QMainWindow): self.__mounting = False self.video_tab.setCurrentIndex(0) + + # ========== BATON DIALOG HANDLING ========== + + @Slot(dict) + def _show_baton_request_dialog(self, payload: dict): + requester = str(payload.get("requester") or "Another user") + timeout = int(payload.get("timeout") or 30) + + 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.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 _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() + 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.""" + if result.get("granted"): + self._waiting_for_baton_response = False + self.alert_banner.show_message("Baton acquired!", False, auto_clear_ms=10000) + logger.info("Baton acquired") + + # 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: + 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...") + + 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") + + 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") + + 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): + """Handle result after we responded to someone else's request.""" + logger.debug(f"Baton response result: {result}") + if result.get("accepted"): + 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._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) + 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) + + # 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") + + 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_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 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") @@ -814,6 +1028,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/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] = [] diff --git a/src/aare/gui/threads/daq_worker.py b/src/aare/gui/threads/daq_worker.py index dffd18e0..c201a938 100644 --- a/src/aare/gui/threads/daq_worker.py +++ b/src/aare/gui/threads/daq_worker.py @@ -9,6 +9,8 @@ from PySide6.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkRe from jfjoch_client import ScanResult, ScanResultImagesInner from aare.common.coordinate import SmargonCoordinate, Coordinate, AerotechCoordinate +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.exception_handler import JFJochCommunicationError from aare.common.models import DAQStatusModel, SampleShortInfoList, SampleShortInfo, SampleCameraSettings, \ @@ -29,7 +31,7 @@ class DAQWorker(QObject): reference_tools = Signal(SampleShortInfoList) http_error = Signal(str) status_message = Signal(str, bool) - + # New dedicated signals for polled device errors and request-time errors polled_devices_status = Signal(str, bool) # (message, is_error) detector_error = Signal(str, bool) # (message, is_error) @@ -59,6 +61,12 @@ class DAQWorker(QObject): workflow_automation_status = Signal(bool) # enabled workflow_event = Signal(object) # WorkflowEvent + 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._active_status_error_key = None @@ -111,10 +119,18 @@ class DAQWorker(QObject): self._last_detector_msg: str | None = None self._last_detector_is_error: bool | 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() def get_last_error_payload(self) -> dict: return dict(self._last_error_payload or {}) @@ -395,7 +411,6 @@ class DAQWorker(QObject): self._server_connected = False self._last_server_error = err_msg - # Clear device states so we show "Server reconnected" on recovery self._last_tell_connected = None self._last_smargon_connected = None self._last_aerotech_connected = None @@ -658,6 +673,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): @@ -1193,7 +1209,6 @@ class DAQWorker(QObject): out[str(k)] = str(v) return out - @Slot() @Slot() def get_error_codes(self) -> None: """ @@ -1262,6 +1277,164 @@ class DAQWorker(QObject): suffix = f"?{'&'.join(query)}" if query else "" 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) + + 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)): + 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"): + 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)) + + @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() + 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.""" + 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") 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)) + + 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}") + # ───────────────────────────────────────────── # Workflow API methods # ───────────────────────────────────────────── diff --git a/src/aare/gui/widgets/alert_banner.py b/src/aare/gui/widgets/alert_banner.py index a0ea3f78..7f18386d 100644 --- a/src/aare/gui/widgets/alert_banner.py +++ b/src/aare/gui/widgets/alert_banner.py @@ -8,18 +8,20 @@ logger = setup_logger("aareGUI") class AlertBanner(QFrame): - def __init__(self, parent=None, error_timeout_ms: int = 15000, recover_timeout_ms: int = 5000): + def __init__(self, parent=None): super().__init__(parent) - self._current_message: str | None = None - self._current_is_error: bool | None = None - self._error_timeout_ms = error_timeout_ms - self._recovery_timeout_ms = recover_timeout_ms - self._clear_timer = QTimer(self) 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) @@ -38,16 +40,15 @@ 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() if not msg: self.clear_message() return - if self._current_message == msg and self._current_is_error == is_error: - return - if is_error: decorated = f"🛑 {msg} 🛑" self.setStyleSheet( @@ -64,7 +65,6 @@ class AlertBanner(QFrame): " padding: 2px 6px 2px 6px;" "}" ) - self._clear_timer.start(self._error_timeout_ms) else: decorated = f"✅ {msg} ✅" self.setStyleSheet( @@ -81,13 +81,82 @@ class AlertBanner(QFrame): " padding: 2px 6px 2px 6px;" "}" ) - self._clear_timer.start(self._recovery_timeout_ms) + timeout = 5000 if auto_clear_ms is None else int(auto_clear_ms) + self._clear_timer.start(timeout) self._current_message = msg self._current_is_error = is_error 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() + self.clear_message() + return + + 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._current_message = None 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..b23bca2e --- /dev/null +++ b/src/aare/gui/widgets/baton_request_dialog.py @@ -0,0 +1,348 @@ +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) + +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/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 b59ae947..a6b546a9 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,18 @@ class StatusBar(QStatusBar): force_session = Signal() end_session = Signal() - close_shutter = Signal() - open_shutter = Signal() + 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) + baton_request_received = Signal(dict) + + close_shutter = Signal() + open_shutter = Signal() def __init__(self, token: TokenData, parent=None): super().__init__(parent) @@ -38,6 +48,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,26 +215,147 @@ 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) + + # 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) + + 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 _emit_incoming_baton_request(self, status: BatonStatus) -> None: + 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, + }) + + @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.""" + 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 + 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") + 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): - 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}") @@ -301,40 +437,90 @@ class StatusBar(QStatusBar): menu.exec() + def _latest_pgroup(self, pgroups: list[str]) -> str | None: + 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 _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 len(pgroups) == 1: + self.set_pgroup.emit(pgroups[0]) + return + + default_pgroup = self._latest_pgroup(pgroups) + curr = None + if self.__status and self.__status.session: + curr = self.__status.session.current_pgroup or default_pgroup + + self._pgroup_dialog_for_baton = PGroupDialog( + curr_pgroup=curr, + 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: + 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): - self.grab_session_clicked() + """Handle grab/request click - baton first, p-group after grant.""" + self.request_baton.emit() - def after_grab(): + def _on_release_clicked(self): + """Handle release click.""" + self.release_baton.emit() - if not self.__status: - logger.info("status is None when session is grabbed") - return + def _on_cancel_request_clicked(self): + """Handle cancel request click.""" + self.cancel_baton_request.emit() - check_state = self.__status.state in ( - BeamlineStateEnum.SampleAlignment, - BeamlineStateEnum.SampleExchange, - BeamlineStateEnum.DewarTransfer, - BeamlineStateEnum.Maintenance, - ) - - 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 allowed: - self.show_change_dialog() - - 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}" - ) - - QTimer.singleShot(300, after_grab) + 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() @@ -359,6 +545,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() @@ -370,6 +557,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)