mx/AareDAQ
0
1
Fork 0
Code Issues 49 Pull Requests Actions Packages Projects Releases 84 Wiki Activity

Merge branch 'x10sa' into x10sa-tutorial

Browse Source 
# Conflicts:
#	src/aare/daq/devices.py
#	src/aare/daq/server.py
#	src/aare/devices/aerotech.py
#	src/aare/devices/tell_client.py
#	src/aare/gui/gui.py
#	src/aare/gui/threads/daq_worker.py
#	src/aare/gui/widgets/alert_banner.py
This commit is contained in:
appleb_m
2026-03-27 11:52:32 +01:00
parent 303173ae82 5949f14599
commit fec26c1186
18 changed files with 1873 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+1 -1
View File
@@ -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",
src/aare/common/auth_models.py
+52
View File
@@ -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
src/aare/daq/aaredb.py
+31 -20
View File
@@ -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
src/aare/daq/auth.py
+269 -1
View File
@@ -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"}
src/aare/daq/config.py
+146
View File
@@ -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
src/aare/daq/daq.py
+1 -1
View File
@@ -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
src/aare/daq/mlbox.py
+1 -1
View File
@@ -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}")
src/aare/daq/server.py
+140 -2
View File
@@ -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:
src/aare/daq/spreadsheetupdater.py
+28 -5
View File
@@ -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__":
src/aare/daq/tellupdater.py
+64 -19
View File
@@ -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()
main()
src/aare/gui/gui.py
+3 -3
View File
@@ -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
src/aare/gui/main_window.py
+220
View File
@@ -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:
src/aare/gui/panels/smargon_trace_panel.py
+1
View File
@@ -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] = []
src/aare/gui/threads/daq_worker.py
+176 -3
View File
@@ -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
# ─────────────────────────────────────────────
src/aare/gui/widgets/alert_banner.py
+81 -12
View File
@@ -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
src/aare/gui/widgets/baton_request_dialog.py
+348
View File
@@ -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"<b>{self._requester}</b> 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(
"<small>If the beamline is busy, transfer will occur after "
"the current operation completes.</small>"
)
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 <b>{self._target_user}</b> 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)
src/aare/gui/widgets/pgroup_dialog.py
+84 -16
View File
@@ -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: <span style='color:#d9534f;'>{message}</span>")
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()
return (self.combo.currentText() or "").strip()
src/aare/gui/widgets/status_bar.py
+227 -39
View File
@@ -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)