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()